pacnpal/Roo-Code
1
0
Fork 0
mirror of https://github.com/pacnpal/Roo-Code.git synced 2025-12-20 04:11:10 -05:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #295 from RooVetGit/add_coverage

Browse Source 
Add test coverage
This commit is contained in:
Matt Rubens
2025-01-09 12:34:05 -05:00
committed by GitHub
parent 0429cc5d97 5fbfe9b775
commit f7759c1be5
19 changed files with 3106 additions and 493 deletions
Download Patch File Download Diff File Expand all files Collapse all files

168
src/api/providers/__tests__/anthropic.test.ts Normal file
View File

@@ -0,0 +1,168 @@
import { AnthropicHandler } from '../anthropic';
import { ApiHandlerOptions } from '../../../shared/api';
import { ApiStream } from '../../transform/stream';
import { Anthropic } from '@anthropic-ai/sdk';
// Mock Anthropic client
const mockBetaCreate = jest.fn();
const mockCreate = jest.fn();
jest.mock('@anthropic-ai/sdk', () => {
return {
Anthropic: jest.fn().mockImplementation(() => ({
beta: {
promptCaching: {
messages: {
create: mockBetaCreate.mockImplementation(async () => ({
async *[Symbol.asyncIterator]() {
yield {
type: 'message_start',
message: {
usage: {
input_tokens: 100,
output_tokens: 50,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 10
}
}
};
yield {
type: 'content_block_start',
index: 0,
content_block: {
type: 'text',
text: 'Hello'
}
};
yield {
type: 'content_block_delta',
delta: {
type: 'text_delta',
text: ' world'
}
};
}
}))
}
}
},
messages: {
create: mockCreate
}
}))
};
});
describe('AnthropicHandler', () => {
let handler: AnthropicHandler;
let mockOptions: ApiHandlerOptions;
beforeEach(() => {
mockOptions = {
apiKey: 'test-api-key',
apiModelId: 'claude-3-5-sonnet-20241022'
};
handler = new AnthropicHandler(mockOptions);
mockBetaCreate.mockClear();
mockCreate.mockClear();
});
describe('constructor', () => {
it('should initialize with provided options', () => {
expect(handler).toBeInstanceOf(AnthropicHandler);
expect(handler.getModel().id).toBe(mockOptions.apiModelId);
});
it('should initialize with undefined API key', () => {
// The SDK will handle API key validation, so we just verify it initializes
const handlerWithoutKey = new AnthropicHandler({
...mockOptions,
apiKey: undefined
});
expect(handlerWithoutKey).toBeInstanceOf(AnthropicHandler);
});
it('should use custom base URL if provided', () => {
const customBaseUrl = 'https://custom.anthropic.com';
const handlerWithCustomUrl = new AnthropicHandler({
...mockOptions,
anthropicBaseUrl: customBaseUrl
});
expect(handlerWithCustomUrl).toBeInstanceOf(AnthropicHandler);
});
});
describe('createMessage', () => {
const systemPrompt = 'You are a helpful assistant.';
const messages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: [{
type: 'text' as const,
text: 'Hello!'
}]
}
];
it('should handle prompt caching for supported models', async () => {
const stream = handler.createMessage(systemPrompt, [
{
role: 'user',
content: [{ type: 'text' as const, text: 'First message' }]
},
{
role: 'assistant',
content: [{ type: 'text' as const, text: 'Response' }]
},
{
role: 'user',
content: [{ type: 'text' as const, text: 'Second message' }]
}
]);
const chunks: any[] = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
// Verify usage information
const usageChunk = chunks.find(chunk => chunk.type === 'usage');
expect(usageChunk).toBeDefined();
expect(usageChunk?.inputTokens).toBe(100);
expect(usageChunk?.outputTokens).toBe(50);
expect(usageChunk?.cacheWriteTokens).toBe(20);
expect(usageChunk?.cacheReadTokens).toBe(10);
// Verify text content
const textChunks = chunks.filter(chunk => chunk.type === 'text');
expect(textChunks).toHaveLength(2);
expect(textChunks[0].text).toBe('Hello');
expect(textChunks[1].text).toBe(' world');
// Verify beta API was used
expect(mockBetaCreate).toHaveBeenCalled();
expect(mockCreate).not.toHaveBeenCalled();
});
});
describe('getModel', () => {
it('should return default model if no model ID is provided', () => {
const handlerWithoutModel = new AnthropicHandler({
...mockOptions,
apiModelId: undefined
});
const model = handlerWithoutModel.getModel();
expect(model.id).toBeDefined();
expect(model.info).toBeDefined();
});
it('should return specified model if valid model ID is provided', () => {
const model = handler.getModel();
expect(model.id).toBe(mockOptions.apiModelId);
expect(model.info).toBeDefined();
expect(model.info.maxTokens).toBe(8192);
expect(model.info.contextWindow).toBe(200_000);
expect(model.info.supportsImages).toBe(true);
expect(model.info.supportsPromptCache).toBe(true);
});
});
});

299
src/api/providers/__tests__/bedrock.test.ts
View File

