feat: add Ollama client for local LLM support
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1 @@
|
||||
export { OllamaClient, type OllamaClientConfig } from './ollama.js';
|
||||
@@ -0,0 +1,30 @@
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import { OllamaClient } from './ollama.js';
|
||||
|
||||
vi.mock('ollama', () => ({
|
||||
Ollama: vi.fn().mockImplementation(() => ({
|
||||
chat: vi.fn().mockResolvedValue({
|
||||
message: { content: 'Hello from Ollama!' },
|
||||
done_reason: 'stop',
|
||||
prompt_eval_count: 10,
|
||||
eval_count: 5,
|
||||
}),
|
||||
})),
|
||||
}));
|
||||
|
||||
describe('OllamaClient', () => {
|
||||
it('sends messages and returns response', async () => {
|
||||
const client = new OllamaClient({
|
||||
model: 'llama3.2',
|
||||
});
|
||||
|
||||
const response = await client.chat({
|
||||
messages: [{ role: 'user', content: 'Hello' }],
|
||||
});
|
||||
|
||||
expect(response.content).toBe('Hello from Ollama!');
|
||||
expect(response.stopReason).toBe('stop');
|
||||
expect(response.usage.inputTokens).toBe(10);
|
||||
expect(response.usage.outputTokens).toBe(5);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,45 @@
|
||||
import { Ollama } from 'ollama';
|
||||
import type { ChatRequest, ChatResponse, ModelClient } from '../types.js';
|
||||
|
||||
export interface OllamaClientConfig {
|
||||
host?: string;
|
||||
model: string;
|
||||
}
|
||||
|
||||
export class OllamaClient implements ModelClient {
|
||||
private client: Ollama;
|
||||
private model: string;
|
||||
|
||||
constructor(config: OllamaClientConfig) {
|
||||
this.client = new Ollama({
|
||||
host: config.host ?? 'http://localhost:11434',
|
||||
});
|
||||
this.model = config.model;
|
||||
}
|
||||
|
||||
async chat(request: ChatRequest): Promise<ChatResponse> {
|
||||
const messages: Array<{ role: 'system' | 'user' | 'assistant'; content: string }> = [];
|
||||
|
||||
if (request.system) {
|
||||
messages.push({ role: 'system', content: request.system });
|
||||
}
|
||||
|
||||
for (const msg of request.messages) {
|
||||
messages.push({ role: msg.role, content: msg.content });
|
||||
}
|
||||
|
||||
const response = await this.client.chat({
|
||||
model: this.model,
|
||||
messages,
|
||||
});
|
||||
|
||||
return {
|
||||
content: response.message.content,
|
||||
stopReason: response.done_reason ?? 'stop',
|
||||
usage: {
|
||||
inputTokens: response.prompt_eval_count ?? 0,
|
||||
outputTokens: response.eval_count ?? 0,
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user