diff --git a/src/core/llm/LlmClient.ts b/src/core/llm/LlmClient.ts new file mode 100644 index 0000000..03616da --- /dev/null +++ b/src/core/llm/LlmClient.ts @@ -0,0 +1,50 @@ +import OpenAI from 'openai'; + +export interface ChatMessage { + role: 'system' | 'user' | 'assistant'; + content: string; +} + +export interface ChatOptions { + model: string; + temperature?: number; + maxTokens?: number; +} + +export interface ChatResult { + content: string; + usage?: { promptTokens?: number; completionTokens?: number; totalTokens?: number }; +} + +export class LlmClient { + private readonly client: OpenAI; + + constructor(opts: { baseUrl: string; apiKey?: string }) { + this.client = new OpenAI({ baseURL: opts.baseUrl, apiKey: opts.apiKey || 'not-needed' }); + } + + async chat(messages: ChatMessage[], options: ChatOptions): Promise { + const response = await this.client.chat.completions.create({ + model: options.model, + temperature: options.temperature, + max_tokens: options.maxTokens, + messages, + }); + + const choice = response.choices[0]; + if (!choice) { + throw new Error('LLM returned no choices'); + } + + return { + content: choice.message?.content ?? '', + usage: response.usage + ? { + promptTokens: response.usage.prompt_tokens, + completionTokens: response.usage.completion_tokens, + totalTokens: response.usage.total_tokens, + } + : undefined, + }; + } +} diff --git a/test/core/llm/LlmClient.test.ts b/test/core/llm/LlmClient.test.ts new file mode 100644 index 0000000..6c73017 --- /dev/null +++ b/test/core/llm/LlmClient.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it, vi } from 'vitest'; + +const createMock = vi.fn(); + +vi.mock('openai', () => ({ + default: vi.fn().mockImplementation(() => ({ + chat: { completions: { create: createMock } }, + })), +})); + +import { LlmClient } from '../../../src/core/llm/LlmClient'; + +describe('LlmClient', () => { + it('sends messages/options to the OpenAI-compatible endpoint and returns content + usage', async () => { + createMock.mockResolvedValue({ + choices: [{ message: { content: 'hello' } }], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }); + + const client = new LlmClient({ baseUrl: 'http://localhost:11434/v1', apiKey: 'key' }); + const result = await client.chat([{ role: 'user', content: 'hi' }], { model: 'llama3', temperature: 0.7 }); + + expect(result).toEqual({ + content: 'hello', + usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15 }, + }); + expect(createMock).toHaveBeenCalledWith({ + model: 'llama3', + temperature: 0.7, + max_tokens: undefined, + messages: [{ role: 'user', content: 'hi' }], + }); + }); + + it('throws when the response has no choices', async () => { + createMock.mockResolvedValue({ choices: [] }); + + const client = new LlmClient({ baseUrl: 'http://localhost:11434/v1' }); + await expect(client.chat([{ role: 'user', content: 'hi' }], { model: 'llama3' })).rejects.toThrow( + /no choices/ + ); + }); +});