@@ -1,191 +1,144 @@
import { AwsBedrockHandler } from '../bedrock'
import { ApiHandlerOptions, ModelInfo } from '../../../shared/api'
import { Anthropic } from '@anthropic-ai/sdk'
import { StreamEvent } from '../bedrock'
// Simplified mock for BedrockRuntimeClient
class MockBedrockRuntimeClient {
private _region: string
private mockStream: StreamEvent[] = []
constructor(config: { region: string }) {
this._region = config.region
}
async send(command: any): Promise<{ stream: AsyncIterableIterator<StreamEvent> }> {
return {
stream: this.createMockStream()
}
}
private createMockStream(): AsyncIterableIterator<StreamEvent> {
const self = this;
return {
async *[Symbol.asyncIterator]() {
for (const event of self.mockStream) {
yield event;
}
},
next: async () => {
const value = this.mockStream.shift();
return value ? { value, done: false } : { value: undefined, done: true };
},
return: async () => ({ value: undefined, done: true }),
throw: async (e) => { throw e; }
};
}
setMockStream(stream: StreamEvent[]) {
this.mockStream = stream;
}
get config() {
return { region: this._region };
}
}
import { AwsBedrockHandler } from '../bedrock';
import { MessageContent } from '../../../shared/api';
import { BedrockRuntimeClient } from '@aws-sdk/client-bedrock-runtime';
import { Anthropic } from '@anthropic-ai/sdk';
describe('AwsBedrockHandler', () => {
const mockOptions: ApiHandlerOptions = {
awsRegion: 'us-east-1',
awsAccessKey: 'mock-access-key',
awsSecretKey: 'mock-secret-key',
apiModelId: 'anthropic.claude-v2',
}
let handler: AwsBedrockHandler;
// Override the BedrockRuntimeClient creation in the constructor
class TestAwsBedrockHandler extends AwsBedrockHandler {
constructor(options: ApiHandlerOptions, mockClient?: MockBedrockRuntimeClient) {
super(options)
if (mockClient) {
// Force type casting to bypass strict type checking
(this as any)['client'] = mockClient
}
}
}
beforeEach(() => {
handler = new AwsBedrockHandler({
apiModelId: 'anthropic.claude-3-5-sonnet-20241022-v2:0',
awsAccessKey: 'test-access-key',
awsSecretKey: 'test-secret-key',
awsRegion: 'us-east-1'
});
});
test('constructor initializes with correct AWS credentials', () => {
const mockClient = new MockBedrockRuntimeClient({
region: 'us-east-1'
})
describe('constructor', () => {
it('should initialize with provided config', () => {
expect(handler['options'].awsAccessKey).toBe('test-access-key');
expect(handler['options'].awsSecretKey).toBe('test-secret-key');
expect(handler['options'].awsRegion).toBe('us-east-1');
expect(handler['options'].apiModelId).toBe('anthropic.claude-3-5-sonnet-20241022-v2:0');
});
const handler = new TestAwsBedrockHandler(mockOptions, mockClient)
// Verify that the client is created with the correct configuration
expect(handler['client']).toBeDefined()
expect(handler['client'].config.region).toBe('us-east-1')
})
it('should initialize with missing AWS credentials', () => {
const handlerWithoutCreds = new AwsBedrockHandler({
apiModelId: 'anthropic.claude-3-5-sonnet-20241022-v2:0',
awsRegion: 'us-east-1'
});
expect(handlerWithoutCreds).toBeInstanceOf(AwsBedrockHandler);
});
});
test('getModel returns correct model info', () => {
const mockClient = new MockBedrockRuntimeClient({
region: 'us-east-1'
})
const handler = new TestAwsBedrockHandler(mockOptions, mockClient)
const result = handler.getModel()
expect(result).toEqual({
id: 'anthropic.claude-v2',
info: {
maxTokens: 5000,
contextWindow: 128_000,
supportsPromptCache: false
}
})
})
test('createMessage handles successful stream events', async () => {
const mockClient = new MockBedrockRuntimeClient({
region: 'us-east-1'
})
// Mock stream events
const mockStreamEvents: StreamEvent[] = [
describe('createMessage', () => {
const mockMessages: Anthropic.Messages.MessageParam[] = [
{
metadata: {
usage: {
inputTokens: 50,
outputTokens: 100
}
}
role: 'user',
content: 'Hello'
},
{
contentBlockStart: {
start: {
text: 'Hello'
}
}
},
{
contentBlockDelta: {
delta: {
text: ' world'
}
}
},
{
messageStop: {
stopReason: 'end_turn'
}
role: 'assistant',
content: 'Hi there!'
}
]
];
mockClient.setMockStream(mockStreamEvents)
const systemPrompt = 'You are a helpful assistant';
const handler = new TestAwsBedrockHandler(mockOptions, mockClient)
it('should handle text messages correctly', async () => {
const mockResponse = {
messages: [{
role: 'assistant',
content: [{ type: 'text', text: 'Hello! How can I help you?' }]
}],
usage: {
input_tokens: 10,
output_tokens: 5
}
};
const systemPrompt = 'You are a helpful assistant'
const messages: Anthropic.Messages.MessageParam[] = [
{ role: 'user', content: 'Say hello' }
]
// Mock AWS SDK invoke
const mockStream = {
[Symbol.asyncIterator]: async function* () {
yield {
metadata: {
usage: {
inputTokens: 10,
outputTokens: 5
}
}
};
}
};
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
const mockInvoke = jest.fn().mockResolvedValue({
stream: mockStream
});
for await (const chunk of generator) {
chunks.push(chunk)
}
handler['client'] = {
send: mockInvoke
} as unknown as BedrockRuntimeClient;
// Verify the chunks match expected stream events
expect(chunks).toHaveLength(3)
expect(chunks[0]).toEqual({
type: 'usage',
inputTokens: 50,
outputTokens: 100
})
expect(chunks[1]).toEqual({
type: 'text',
text: 'Hello'
})
expect(chunks[2]).toEqual({
type: 'text',
text: ' world'
})
})
test('createMessage handles error scenarios', async () => {
const mockClient = new MockBedrockRuntimeClient({
region: 'us-east-1'
})
// Simulate an error by overriding the send method
mockClient.send = () => {
throw new Error('API request failed')
}
const handler = new TestAwsBedrockHandler(mockOptions, mockClient)
const systemPrompt = 'You are a helpful assistant'
const messages: Anthropic.Messages.MessageParam[] = [
{ role: 'user', content: 'Cause an error' }
]
await expect(async () => {
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
const stream = handler.createMessage(systemPrompt, mockMessages);
const chunks = [];
for await (const chunk of generator) {
chunks.push(chunk)
for await (const chunk of stream) {
chunks.push(chunk);
}
}).rejects.toThrow('API request failed')
})
})
expect(chunks.length).toBeGreaterThan(0);
expect(chunks[0]).toEqual({
type: 'usage',
inputTokens: 10,
outputTokens: 5
});
expect(mockInvoke).toHaveBeenCalledWith(expect.objectContaining({
input: expect.objectContaining({
modelId: 'anthropic.claude-3-5-sonnet-20241022-v2:0'
})
}));
});
it('should handle API errors', async () => {
// Mock AWS SDK invoke with error
const mockInvoke = jest.fn().mockRejectedValue(new Error('AWS Bedrock error'));
handler['client'] = {
send: mockInvoke
} as unknown as BedrockRuntimeClient;
const stream = handler.createMessage(systemPrompt, mockMessages);
await expect(async () => {
for await (const chunk of stream) {
// Should throw before yielding any chunks
}
}).rejects.toThrow('AWS Bedrock error');
});
});
describe('getModel', () => {
it('should return correct model info in test environment', () => {
const modelInfo = handler.getModel();
expect(modelInfo.id).toBe('anthropic.claude-3-5-sonnet-20241022-v2:0');
expect(modelInfo.info).toBeDefined();
expect(modelInfo.info.maxTokens).toBe(5000); // Test environment value
expect(modelInfo.info.contextWindow).toBe(128_000); // Test environment value
});
it('should return test model info for invalid model in test environment', () => {
const invalidHandler = new AwsBedrockHandler({
apiModelId: 'invalid-model',
awsAccessKey: 'test-access-key',
awsSecretKey: 'test-secret-key',
awsRegion: 'us-east-1'
});
const modelInfo = invalidHandler.getModel();
expect(modelInfo.id).toBe('invalid-model'); // In test env, returns whatever is passed
expect(modelInfo.info.maxTokens).toBe(5000);
expect(modelInfo.info.contextWindow).toBe(128_000);
});
});
});

324
src/api/providers/__tests__/deepseek.test.ts
View File

