Add non-streaming completePrompt to all providers

This commit is contained in:
Matt Rubens
2025-01-13 16:16:58 -05:00
parent 2d176e5c92
commit 4027e1c10c
18 changed files with 1235 additions and 438 deletions

View File

@@ -6,7 +6,12 @@ import { GoogleGenerativeAI } from '@google/generative-ai';
jest.mock('@google/generative-ai', () => ({
GoogleGenerativeAI: jest.fn().mockImplementation(() => ({
getGenerativeModel: jest.fn().mockReturnValue({
generateContentStream: jest.fn()
generateContentStream: jest.fn(),
generateContent: jest.fn().mockResolvedValue({
response: {
text: () => 'Test response'
}
})
})
}))
}));
@@ -133,6 +138,59 @@ describe('GeminiHandler', () => {
});
});
describe('completePrompt', () => {
it('should complete prompt successfully', async () => {
const mockGenerateContent = jest.fn().mockResolvedValue({
response: {
text: () => 'Test response'
}
});
const mockGetGenerativeModel = jest.fn().mockReturnValue({
generateContent: mockGenerateContent
});
(handler['client'] as any).getGenerativeModel = mockGetGenerativeModel;
const result = await handler.completePrompt('Test prompt');
expect(result).toBe('Test response');
expect(mockGetGenerativeModel).toHaveBeenCalledWith({
model: 'gemini-2.0-flash-thinking-exp-1219'
});
expect(mockGenerateContent).toHaveBeenCalledWith({
contents: [{ role: 'user', parts: [{ text: 'Test prompt' }] }],
generationConfig: {
temperature: 0
}
});
});
it('should handle API errors', async () => {
const mockError = new Error('Gemini API error');
const mockGenerateContent = jest.fn().mockRejectedValue(mockError);
const mockGetGenerativeModel = jest.fn().mockReturnValue({
generateContent: mockGenerateContent
});
(handler['client'] as any).getGenerativeModel = mockGetGenerativeModel;
await expect(handler.completePrompt('Test prompt'))
.rejects.toThrow('Gemini completion error: Gemini API error');
});
it('should handle empty response', async () => {
const mockGenerateContent = jest.fn().mockResolvedValue({
response: {
text: () => ''
}
});
const mockGetGenerativeModel = jest.fn().mockReturnValue({
generateContent: mockGenerateContent
});
(handler['client'] as any).getGenerativeModel = mockGetGenerativeModel;
const result = await handler.completePrompt('Test prompt');
expect(result).toBe('');
});
});
describe('getModel', () => {
it('should return correct model info', () => {
const modelInfo = handler.getModel();