diff --git a/src/models/index.ts b/src/models/index.ts index 6e2626a..11c7316 100644 --- a/src/models/index.ts +++ b/src/models/index.ts @@ -1,2 +1,3 @@ export { AnthropicClient, type AnthropicClientConfig } from './anthropic.js'; +export { OpenAIClient, type OpenAIClientConfig } from './openai.js'; export type { Message, ChatRequest, ChatResponse, ModelClient } from './types.js'; diff --git a/src/models/openai.test.ts b/src/models/openai.test.ts new file mode 100644 index 0000000..1b37a0d --- /dev/null +++ b/src/models/openai.test.ts @@ -0,0 +1,33 @@ +import { describe, it, expect, vi } from 'vitest'; +import { OpenAIClient } from './openai.js'; + +vi.mock('openai', () => ({ + default: vi.fn().mockImplementation(() => ({ + chat: { + completions: { + create: vi.fn().mockResolvedValue({ + choices: [{ message: { content: 'Hello from GPT!' }, finish_reason: 'stop' }], + usage: { prompt_tokens: 10, completion_tokens: 5 }, + }), + }, + }, + })), +})); + +describe('OpenAIClient', () => { + it('sends messages and returns response', async () => { + const client = new OpenAIClient({ + apiKey: 'test-key', + model: 'gpt-4o', + }); + + const response = await client.chat({ + messages: [{ role: 'user', content: 'Hello' }], + }); + + expect(response.content).toBe('Hello from GPT!'); + expect(response.stopReason).toBe('stop'); + expect(response.usage.inputTokens).toBe(10); + expect(response.usage.outputTokens).toBe(5); + }); +}); diff --git a/src/models/openai.ts b/src/models/openai.ts new file mode 100644 index 0000000..3b75ca0 --- /dev/null +++ b/src/models/openai.ts @@ -0,0 +1,54 @@ +import OpenAI from 'openai'; +import type { ChatRequest, ChatResponse, ModelClient } from './types.js'; + +export interface OpenAIClientConfig { + apiKey?: string; + model: string; + maxTokens?: number; + baseURL?: string; +} + +export class OpenAIClient implements ModelClient { + private client: OpenAI; + private model: string; + private defaultMaxTokens: number; + + constructor(config: OpenAIClientConfig) { + this.client = new OpenAI({ + apiKey: config.apiKey, + baseURL: config.baseURL, + }); + this.model = config.model; + this.defaultMaxTokens = config.maxTokens ?? 4096; + } + + async chat(request: ChatRequest): Promise { + const messages: OpenAI.ChatCompletionMessageParam[] = []; + + 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.completions.create({ + model: this.model, + max_tokens: request.maxTokens ?? this.defaultMaxTokens, + messages, + }); + + const choice = response.choices[0]; + const content = choice?.message?.content ?? ''; + + return { + content, + stopReason: choice?.finish_reason ?? 'stop', + usage: { + inputTokens: response.usage?.prompt_tokens ?? 0, + outputTokens: response.usage?.completion_tokens ?? 0, + }, + }; + } +}