@@ -1,167 +1,203 @@
import { DeepSeekHandler } from '../deepseek'
import { ApiHandlerOptions } from '../../../shared/api'
import OpenAI from 'openai'
import { Anthropic } from '@anthropic-ai/sdk'
import { DeepSeekHandler } from '../deepseek';
import { ApiHandlerOptions, deepSeekDefaultModelId } from '../../../shared/api';
import OpenAI from 'openai';
import { Anthropic } from '@anthropic-ai/sdk';
// Mock dependencies
jest.mock('openai')
// Mock OpenAI client
const mockCreate = jest.fn();
jest.mock('openai', () => {
return {
__esModule: true,
default: jest.fn().mockImplementation(() => ({
chat: {
completions: {
create: mockCreate.mockImplementation(async (options) => {
if (!options.stream) {
return {
id: 'test-completion',
choices: [{
message: { role: 'assistant', content: 'Test response', refusal: null },
finish_reason: 'stop',
index: 0
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15
}
};
}
// Return async iterator for streaming
return {
[Symbol.asyncIterator]: async function* () {
yield {
choices: [{
delta: { content: 'Test response' },
index: 0
}],
usage: null
};
yield {
choices: [{
delta: {},
index: 0
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15
}
};
}
};
})
}
}
}))
};
});
describe('DeepSeekHandler', () => {
const mockOptions: ApiHandlerOptions = {
deepSeekApiKey: 'test-key',
deepSeekModelId: 'deepseek-chat',
}
let handler: DeepSeekHandler;
let mockOptions: ApiHandlerOptions;
beforeEach(() => {
jest.clearAllMocks()
})
mockOptions = {
deepSeekApiKey: 'test-api-key',
deepSeekModelId: 'deepseek-chat',
deepSeekBaseUrl: 'https://api.deepseek.com/v1'
};
handler = new DeepSeekHandler(mockOptions);
mockCreate.mockClear();
});
test('constructor initializes with correct options', () => {
const handler = new DeepSeekHandler(mockOptions)
expect(handler).toBeInstanceOf(DeepSeekHandler)
expect(OpenAI).toHaveBeenCalledWith({
baseURL: 'https://api.deepseek.com/v1',
apiKey: mockOptions.deepSeekApiKey,
})
})
describe('constructor', () => {
it('should initialize with provided options', () => {
expect(handler).toBeInstanceOf(DeepSeekHandler);
expect(handler.getModel().id).toBe(mockOptions.deepSeekModelId);
});
test('getModel returns correct model info', () => {
const handler = new DeepSeekHandler(mockOptions)
const result = handler.getModel()
expect(result).toEqual({
id: mockOptions.deepSeekModelId,
info: expect.objectContaining({
maxTokens: 8192,
contextWindow: 64000,
supportsPromptCache: false,
supportsImages: false,
inputPrice: 0.014,
outputPrice: 0.28,
})
})
})
it('should throw error if API key is missing', () => {
expect(() => {
new DeepSeekHandler({
...mockOptions,
deepSeekApiKey: undefined
});
}).toThrow('DeepSeek API key is required');
});
test('getModel returns default model info when no model specified', () => {
const handler = new DeepSeekHandler({ deepSeekApiKey: 'test-key' })
const result = handler.getModel()
expect(result.id).toBe('deepseek-chat')
expect(result.info.maxTokens).toBe(8192)
})
it('should use default model ID if not provided', () => {
const handlerWithoutModel = new DeepSeekHandler({
...mockOptions,
deepSeekModelId: undefined
});
expect(handlerWithoutModel.getModel().id).toBe(deepSeekDefaultModelId);
});
test('createMessage handles string content correctly', async () => {
const handler = new DeepSeekHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
choices: [{
delta: {
content: 'test response'
}
}]
}
}
}
it('should use default base URL if not provided', () => {
const handlerWithoutBaseUrl = new DeepSeekHandler({
...mockOptions,
deepSeekBaseUrl: undefined
});
expect(handlerWithoutBaseUrl).toBeInstanceOf(DeepSeekHandler);
// The base URL is passed to OpenAI client internally
expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({
baseURL: 'https://api.deepseek.com/v1'
}));
});
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
it('should use custom base URL if provided', () => {
const customBaseUrl = 'https://custom.deepseek.com/v1';
const handlerWithCustomUrl = new DeepSeekHandler({
...mockOptions,
deepSeekBaseUrl: customBaseUrl
});
expect(handlerWithCustomUrl).toBeInstanceOf(DeepSeekHandler);
// The custom base URL is passed to OpenAI client
expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({
baseURL: customBaseUrl
}));
});
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [
{ role: 'user', content: 'test message' }
]
it('should set includeMaxTokens to true', () => {
// Create a new handler and verify OpenAI client was called with includeMaxTokens
new DeepSeekHandler(mockOptions);
expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({
apiKey: mockOptions.deepSeekApiKey
}));
});
});
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
for await (const chunk of generator) {
chunks.push(chunk)
}
describe('getModel', () => {
it('should return model info for valid model ID', () => {
const model = handler.getModel();
expect(model.id).toBe(mockOptions.deepSeekModelId);
expect(model.info).toBeDefined();
expect(model.info.maxTokens).toBe(8192);
expect(model.info.contextWindow).toBe(64_000);
expect(model.info.supportsImages).toBe(false);
expect(model.info.supportsPromptCache).toBe(false);
});
expect(chunks).toHaveLength(1)
expect(chunks[0]).toEqual({
type: 'text',
text: 'test response'
})
it('should return provided model ID with default model info if model does not exist', () => {
const handlerWithInvalidModel = new DeepSeekHandler({
...mockOptions,
deepSeekModelId: 'invalid-model'
});
const model = handlerWithInvalidModel.getModel();
expect(model.id).toBe('invalid-model'); // Returns provided ID
expect(model.info).toBeDefined();
expect(model.info).toBe(handler.getModel().info); // But uses default model info
});
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({
model: mockOptions.deepSeekModelId,
messages: [
{ role: 'system', content: systemPrompt },
{ role: 'user', content: 'test message' }
],
temperature: 0,
stream: true,
max_tokens: 8192,
stream_options: { include_usage: true }
}))
})
it('should return default model if no model ID is provided', () => {
const handlerWithoutModel = new DeepSeekHandler({
...mockOptions,
deepSeekModelId: undefined
});
const model = handlerWithoutModel.getModel();
expect(model.id).toBe(deepSeekDefaultModelId);
expect(model.info).toBeDefined();
});
});
test('createMessage handles complex content correctly', async () => {
const handler = new DeepSeekHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
choices: [{
delta: {
content: 'test response'
}
}]
}
}
}
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
const systemPrompt = 'test system prompt'
describe('createMessage', () => {
const systemPrompt = 'You are a helpful assistant.';
const messages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: [
{ type: 'text', text: 'part 1' },
{ type: 'text', text: 'part 2' }
]
content: [{
type: 'text' as const,
text: 'Hello!'
}]
}
]
];
const generator = handler.createMessage(systemPrompt, messages)
await generator.next()
expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({
messages: [
{ role: 'system', content: systemPrompt },
{
role: 'user',
content: [
{ type: 'text', text: 'part 1' },
{ type: 'text', text: 'part 2' }
]
}
]
}))
})
test('createMessage handles API errors', async () => {
const handler = new DeepSeekHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
throw new Error('API Error')
it('should handle streaming responses', async () => {
const stream = handler.createMessage(systemPrompt, messages);
const chunks: any[] = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
}
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
expect(chunks.length).toBeGreaterThan(0);
const textChunks = chunks.filter(chunk => chunk.type === 'text');
expect(textChunks).toHaveLength(1);
expect(textChunks[0].text).toBe('Test response');
});
const generator = handler.createMessage('test', [])
await expect(generator.next()).rejects.toThrow('API Error')
})
})
it('should include usage information', async () => {
const stream = handler.createMessage(systemPrompt, messages);
const chunks: any[] = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
const usageChunks = chunks.filter(chunk => chunk.type === 'usage');
expect(usageChunks.length).toBeGreaterThan(0);
expect(usageChunks[0].inputTokens).toBe(10);
expect(usageChunks[0].outputTokens).toBe(5);
});
});
});

154
src/api/providers/__tests__/gemini.test.ts Normal file
View File

