/** * @vitest-environment node */ import { beforeEach, describe, expect, it, vi } from 'vitest' const { mockCreate, mockExecuteTool } = vi.hoisted(() => ({ mockCreate: vi.fn(), mockExecuteTool: vi.fn(), })) vi.mock('openai', () => ({ default: vi.fn().mockImplementation( class { chat = { completions: { create: mockCreate } } } ), })) vi.mock('@/tools', () => ({ executeTool: mockExecuteTool })) vi.mock('@/providers', () => ({ MAX_TOOL_ITERATIONS: 20 })) vi.mock('@/lib/core/config/env', () => ({ env: { LITELLM_BASE_URL: 'http://litellm.test', LITELLM_API_KEY: '' }, })) vi.mock('@/stores/providers', () => ({ useProvidersStore: { getState: () => ({ setProviderModels: vi.fn() }) }, })) vi.mock('@/providers/models', () => ({ getProviderFileAttachment: vi .fn() .mockReturnValue({ maxBytes: 10 * 1024 * 1024, strategy: 'inline' }), INLINE_ATTACHMENT_MAX_BYTES: 10 * 1024 * 1024, getProviderModels: () => [], getProviderDefaultModel: () => '', })) vi.mock('@/providers/attachments', () => ({ formatMessagesForProvider: (messages: unknown) => messages, })) vi.mock('@/providers/trace-enrichment', () => ({ enrichLastModelSegmentFromChatCompletions: vi.fn(), })) vi.mock('@/providers/litellm/utils', () => ({ createReadableStreamFromLiteLLMStream: vi.fn( () => new ReadableStream({ start: (c) => c.close() }) ), })) vi.mock('@/providers/utils', () => ({ calculateCost: vi.fn(() => ({ input: 0, output: 0, total: 0 })), sumToolCosts: vi.fn(() => 0), prepareToolExecution: vi.fn((_tool, toolArgs) => ({ toolParams: toolArgs, executionParams: toolArgs, })), prepareToolsWithUsageControl: vi.fn((tools) => ({ tools, toolChoice: 'auto', forcedTools: [], hasFilteredTools: false, })), trackForcedToolUsage: vi.fn(() => ({ hasUsedForcedTool: false, usedForcedTools: [] })), enforceStrictSchema: vi.fn((schema) => ({ ...schema, additionalProperties: false })), })) import { litellmProvider } from '@/providers/litellm' import { ProviderError } from '@/providers/types' interface ChatOptions { content?: string | null toolCalls?: Array<{ id: string; function: { name: string; arguments: string } }> usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number } } function chat({ content = null, toolCalls, usage }: ChatOptions = {}) { return { choices: [ { message: { content, tool_calls: toolCalls }, finish_reason: toolCalls ? 'tool_calls' : 'stop', }, ], usage: usage ?? { prompt_tokens: 5, completion_tokens: 3, total_tokens: 8 }, } } function tool(name: string) { return { id: name, name, description: 'd', parameters: {} } } function run(request: Record) { return litellmProvider.executeRequest!({ model: 'litellm/llama-3', messages: [{ role: 'user', content: 'Hi' }], ...request, } as never) as Promise } const firstPayload = () => mockCreate.mock.calls[0][0] const lastPayload = () => mockCreate.mock.calls.at(-1)![0] describe('litellmProvider.executeRequest', () => { beforeEach(() => { vi.clearAllMocks() mockCreate.mockResolvedValue(chat({ content: 'hello' })) mockExecuteTool.mockResolvedValue({ success: true, output: { ok: true } }) }) it('assembles messages, strips the model prefix, and maps params', async () => { const result = await run({ systemPrompt: 'You are helpful.', context: 'Some context', temperature: 0.5, maxTokens: 256, }) const payload = firstPayload() expect(payload.model).toBe('llama-3') expect(payload.messages).toEqual([ { role: 'system', content: 'You are helpful.' }, { role: 'user', content: 'Some context' }, { role: 'user', content: 'Hi' }, ]) expect(payload.temperature).toBe(0.5) expect(payload.max_completion_tokens).toBe(256) expect(result.content).toBe('hello') expect(result.tokens).toEqual({ input: 5, output: 3, total: 8 }) }) it('forwards reasoning_effort only when set to a non-default value', async () => { await run({ reasoningEffort: 'high' }) expect(firstPayload().reasoning_effort).toBe('high') mockCreate.mockClear() await run({ reasoningEffort: 'auto' }) expect(firstPayload().reasoning_effort).toBeUndefined() mockCreate.mockClear() await run({}) expect(firstPayload().reasoning_effort).toBeUndefined() }) it('sanitizes the schema for strict response_format and passes it through otherwise', async () => { await run({ responseFormat: { name: 'r', schema: { type: 'object', properties: {} } } }) let rf = firstPayload().response_format expect(rf.type).toBe('json_schema') expect(rf.json_schema.strict).toBe(true) expect(rf.json_schema.schema.additionalProperties).toBe(false) mockCreate.mockClear() await run({ responseFormat: { name: 'r', schema: { type: 'object', properties: {} }, strict: false }, }) rf = firstPayload().response_format expect(rf.json_schema.strict).toBe(false) expect(rf.json_schema.schema.additionalProperties).toBeUndefined() }) it('defers response_format past the tool loop and keeps tools on the final call', async () => { mockCreate .mockResolvedValueOnce( chat({ toolCalls: [{ id: 'c1', function: { name: 'known', arguments: '{"q":1}' } }] }) ) .mockResolvedValueOnce(chat({ content: 'mid' })) .mockResolvedValueOnce(chat({ content: '{"answer":1}' })) const result = await run({ tools: [tool('known')], reasoningEffort: 'high', responseFormat: { name: 'r', schema: { type: 'object', properties: {} } }, }) expect(firstPayload().response_format).toBeUndefined() expect(firstPayload().tools).toBeDefined() const final = lastPayload() expect(final.response_format.type).toBe('json_schema') expect(final.tools).toBeDefined() expect(final.tool_choice).toBe('none') expect(final.parallel_tool_calls).toBe(false) expect(final.reasoning_effort).toBe('high') expect(result.content).toBe('{"answer":1}') }) it('defers response_format into the final streaming call while keeping tools', async () => { mockCreate .mockResolvedValueOnce( chat({ toolCalls: [{ id: 'c1', function: { name: 'known', arguments: '{}' } }] }) ) .mockResolvedValueOnce(chat({ content: 'mid' })) const result = await run({ stream: true, tools: [tool('known')], responseFormat: { name: 'r', schema: { type: 'object', properties: {} } }, }) const final = lastPayload() expect(final.stream).toBe(true) expect(final.response_format.type).toBe('json_schema') expect(final.tools).toBeDefined() expect(final.tool_choice).toBe('none') expect(final.parallel_tool_calls).toBe(false) expect(result.execution.isStreaming).toBe(true) }) it('threads assistant tool_calls and a named tool response, and reports toolCalls', async () => { mockCreate .mockResolvedValueOnce( chat({ toolCalls: [{ id: 'c1', function: { name: 'known', arguments: '{}' } }] }) ) .mockResolvedValueOnce(chat({ content: 'done' })) mockExecuteTool.mockResolvedValue({ success: true, output: { temp: 72 } }) const result = await run({ tools: [tool('known')] }) const followupMessages = mockCreate.mock.calls[1][0].messages expect(followupMessages).toContainEqual({ role: 'assistant', content: null, tool_calls: [{ id: 'c1', type: 'function', function: { name: 'known', arguments: '{}' } }], }) expect(followupMessages).toContainEqual({ role: 'tool', tool_call_id: 'c1', name: 'known', content: JSON.stringify({ temp: 72 }), }) expect(result.toolCalls).toHaveLength(1) expect(result.content).toBe('done') }) it('emits a stub tool response for an unanswered tool_call_id', async () => { mockCreate .mockResolvedValueOnce( chat({ toolCalls: [{ id: 'cX', function: { name: 'ghost', arguments: '{}' } }] }) ) .mockResolvedValueOnce(chat({ content: 'recovered' })) await run({ tools: [tool('known')] }) expect(mockExecuteTool).not.toHaveBeenCalled() const followupMessages = mockCreate.mock.calls[1][0].messages const toolMsg = followupMessages.find((m: any) => m.role === 'tool' && m.tool_call_id === 'cX') expect(toolMsg).toBeDefined() expect(toolMsg.content).toContain('not available') }) it('executes a tool with empty arguments without failing', async () => { mockCreate .mockResolvedValueOnce( chat({ toolCalls: [{ id: 'c1', function: { name: 'ping', arguments: '' } }] }) ) .mockResolvedValueOnce(chat({ content: 'pong' })) await run({ tools: [tool('ping')] }) expect(mockExecuteTool).toHaveBeenCalledTimes(1) const toolMsg = mockCreate.mock.calls[1][0].messages.find((m: any) => m.role === 'tool') expect(toolMsg.content).not.toContain('"error":true') }) it('stops the tool loop at MAX_TOOL_ITERATIONS', async () => { mockCreate.mockResolvedValue( chat({ toolCalls: [{ id: 'c1', function: { name: 'known', arguments: '{}' } }] }) ) await run({ tools: [tool('known')] }) expect(mockCreate).toHaveBeenCalledTimes(1 + 20) expect(mockExecuteTool).toHaveBeenCalledTimes(20) }) it('returns a streaming execution when streaming without active tools', async () => { const result = await run({ stream: true }) expect(firstPayload().stream).toBe(true) expect(firstPayload().stream_options).toEqual({ include_usage: true }) expect(result.stream).toBeInstanceOf(ReadableStream) expect(result.execution.isStreaming).toBe(true) }) it('wraps API errors in a ProviderError using the error envelope message', async () => { mockCreate.mockRejectedValue({ error: { message: 'rate limited', type: 'rate_limit_error', code: '429' }, }) await expect(run({})).rejects.toBeInstanceOf(ProviderError) await expect(run({})).rejects.toThrow('rate limited') }) })