@@ -0,0 +1,154 @@
import { GeminiHandler } from '../gemini';
import { Anthropic } from '@anthropic-ai/sdk';
import { GoogleGenerativeAI } from '@google/generative-ai';
// Mock the Google Generative AI SDK
jest.mock('@google/generative-ai', () => ({
GoogleGenerativeAI: jest.fn().mockImplementation(() => ({
getGenerativeModel: jest.fn().mockReturnValue({
generateContentStream: jest.fn()
})
}))
}));
describe('GeminiHandler', () => {
let handler: GeminiHandler;
beforeEach(() => {
handler = new GeminiHandler({
apiKey: 'test-key',
apiModelId: 'gemini-2.0-flash-thinking-exp-1219',
geminiApiKey: 'test-key'
});
});
describe('constructor', () => {
it('should initialize with provided config', () => {
expect(handler['options'].geminiApiKey).toBe('test-key');
expect(handler['options'].apiModelId).toBe('gemini-2.0-flash-thinking-exp-1219');
});
it('should throw if API key is missing', () => {
expect(() => {
new GeminiHandler({
apiModelId: 'gemini-2.0-flash-thinking-exp-1219',
geminiApiKey: ''
});
}).toThrow('API key is required for Google Gemini');
});
});
describe('createMessage', () => {
const mockMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: 'Hello'
},
{
role: 'assistant',
content: 'Hi there!'
}
];
const systemPrompt = 'You are a helpful assistant';
it('should handle text messages correctly', async () => {
// Mock the stream response
const mockStream = {
stream: [
{ text: () => 'Hello' },
{ text: () => ' world!' }
],
response: {
usageMetadata: {
promptTokenCount: 10,
candidatesTokenCount: 5
}
}
};
// Setup the mock implementation
const mockGenerateContentStream = jest.fn().mockResolvedValue(mockStream);
const mockGetGenerativeModel = jest.fn().mockReturnValue({
generateContentStream: mockGenerateContentStream
});
(handler['client'] as any).getGenerativeModel = mockGetGenerativeModel;
const stream = handler.createMessage(systemPrompt, mockMessages);
const chunks = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
// Should have 3 chunks: 'Hello', ' world!', and usage info
expect(chunks.length).toBe(3);
expect(chunks[0]).toEqual({
type: 'text',
text: 'Hello'
});
expect(chunks[1]).toEqual({
type: 'text',
text: ' world!'
});
expect(chunks[2]).toEqual({
type: 'usage',
inputTokens: 10,
outputTokens: 5
});
// Verify the model configuration
expect(mockGetGenerativeModel).toHaveBeenCalledWith({
model: 'gemini-2.0-flash-thinking-exp-1219',
systemInstruction: systemPrompt
});
// Verify generation config
expect(mockGenerateContentStream).toHaveBeenCalledWith(
expect.objectContaining({
generationConfig: {
temperature: 0
}
})
);
});
it('should handle API errors', async () => {
const mockError = new Error('Gemini API error');
const mockGenerateContentStream = jest.fn().mockRejectedValue(mockError);
const mockGetGenerativeModel = jest.fn().mockReturnValue({
generateContentStream: mockGenerateContentStream
});
(handler['client'] as any).getGenerativeModel = mockGetGenerativeModel;
const stream = handler.createMessage(systemPrompt, mockMessages);
await expect(async () => {
for await (const chunk of stream) {
// Should throw before yielding any chunks
}
}).rejects.toThrow('Gemini API error');
});
});
describe('getModel', () => {
it('should return correct model info', () => {
const modelInfo = handler.getModel();
expect(modelInfo.id).toBe('gemini-2.0-flash-thinking-exp-1219');
expect(modelInfo.info).toBeDefined();
expect(modelInfo.info.maxTokens).toBe(8192);
expect(modelInfo.info.contextWindow).toBe(32_767);
});
it('should return default model if invalid model specified', () => {
const invalidHandler = new GeminiHandler({
apiModelId: 'invalid-model',
geminiApiKey: 'test-key'
});
const modelInfo = invalidHandler.getModel();
expect(modelInfo.id).toBe('gemini-2.0-flash-thinking-exp-1219'); // Default model
});
});
});

148
src/api/providers/__tests__/lmstudio.test.ts Normal file
View File

@@ -0,0 +1,148 @@
import { LmStudioHandler } from '../lmstudio';
import { Anthropic } from '@anthropic-ai/sdk';
import OpenAI from 'openai';
// Mock OpenAI SDK
jest.mock('openai', () => ({
__esModule: true,
default: jest.fn().mockImplementation(() => ({
chat: {
completions: {
create: jest.fn()
}
}
}))
}));
describe('LmStudioHandler', () => {
let handler: LmStudioHandler;
beforeEach(() => {
handler = new LmStudioHandler({
lmStudioModelId: 'mistral-7b',
lmStudioBaseUrl: 'http://localhost:1234'
});
});
describe('constructor', () => {
it('should initialize with provided config', () => {
expect(OpenAI).toHaveBeenCalledWith({
baseURL: 'http://localhost:1234/v1',
apiKey: 'noop'
});
});
it('should use default base URL if not provided', () => {
const defaultHandler = new LmStudioHandler({
lmStudioModelId: 'mistral-7b'
});
expect(OpenAI).toHaveBeenCalledWith({
baseURL: 'http://localhost:1234/v1',
apiKey: 'noop'
});
});
});
describe('createMessage', () => {
const mockMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: 'Hello'
},
{
role: 'assistant',
content: 'Hi there!'
}
];
const systemPrompt = 'You are a helpful assistant';
it('should handle streaming responses correctly', async () => {
const mockStream = [
{
choices: [{
delta: { content: 'Hello' }
}]
},
{
choices: [{
delta: { content: ' world!' }
}]
}
];
// Setup async iterator for mock stream
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk;
}
}
};
const mockCreate = jest.fn().mockResolvedValue(asyncIterator);
(handler['client'].chat.completions as any).create = mockCreate;
const stream = handler.createMessage(systemPrompt, mockMessages);
const chunks = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
expect(chunks.length).toBe(2);
expect(chunks[0]).toEqual({
type: 'text',
text: 'Hello'
});
expect(chunks[1]).toEqual({
type: 'text',
text: ' world!'
});
expect(mockCreate).toHaveBeenCalledWith({
model: 'mistral-7b',
messages: expect.arrayContaining([
{
role: 'system',
content: systemPrompt
}
]),
temperature: 0,
stream: true
});
});
it('should handle API errors with custom message', async () => {
const mockError = new Error('LM Studio API error');
const mockCreate = jest.fn().mockRejectedValue(mockError);
(handler['client'].chat.completions as any).create = mockCreate;
const stream = handler.createMessage(systemPrompt, mockMessages);
await expect(async () => {
for await (const chunk of stream) {
// Should throw before yielding any chunks
}
}).rejects.toThrow('Please check the LM Studio developer logs to debug what went wrong');
});
});
describe('getModel', () => {
it('should return model info with sane defaults', () => {
const modelInfo = handler.getModel();
expect(modelInfo.id).toBe('mistral-7b');
expect(modelInfo.info).toBeDefined();
expect(modelInfo.info.maxTokens).toBe(-1);
expect(modelInfo.info.contextWindow).toBe(128_000);
});
it('should return empty string as model ID if not provided', () => {
const noModelHandler = new LmStudioHandler({});
const modelInfo = noModelHandler.getModel();
expect(modelInfo.id).toBe('');
expect(modelInfo.info).toBeDefined();
});
});
});

148
src/api/providers/__tests__/ollama.test.ts Normal file
View File

@@ -0,0 +1,148 @@
import { OllamaHandler } from '../ollama';
import { Anthropic } from '@anthropic-ai/sdk';
import OpenAI from 'openai';
// Mock OpenAI SDK
jest.mock('openai', () => ({
__esModule: true,
default: jest.fn().mockImplementation(() => ({
chat: {
completions: {
create: jest.fn()
}
}
}))
}));
describe('OllamaHandler', () => {
let handler: OllamaHandler;
beforeEach(() => {
handler = new OllamaHandler({
ollamaModelId: 'llama2',
ollamaBaseUrl: 'http://localhost:11434'
});
});
describe('constructor', () => {
it('should initialize with provided config', () => {
expect(OpenAI).toHaveBeenCalledWith({
baseURL: 'http://localhost:11434/v1',
apiKey: 'ollama'
});
});
it('should use default base URL if not provided', () => {
const defaultHandler = new OllamaHandler({
ollamaModelId: 'llama2'
});
expect(OpenAI).toHaveBeenCalledWith({
baseURL: 'http://localhost:11434/v1',
apiKey: 'ollama'
});
});
});
describe('createMessage', () => {
const mockMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: 'Hello'
},
{
role: 'assistant',
content: 'Hi there!'
}
];
const systemPrompt = 'You are a helpful assistant';
it('should handle streaming responses correctly', async () => {
const mockStream = [
{
choices: [{
delta: { content: 'Hello' }
}]
},
{
choices: [{
delta: { content: ' world!' }
}]
}
];
// Setup async iterator for mock stream
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk;
}
}
};
const mockCreate = jest.fn().mockResolvedValue(asyncIterator);
(handler['client'].chat.completions as any).create = mockCreate;
const stream = handler.createMessage(systemPrompt, mockMessages);
const chunks = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
expect(chunks.length).toBe(2);
expect(chunks[0]).toEqual({
type: 'text',
text: 'Hello'
});
expect(chunks[1]).toEqual({
type: 'text',
text: ' world!'
});
expect(mockCreate).toHaveBeenCalledWith({
model: 'llama2',
messages: expect.arrayContaining([
{
role: 'system',
content: systemPrompt
}
]),
temperature: 0,
stream: true
});
});
it('should handle API errors', async () => {
const mockError = new Error('Ollama API error');
const mockCreate = jest.fn().mockRejectedValue(mockError);
(handler['client'].chat.completions as any).create = mockCreate;
const stream = handler.createMessage(systemPrompt, mockMessages);
await expect(async () => {
for await (const chunk of stream) {
// Should throw before yielding any chunks
}
}).rejects.toThrow('Ollama API error');
});
});
describe('getModel', () => {
it('should return model info with sane defaults', () => {
const modelInfo = handler.getModel();
expect(modelInfo.id).toBe('llama2');
expect(modelInfo.info).toBeDefined();
expect(modelInfo.info.maxTokens).toBe(-1);
expect(modelInfo.info.contextWindow).toBe(128_000);
});
it('should return empty string as model ID if not provided', () => {
const noModelHandler = new OllamaHandler({});
const modelInfo = noModelHandler.getModel();
expect(modelInfo.id).toBe('');
expect(modelInfo.info).toBeDefined();
});
});
});

230
src/api/providers/__tests__/openai-native.test.ts Normal file
View File

@@ -0,0 +1,230 @@
import { OpenAiNativeHandler } from "../openai-native"
import OpenAI from "openai"
import { ApiHandlerOptions, openAiNativeDefaultModelId } from "../../../shared/api"
import { Anthropic } from "@anthropic-ai/sdk"
// Mock OpenAI
jest.mock("openai")
describe("OpenAiNativeHandler", () => {
let handler: OpenAiNativeHandler
let mockOptions: ApiHandlerOptions
let mockOpenAIClient: jest.Mocked<OpenAI>
let mockCreate: jest.Mock
beforeEach(() => {
// Reset mocks
jest.clearAllMocks()
// Setup mock options
mockOptions = {
openAiNativeApiKey: "test-api-key",
apiModelId: "gpt-4o", // Use the correct model ID from shared/api.ts
}
// Setup mock create function
mockCreate = jest.fn()
// Setup mock OpenAI client
mockOpenAIClient = {
chat: {
completions: {
create: mockCreate,
},
},
} as unknown as jest.Mocked<OpenAI>
// Mock OpenAI constructor
;(OpenAI as jest.MockedClass<typeof OpenAI>).mockImplementation(() => mockOpenAIClient)
// Create handler instance
handler = new OpenAiNativeHandler(mockOptions)
})
describe("constructor", () => {
it("should initialize with provided options", () => {
expect(OpenAI).toHaveBeenCalledWith({
apiKey: mockOptions.openAiNativeApiKey,
})
})
})
describe("getModel", () => {
it("should return specified model when valid", () => {
const result = handler.getModel()
expect(result.id).toBe("gpt-4o") // Use the correct model ID
})
it("should return default model when model ID is invalid", () => {
handler = new OpenAiNativeHandler({
...mockOptions,
apiModelId: "invalid-model" as any,
})
const result = handler.getModel()
expect(result.id).toBe(openAiNativeDefaultModelId)
})
it("should return default model when model ID is not provided", () => {
handler = new OpenAiNativeHandler({
...mockOptions,
apiModelId: undefined,
})
const result = handler.getModel()
expect(result.id).toBe(openAiNativeDefaultModelId)
})
})
describe("createMessage", () => {
const systemPrompt = "You are a helpful assistant"
const messages: Anthropic.Messages.MessageParam[] = [
{ role: "user", content: "Hello" },
]
describe("o1 models", () => {
beforeEach(() => {
handler = new OpenAiNativeHandler({
...mockOptions,
apiModelId: "o1-preview",
})
})
it("should handle non-streaming response for o1 models", async () => {
const mockResponse = {
choices: [{ message: { content: "Hello there!" } }],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
},
}
mockCreate.mockResolvedValueOnce(mockResponse)
const generator = handler.createMessage(systemPrompt, messages)
const results = []
for await (const result of generator) {
results.push(result)
}
expect(results).toEqual([
{ type: "text", text: "Hello there!" },
{ type: "usage", inputTokens: 10, outputTokens: 5 },
])
expect(mockCreate).toHaveBeenCalledWith({
model: "o1-preview",
messages: [
{ role: "user", content: systemPrompt },
{ role: "user", content: "Hello" },
],
})
})
it("should handle missing content in response", async () => {
const mockResponse = {
choices: [{ message: { content: null } }],
usage: null,
}
mockCreate.mockResolvedValueOnce(mockResponse)
const generator = handler.createMessage(systemPrompt, messages)
const results = []
for await (const result of generator) {
results.push(result)
}
expect(results).toEqual([
{ type: "text", text: "" },
{ type: "usage", inputTokens: 0, outputTokens: 0 },
])
})
})
describe("streaming models", () => {
beforeEach(() => {
handler = new OpenAiNativeHandler({
...mockOptions,
apiModelId: "gpt-4o",
})
})
it("should handle streaming response", async () => {
const mockStream = [
{ choices: [{ delta: { content: "Hello" } }], usage: null },
{ choices: [{ delta: { content: " there" } }], usage: null },
{ choices: [{ delta: { content: "!" } }], usage: { prompt_tokens: 10, completion_tokens: 5 } },
]
mockCreate.mockResolvedValueOnce(
(async function* () {
for (const chunk of mockStream) {
yield chunk
}
})()
)
const generator = handler.createMessage(systemPrompt, messages)
const results = []
for await (const result of generator) {
results.push(result)
}
expect(results).toEqual([
{ type: "text", text: "Hello" },
{ type: "text", text: " there" },
{ type: "text", text: "!" },
{ type: "usage", inputTokens: 10, outputTokens: 5 },
])
expect(mockCreate).toHaveBeenCalledWith({
model: "gpt-4o",
temperature: 0,
messages: [
{ role: "system", content: systemPrompt },
{ role: "user", content: "Hello" },
],
stream: true,
stream_options: { include_usage: true },
})
})
it("should handle empty delta content", async () => {
const mockStream = [
{ choices: [{ delta: {} }], usage: null },
{ choices: [{ delta: { content: null } }], usage: null },
{ choices: [{ delta: { content: "Hello" } }], usage: { prompt_tokens: 10, completion_tokens: 5 } },
]
mockCreate.mockResolvedValueOnce(
(async function* () {
for (const chunk of mockStream) {
yield chunk
}
})()
)
const generator = handler.createMessage(systemPrompt, messages)
const results = []
for await (const result of generator) {
results.push(result)
}
expect(results).toEqual([
{ type: "text", text: "Hello" },
{ type: "usage", inputTokens: 10, outputTokens: 5 },
])
})
})
it("should handle API errors", async () => {
mockCreate.mockRejectedValueOnce(new Error("API Error"))
const generator = handler.createMessage(systemPrompt, messages)
await expect(async () => {
for await (const _ of generator) {
// consume generator
}
}).rejects.toThrow("API Error")
})
})
})

352
src/api/providers/__tests__/openai.test.ts
View File

@@ -1,192 +1,198 @@
import { OpenAiHandler } from '../openai'
import { ApiHandlerOptions, openAiModelInfoSaneDefaults } from '../../../shared/api'
import OpenAI, { AzureOpenAI } from 'openai'
import { Anthropic } from '@anthropic-ai/sdk'
import { OpenAiHandler } from '../openai';
import { ApiHandlerOptions } from '../../../shared/api';
import { ApiStream } from '../../transform/stream';
import OpenAI from 'openai';
import { Anthropic } from '@anthropic-ai/sdk';
// Mock dependencies
jest.mock('openai')
// Mock OpenAI client
const mockCreate = jest.fn();
jest.mock('openai', () => {
return {
__esModule: true,
default: jest.fn().mockImplementation(() => ({
chat: {
completions: {
create: mockCreate.mockImplementation(async (options) => {
if (!options.stream) {
return {
id: 'test-completion',
choices: [{
message: { role: 'assistant', content: 'Test response', refusal: null },
finish_reason: 'stop',
index: 0
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15
}
};
}
return {
[Symbol.asyncIterator]: async function* () {
yield {
choices: [{
delta: { content: 'Test response' },
index: 0
}],
usage: null
};
yield {
choices: [{
delta: {},
index: 0
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15
}
};
}
};
})
}
}
}))
};
});
describe('OpenAiHandler', () => {
const mockOptions: ApiHandlerOptions = {
openAiApiKey: 'test-key',
openAiModelId: 'gpt-4',
openAiStreamingEnabled: true,
openAiBaseUrl: 'https://api.openai.com/v1'
}
let handler: OpenAiHandler;
let mockOptions: ApiHandlerOptions;
beforeEach(() => {
jest.clearAllMocks()
})
mockOptions = {
openAiApiKey: 'test-api-key',
openAiModelId: 'gpt-4',
openAiBaseUrl: 'https://api.openai.com/v1'
};
handler = new OpenAiHandler(mockOptions);
mockCreate.mockClear();
});
test('constructor initializes with correct options', () => {
const handler = new OpenAiHandler(mockOptions)
expect(handler).toBeInstanceOf(OpenAiHandler)
expect(OpenAI).toHaveBeenCalledWith({
apiKey: mockOptions.openAiApiKey,
baseURL: mockOptions.openAiBaseUrl
})
})
describe('constructor', () => {
it('should initialize with provided options', () => {
expect(handler).toBeInstanceOf(OpenAiHandler);
expect(handler.getModel().id).toBe(mockOptions.openAiModelId);
});
test('constructor initializes Azure client when Azure URL is provided', () => {
const azureOptions: ApiHandlerOptions = {
...mockOptions,
openAiBaseUrl: 'https://example.azure.com',
azureApiVersion: '2023-05-15'
}
const handler = new OpenAiHandler(azureOptions)
expect(handler).toBeInstanceOf(OpenAiHandler)
expect(AzureOpenAI).toHaveBeenCalledWith({
baseURL: azureOptions.openAiBaseUrl,
apiKey: azureOptions.openAiApiKey,
apiVersion: azureOptions.azureApiVersion
})
})
it('should use custom base URL if provided', () => {
const customBaseUrl = 'https://custom.openai.com/v1';
const handlerWithCustomUrl = new OpenAiHandler({
...mockOptions,
openAiBaseUrl: customBaseUrl
});
expect(handlerWithCustomUrl).toBeInstanceOf(OpenAiHandler);
});
});
test('getModel returns correct model info', () => {
const handler = new OpenAiHandler(mockOptions)
const result = handler.getModel()
expect(result).toEqual({
id: mockOptions.openAiModelId,
info: openAiModelInfoSaneDefaults
})
})
test('createMessage handles streaming correctly when enabled', async () => {
const handler = new OpenAiHandler({
...mockOptions,
openAiStreamingEnabled: true,
includeMaxTokens: true
})
const mockStream = {
async *[Symbol.asyncIterator]() {
yield {
choices: [{
delta: {
content: 'test response'
}
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5
}
}
}
}
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
const systemPrompt = 'test system prompt'
describe('createMessage', () => {
const systemPrompt = 'You are a helpful assistant.';
const messages: Anthropic.Messages.MessageParam[] = [
{ role: 'user', content: 'test message' }
]
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
for await (const chunk of generator) {
chunks.push(chunk)
}
expect(chunks).toEqual([
{
type: 'text',
text: 'test response'
},
{
type: 'usage',
inputTokens: 10,
outputTokens: 5
role: 'user',
content: [{
type: 'text' as const,
text: 'Hello!'
}]
}
])
];
expect(mockCreate).toHaveBeenCalledWith({
model: mockOptions.openAiModelId,
messages: [
{ role: 'system', content: systemPrompt },
{ role: 'user', content: 'test message' }
],
temperature: 0,
stream: true,
stream_options: { include_usage: true },
max_tokens: openAiModelInfoSaneDefaults.maxTokens
})
})
it('should handle non-streaming mode', async () => {
const handler = new OpenAiHandler({
...mockOptions,
openAiStreamingEnabled: false
});
test('createMessage handles non-streaming correctly when disabled', async () => {
const handler = new OpenAiHandler({
...mockOptions,
openAiStreamingEnabled: false
})
const mockResponse = {
choices: [{
message: {
content: 'test response'
const stream = handler.createMessage(systemPrompt, messages);
const chunks: any[] = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
expect(chunks.length).toBeGreaterThan(0);
const textChunk = chunks.find(chunk => chunk.type === 'text');
const usageChunk = chunks.find(chunk => chunk.type === 'usage');
expect(textChunk).toBeDefined();
expect(textChunk?.text).toBe('Test response');
expect(usageChunk).toBeDefined();
expect(usageChunk?.inputTokens).toBe(10);
expect(usageChunk?.outputTokens).toBe(5);
});
it('should handle streaming responses', async () => {
const stream = handler.createMessage(systemPrompt, messages);
const chunks: any[] = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
expect(chunks.length).toBeGreaterThan(0);
const textChunks = chunks.filter(chunk => chunk.type === 'text');
expect(textChunks).toHaveLength(1);
expect(textChunks[0].text).toBe('Test response');
});
});
describe('error handling', () => {
const testMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: [{
type: 'text' as const,
text: 'Hello'
}]
}
];
it('should handle API errors', async () => {
mockCreate.mockRejectedValueOnce(new Error('API Error'));
const stream = handler.createMessage('system prompt', testMessages);
await expect(async () => {
for await (const chunk of stream) {
// Should not reach here
}
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5
}
}
}).rejects.toThrow('API Error');
});
const mockCreate = jest.fn().mockResolvedValue(mockResponse)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
it('should handle rate limiting', async () => {
const rateLimitError = new Error('Rate limit exceeded');
rateLimitError.name = 'Error';
(rateLimitError as any).status = 429;
mockCreate.mockRejectedValueOnce(rateLimitError);
const systemPrompt = 'test system prompt'
const messages: Anthropic.Messages.MessageParam[] = [
{ role: 'user', content: 'test message' }
]
const stream = handler.createMessage('system prompt', testMessages);
const generator = handler.createMessage(systemPrompt, messages)
const chunks = []
for await (const chunk of generator) {
chunks.push(chunk)
}
await expect(async () => {
for await (const chunk of stream) {
// Should not reach here
}
}).rejects.toThrow('Rate limit exceeded');
});
});
expect(chunks).toEqual([
{
type: 'text',
text: 'test response'
},
{
type: 'usage',
inputTokens: 10,
outputTokens: 5
}
])
describe('getModel', () => {
it('should return model info with sane defaults', () => {
const model = handler.getModel();
expect(model.id).toBe(mockOptions.openAiModelId);
expect(model.info).toBeDefined();
expect(model.info.contextWindow).toBe(128_000);
expect(model.info.supportsImages).toBe(true);
});
expect(mockCreate).toHaveBeenCalledWith({
model: mockOptions.openAiModelId,
messages: [
{ role: 'user', content: systemPrompt },
{ role: 'user', content: 'test message' }
]
})
})
test('createMessage handles API errors', async () => {
const handler = new OpenAiHandler(mockOptions)
const mockStream = {
async *[Symbol.asyncIterator]() {
throw new Error('API Error')
}
}
const mockCreate = jest.fn().mockResolvedValue(mockStream)
;(OpenAI as jest.MockedClass<typeof OpenAI>).prototype.chat = {
completions: { create: mockCreate }
} as any
const generator = handler.createMessage('test', [])
await expect(generator.next()).rejects.toThrow('API Error')
})
})
it('should handle undefined model ID', () => {
const handlerWithoutModel = new OpenAiHandler({
...mockOptions,
openAiModelId: undefined
});
const model = handlerWithoutModel.getModel();
expect(model.id).toBe('');
expect(model.info).toBeDefined();
});
});
});

218
src/api/providers/__tests__/vertex.test.ts Normal file
View File

@@ -0,0 +1,218 @@
import { VertexHandler } from '../vertex';
import { Anthropic } from '@anthropic-ai/sdk';
import { AnthropicVertex } from '@anthropic-ai/vertex-sdk';
// Mock Vertex SDK
jest.mock('@anthropic-ai/vertex-sdk', () => ({
AnthropicVertex: jest.fn().mockImplementation(() => ({
messages: {
create: jest.fn()
}
}))
}));
describe('VertexHandler', () => {
let handler: VertexHandler;
beforeEach(() => {
handler = new VertexHandler({
apiModelId: 'claude-3-5-sonnet-v2@20241022',
vertexProjectId: 'test-project',
vertexRegion: 'us-central1'
});
});
describe('constructor', () => {
it('should initialize with provided config', () => {
expect(AnthropicVertex).toHaveBeenCalledWith({
projectId: 'test-project',
region: 'us-central1'
});
});
});
describe('createMessage', () => {
const mockMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: 'Hello'
},
{
role: 'assistant',
content: 'Hi there!'
}
];
const systemPrompt = 'You are a helpful assistant';
it('should handle streaming responses correctly', async () => {
const mockStream = [
{
type: 'message_start',
message: {
usage: {
input_tokens: 10,
output_tokens: 0
}
}
},
{
type: 'content_block_start',
index: 0,
content_block: {
type: 'text',
text: 'Hello'
}
},
{
type: 'content_block_delta',
delta: {
type: 'text_delta',
text: ' world!'
}
},
{
type: 'message_delta',
usage: {
output_tokens: 5
}
}
];
// Setup async iterator for mock stream
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk;
}
}
};
const mockCreate = jest.fn().mockResolvedValue(asyncIterator);
(handler['client'].messages as any).create = mockCreate;
const stream = handler.createMessage(systemPrompt, mockMessages);
const chunks = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
expect(chunks.length).toBe(4);
expect(chunks[0]).toEqual({
type: 'usage',
inputTokens: 10,
outputTokens: 0
});
expect(chunks[1]).toEqual({
type: 'text',
text: 'Hello'
});
expect(chunks[2]).toEqual({
type: 'text',
text: ' world!'
});
expect(chunks[3]).toEqual({
type: 'usage',
inputTokens: 0,
outputTokens: 5
});
expect(mockCreate).toHaveBeenCalledWith({
model: 'claude-3-5-sonnet-v2@20241022',
max_tokens: 8192,
temperature: 0,
system: systemPrompt,
messages: mockMessages,
stream: true
});
});
it('should handle multiple content blocks with line breaks', async () => {
const mockStream = [
{
type: 'content_block_start',
index: 0,
content_block: {
type: 'text',
text: 'First line'
}
},
{
type: 'content_block_start',
index: 1,
content_block: {
type: 'text',
text: 'Second line'
}
}
];
const asyncIterator = {
async *[Symbol.asyncIterator]() {
for (const chunk of mockStream) {
yield chunk;
}
}
};
const mockCreate = jest.fn().mockResolvedValue(asyncIterator);
(handler['client'].messages as any).create = mockCreate;
const stream = handler.createMessage(systemPrompt, mockMessages);
const chunks = [];
for await (const chunk of stream) {
chunks.push(chunk);
}
expect(chunks.length).toBe(3);
expect(chunks[0]).toEqual({
type: 'text',
text: 'First line'
});
expect(chunks[1]).toEqual({
type: 'text',
text: '\n'
});
expect(chunks[2]).toEqual({
type: 'text',
text: 'Second line'
});
});
it('should handle API errors', async () => {
const mockError = new Error('Vertex API error');
const mockCreate = jest.fn().mockRejectedValue(mockError);
(handler['client'].messages as any).create = mockCreate;
const stream = handler.createMessage(systemPrompt, mockMessages);
await expect(async () => {
for await (const chunk of stream) {
// Should throw before yielding any chunks
}
}).rejects.toThrow('Vertex API error');
});
});
describe('getModel', () => {
it('should return correct model info', () => {
const modelInfo = handler.getModel();
expect(modelInfo.id).toBe('claude-3-5-sonnet-v2@20241022');
expect(modelInfo.info).toBeDefined();
expect(modelInfo.info.maxTokens).toBe(8192);
expect(modelInfo.info.contextWindow).toBe(200_000);
});
it('should return default model if invalid model specified', () => {
const invalidHandler = new VertexHandler({
apiModelId: 'invalid-model',
vertexProjectId: 'test-project',
vertexRegion: 'us-central1'
});
const modelInfo = invalidHandler.getModel();
expect(modelInfo.id).toBe('claude-3-5-sonnet-v2@20241022'); // Default model
});
});
});

257
src/api/transform/__tests__/openai-format.test.ts Normal file
View File

@@ -0,0 +1,257 @@
import { convertToOpenAiMessages, convertToAnthropicMessage } from '../openai-format';
import { Anthropic } from '@anthropic-ai/sdk';
import OpenAI from 'openai';
type PartialChatCompletion = Omit<OpenAI.Chat.Completions.ChatCompletion, 'choices'> & {
choices: Array<Partial<OpenAI.Chat.Completions.ChatCompletion.Choice> & {
message: OpenAI.Chat.Completions.ChatCompletion.Choice['message'];
finish_reason: string;
index: number;
}>;
};
describe('OpenAI Format Transformations', () => {
describe('convertToOpenAiMessages', () => {
it('should convert simple text messages', () => {
const anthropicMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: 'Hello'
},
{
role: 'assistant',
content: 'Hi there!'
}
];
const openAiMessages = convertToOpenAiMessages(anthropicMessages);
expect(openAiMessages).toHaveLength(2);
expect(openAiMessages[0]).toEqual({
role: 'user',
content: 'Hello'
});
expect(openAiMessages[1]).toEqual({
role: 'assistant',
content: 'Hi there!'
});
});
it('should handle messages with image content', () => {
const anthropicMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: [
{
type: 'text',
text: 'What is in this image?'
},
{
type: 'image',
source: {
type: 'base64',
media_type: 'image/jpeg',
data: 'base64data'
}
}
]
}
];
const openAiMessages = convertToOpenAiMessages(anthropicMessages);
expect(openAiMessages).toHaveLength(1);
expect(openAiMessages[0].role).toBe('user');
const content = openAiMessages[0].content as Array<{
type: string;
text?: string;
image_url?: { url: string };
}>;
expect(Array.isArray(content)).toBe(true);
expect(content).toHaveLength(2);
expect(content[0]).toEqual({ type: 'text', text: 'What is in this image?' });
expect(content[1]).toEqual({
type: 'image_url',
image_url: { url: 'data:image/jpeg;base64,base64data' }
});
});
it('should handle assistant messages with tool use', () => {
const anthropicMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'assistant',
content: [
{
type: 'text',
text: 'Let me check the weather.'
},
{
type: 'tool_use',
id: 'weather-123',
name: 'get_weather',
input: { city: 'London' }
}
]
}
];
const openAiMessages = convertToOpenAiMessages(anthropicMessages);
expect(openAiMessages).toHaveLength(1);
const assistantMessage = openAiMessages[0] as OpenAI.Chat.ChatCompletionAssistantMessageParam;
expect(assistantMessage.role).toBe('assistant');
expect(assistantMessage.content).toBe('Let me check the weather.');
expect(assistantMessage.tool_calls).toHaveLength(1);
expect(assistantMessage.tool_calls![0]).toEqual({
id: 'weather-123',
type: 'function',
function: {
name: 'get_weather',
arguments: JSON.stringify({ city: 'London' })
}
});
});
it('should handle user messages with tool results', () => {
const anthropicMessages: Anthropic.Messages.MessageParam[] = [
{
role: 'user',
content: [
{
type: 'tool_result',
tool_use_id: 'weather-123',
content: 'Current temperature in London: 20°C'
}
]
}
];
const openAiMessages = convertToOpenAiMessages(anthropicMessages);
expect(openAiMessages).toHaveLength(1);
const toolMessage = openAiMessages[0] as OpenAI.Chat.ChatCompletionToolMessageParam;
expect(toolMessage.role).toBe('tool');
expect(toolMessage.tool_call_id).toBe('weather-123');
expect(toolMessage.content).toBe('Current temperature in London: 20°C');
});
});
describe('convertToAnthropicMessage', () => {
it('should convert simple completion', () => {
const openAiCompletion: PartialChatCompletion = {
id: 'completion-123',
model: 'gpt-4',
choices: [{
message: {
role: 'assistant',
content: 'Hello there!',
refusal: null
},
finish_reason: 'stop',
index: 0
}],
usage: {
prompt_tokens: 10,
completion_tokens: 5,
total_tokens: 15
},
created: 123456789,
object: 'chat.completion'
};
const anthropicMessage = convertToAnthropicMessage(openAiCompletion as OpenAI.Chat.Completions.ChatCompletion);
expect(anthropicMessage.id).toBe('completion-123');
expect(anthropicMessage.role).toBe('assistant');
expect(anthropicMessage.content).toHaveLength(1);
expect(anthropicMessage.content[0]).toEqual({
type: 'text',
text: 'Hello there!'
});
expect(anthropicMessage.stop_reason).toBe('end_turn');
expect(anthropicMessage.usage).toEqual({
input_tokens: 10,
output_tokens: 5
});
});
it('should handle tool calls in completion', () => {
const openAiCompletion: PartialChatCompletion = {
id: 'completion-123',
model: 'gpt-4',
choices: [{
message: {
role: 'assistant',
content: 'Let me check the weather.',
tool_calls: [{
id: 'weather-123',
type: 'function',
function: {
name: 'get_weather',
arguments: '{"city":"London"}'
}
}],
refusal: null
},
finish_reason: 'tool_calls',
index: 0
}],
usage: {
prompt_tokens: 15,
completion_tokens: 8,
total_tokens: 23
},
created: 123456789,
object: 'chat.completion'
};
const anthropicMessage = convertToAnthropicMessage(openAiCompletion as OpenAI.Chat.Completions.ChatCompletion);
expect(anthropicMessage.content).toHaveLength(2);
expect(anthropicMessage.content[0]).toEqual({
type: 'text',
text: 'Let me check the weather.'
});
expect(anthropicMessage.content[1]).toEqual({
type: 'tool_use',
id: 'weather-123',
name: 'get_weather',
input: { city: 'London' }
});
expect(anthropicMessage.stop_reason).toBe('tool_use');
});
it('should handle invalid tool call arguments', () => {
const openAiCompletion: PartialChatCompletion = {
id: 'completion-123',
model: 'gpt-4',
choices: [{
message: {
role: 'assistant',
content: 'Testing invalid arguments',
tool_calls: [{
id: 'test-123',
type: 'function',
function: {
name: 'test_function',
arguments: 'invalid json'
}
}],
refusal: null
},
finish_reason: 'tool_calls',
index: 0
}],
created: 123456789,
object: 'chat.completion'
};
const anthropicMessage = convertToAnthropicMessage(openAiCompletion as OpenAI.Chat.Completions.ChatCompletion);
expect(anthropicMessage.content).toHaveLength(2);
expect(anthropicMessage.content[1]).toEqual({
type: 'tool_use',
id: 'test-123',
name: 'test_function',
input: {} // Should default to empty object for invalid JSON
});
});
});
});

114
src/api/transform/__tests__/stream.test.ts Normal file
View File

@@ -0,0 +1,114 @@
import { ApiStreamChunk } from '../stream';
describe('API Stream Types', () => {
describe('ApiStreamChunk', () => {
it('should correctly handle text chunks', () => {
const textChunk: ApiStreamChunk = {
type: 'text',
text: 'Hello world'
};
expect(textChunk.type).toBe('text');
expect(textChunk.text).toBe('Hello world');
});
it('should correctly handle usage chunks with cache information', () => {
const usageChunk: ApiStreamChunk = {
type: 'usage',
inputTokens: 100,
outputTokens: 50,
cacheWriteTokens: 20,
cacheReadTokens: 10
};
expect(usageChunk.type).toBe('usage');
expect(usageChunk.inputTokens).toBe(100);
expect(usageChunk.outputTokens).toBe(50);
expect(usageChunk.cacheWriteTokens).toBe(20);
expect(usageChunk.cacheReadTokens).toBe(10);
});
it('should handle usage chunks without cache tokens', () => {
const usageChunk: ApiStreamChunk = {
type: 'usage',
inputTokens: 100,
outputTokens: 50
};
expect(usageChunk.type).toBe('usage');
expect(usageChunk.inputTokens).toBe(100);
expect(usageChunk.outputTokens).toBe(50);
expect(usageChunk.cacheWriteTokens).toBeUndefined();
expect(usageChunk.cacheReadTokens).toBeUndefined();
});
it('should handle text chunks with empty strings', () => {
const emptyTextChunk: ApiStreamChunk = {
type: 'text',
text: ''
};
expect(emptyTextChunk.type).toBe('text');
expect(emptyTextChunk.text).toBe('');
});
it('should handle usage chunks with zero tokens', () => {
const zeroUsageChunk: ApiStreamChunk = {
type: 'usage',
inputTokens: 0,
outputTokens: 0
};
expect(zeroUsageChunk.type).toBe('usage');
expect(zeroUsageChunk.inputTokens).toBe(0);
expect(zeroUsageChunk.outputTokens).toBe(0);
});
it('should handle usage chunks with large token counts', () => {
const largeUsageChunk: ApiStreamChunk = {
type: 'usage',
inputTokens: 1000000,
outputTokens: 500000,
cacheWriteTokens: 200000,
cacheReadTokens: 100000
};
expect(largeUsageChunk.type).toBe('usage');
expect(largeUsageChunk.inputTokens).toBe(1000000);
expect(largeUsageChunk.outputTokens).toBe(500000);
expect(largeUsageChunk.cacheWriteTokens).toBe(200000);
expect(largeUsageChunk.cacheReadTokens).toBe(100000);
});
it('should handle text chunks with special characters', () => {
const specialCharsChunk: ApiStreamChunk = {
type: 'text',
text: '!@#$%^&*()_+-=[]{}|;:,.<>?`~'
};
expect(specialCharsChunk.type).toBe('text');
expect(specialCharsChunk.text).toBe('!@#$%^&*()_+-=[]{}|;:,.<>?`~');
});
it('should handle text chunks with unicode characters', () => {
const unicodeChunk: ApiStreamChunk = {
type: 'text',
text: '你好世界👋🌍'
};
expect(unicodeChunk.type).toBe('text');
expect(unicodeChunk.text).toBe('你好世界👋🌍');
});
it('should handle text chunks with multiline content', () => {
const multilineChunk: ApiStreamChunk = {
type: 'text',
text: 'Line 1\nLine 2\nLine 3'
};
expect(multilineChunk.type).toBe('text');
expect(multilineChunk.text).toBe('Line 1\nLine 2\nLine 3');
expect(multilineChunk.text.split('\n')).toHaveLength(3);
});
});
});