diff --git a/src/api/providers/__tests__/vercel-ai-gateway.spec.ts b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts new file mode 100644 index 0000000000..a3e7c9e7d5 --- /dev/null +++ b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts @@ -0,0 +1,383 @@ +// npx vitest run src/api/providers/__tests__/vercel-ai-gateway.spec.ts + +// Mock vscode first to avoid import errors +vitest.mock("vscode", () => ({})) + +import { Anthropic } from "@anthropic-ai/sdk" +import OpenAI from "openai" + +import { VercelAiGatewayHandler } from "../vercel-ai-gateway" +import { ApiHandlerOptions } from "../../../shared/api" +import { vercelAiGatewayDefaultModelId, VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types" + +// Mock dependencies +vitest.mock("openai") +vitest.mock("delay", () => ({ default: vitest.fn(() => Promise.resolve()) })) +vitest.mock("../fetchers/modelCache", () => ({ + getModels: vitest.fn().mockImplementation(() => { + return Promise.resolve({ + "anthropic/claude-sonnet-4": { + maxTokens: 64000, + contextWindow: 200000, + supportsImages: true, + supportsPromptCache: true, + inputPrice: 3, + outputPrice: 15, + cacheWritesPrice: 3.75, + cacheReadsPrice: 0.3, + description: "Claude Sonnet 4", + supportsComputerUse: true, + }, + "anthropic/claude-3.5-haiku": { + maxTokens: 32000, + contextWindow: 200000, + supportsImages: true, + supportsPromptCache: true, + inputPrice: 1, + outputPrice: 5, + cacheWritesPrice: 1.25, + cacheReadsPrice: 0.1, + description: "Claude 3.5 Haiku", + supportsComputerUse: false, + }, + "openai/gpt-4o": { + maxTokens: 16000, + contextWindow: 128000, + supportsImages: true, + supportsPromptCache: true, + inputPrice: 2.5, + outputPrice: 10, + cacheWritesPrice: 3.125, + cacheReadsPrice: 0.25, + description: "GPT-4o", + supportsComputerUse: true, + }, + }) + }), +})) + +vitest.mock("../../transform/caching/vercel-ai-gateway", () => ({ + addCacheBreakpoints: vitest.fn(), +})) + +const mockCreate = vitest.fn() +const mockConstructor = vitest.fn() + +;(OpenAI as any).mockImplementation(() => ({ + chat: { + completions: { + create: mockCreate, + }, + }, +})) +;(OpenAI as any).mockImplementation = mockConstructor.mockReturnValue({ + chat: { + completions: { + create: mockCreate, + }, + }, +}) + +describe("VercelAiGatewayHandler", () => { + const mockOptions: ApiHandlerOptions = { + vercelAiGatewayApiKey: "test-key", + vercelAiGatewayModelId: "anthropic/claude-sonnet-4", + } + + beforeEach(() => { + vitest.clearAllMocks() + mockCreate.mockClear() + mockConstructor.mockClear() + }) + + it("initializes with correct options", () => { + const handler = new VercelAiGatewayHandler(mockOptions) + expect(handler).toBeInstanceOf(VercelAiGatewayHandler) + + expect(OpenAI).toHaveBeenCalledWith({ + baseURL: "https://ai-gateway.vercel.sh/v1", + apiKey: mockOptions.vercelAiGatewayApiKey, + defaultHeaders: expect.objectContaining({ + "HTTP-Referer": "https://github.com/RooVetGit/Roo-Cline", + "X-Title": "Roo Code", + "User-Agent": expect.stringContaining("RooCode/"), + }), + }) + }) + + describe("fetchModel", () => { + it("returns correct model info when options are provided", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + const result = await handler.fetchModel() + + expect(result.id).toBe(mockOptions.vercelAiGatewayModelId) + expect(result.info.maxTokens).toBe(64000) + expect(result.info.contextWindow).toBe(200000) + expect(result.info.supportsImages).toBe(true) + expect(result.info.supportsPromptCache).toBe(true) + expect(result.info.supportsComputerUse).toBe(true) + }) + + it("returns default model info when options are not provided", async () => { + const handler = new VercelAiGatewayHandler({}) + const result = await handler.fetchModel() + expect(result.id).toBe(vercelAiGatewayDefaultModelId) + expect(result.info.supportsPromptCache).toBe(true) + }) + + it("uses vercel ai gateway default model when no model specified", async () => { + const handler = new VercelAiGatewayHandler({ vercelAiGatewayApiKey: "test-key" }) + const result = await handler.fetchModel() + expect(result.id).toBe("anthropic/claude-sonnet-4") + }) + }) + + describe("createMessage", () => { + beforeEach(() => { + mockCreate.mockImplementation(async () => ({ + [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, + cache_creation_input_tokens: 2, + prompt_tokens_details: { + cached_tokens: 3, + }, + cost: 0.005, + }, + } + }, + })) + }) + + it("streams text content correctly", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + + const stream = handler.createMessage(systemPrompt, messages) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks).toHaveLength(2) + expect(chunks[0]).toEqual({ + type: "text", + text: "Test response", + }) + expect(chunks[1]).toEqual({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + cacheWriteTokens: 2, + cacheReadTokens: 3, + totalCost: 0.005, + }) + }) + + it("uses correct temperature from options", async () => { + const customTemp = 0.5 + const handler = new VercelAiGatewayHandler({ + ...mockOptions, + modelTemperature: customTemp, + }) + + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + + await handler.createMessage(systemPrompt, messages).next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + temperature: customTemp, + }), + ) + }) + + it("uses default temperature when none provided", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + + await handler.createMessage(systemPrompt, messages).next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + temperature: VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE, + }), + ) + }) + + it("adds cache breakpoints for supported models", async () => { + const { addCacheBreakpoints } = await import("../../transform/caching/vercel-ai-gateway") + const handler = new VercelAiGatewayHandler({ + ...mockOptions, + vercelAiGatewayModelId: "anthropic/claude-3.5-haiku", + }) + + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + + await handler.createMessage(systemPrompt, messages).next() + + expect(addCacheBreakpoints).toHaveBeenCalled() + }) + + it("sets correct max_completion_tokens", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + + await handler.createMessage(systemPrompt, messages).next() + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + max_completion_tokens: 64000, // max tokens for sonnet 4 + }), + ) + }) + + it("handles usage info correctly with all Vercel AI Gateway specific fields", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] + + const stream = handler.createMessage(systemPrompt, messages) + const chunks = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toEqual({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + cacheWriteTokens: 2, + cacheReadTokens: 3, + totalCost: 0.005, + }) + }) + }) + + describe("completePrompt", () => { + beforeEach(() => { + mockCreate.mockImplementation(async () => ({ + choices: [ + { + message: { role: "assistant", content: "Test completion response" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 8, + completion_tokens: 4, + total_tokens: 12, + }, + })) + }) + + it("completes prompt correctly", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + const prompt = "Complete this: Hello" + + const result = await handler.completePrompt(prompt) + + expect(result).toBe("Test completion response") + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "anthropic/claude-sonnet-4", + messages: [{ role: "user", content: prompt }], + stream: false, + temperature: VERCEL_AI_GATEWAY_DEFAULT_TEMPERATURE, + max_completion_tokens: 64000, + }), + ) + }) + + it("uses custom temperature for completion", async () => { + const customTemp = 0.8 + const handler = new VercelAiGatewayHandler({ + ...mockOptions, + modelTemperature: customTemp, + }) + + await handler.completePrompt("Test prompt") + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + temperature: customTemp, + }), + ) + }) + + it("handles completion errors correctly", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + const errorMessage = "API error" + + mockCreate.mockImplementation(() => { + throw new Error(errorMessage) + }) + + await expect(handler.completePrompt("Test")).rejects.toThrow( + `Vercel AI Gateway completion error: ${errorMessage}`, + ) + }) + + it("returns empty string when no content in response", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + + mockCreate.mockImplementation(async () => ({ + choices: [ + { + message: { role: "assistant", content: null }, + finish_reason: "stop", + index: 0, + }, + ], + })) + + const result = await handler.completePrompt("Test") + expect(result).toBe("") + }) + }) + + describe("temperature support", () => { + it("applies temperature for supported models", async () => { + const handler = new VercelAiGatewayHandler({ + ...mockOptions, + vercelAiGatewayModelId: "anthropic/claude-sonnet-4", + modelTemperature: 0.9, + }) + + await handler.completePrompt("Test") + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + temperature: 0.9, + }), + ) + }) + }) +}) diff --git a/src/api/providers/fetchers/__tests__/vercel-ai-gateway.spec.ts b/src/api/providers/fetchers/__tests__/vercel-ai-gateway.spec.ts new file mode 100644 index 0000000000..657d335b61 --- /dev/null +++ b/src/api/providers/fetchers/__tests__/vercel-ai-gateway.spec.ts @@ -0,0 +1,317 @@ +// npx vitest run src/api/providers/fetchers/__tests__/vercel-ai-gateway.spec.ts + +import axios from "axios" +import { VERCEL_AI_GATEWAY_VISION_ONLY_MODELS, VERCEL_AI_GATEWAY_VISION_AND_TOOLS_MODELS } from "@roo-code/types" + +import { getVercelAiGatewayModels, parseVercelAiGatewayModel } from "../vercel-ai-gateway" + +vitest.mock("axios") +const mockedAxios = axios as any + +describe("Vercel AI Gateway Fetchers", () => { + beforeEach(() => { + vitest.clearAllMocks() + }) + + describe("getVercelAiGatewayModels", () => { + const mockResponse = { + data: { + object: "list", + data: [ + { + id: "anthropic/claude-sonnet-4", + object: "model", + created: 1640995200, + owned_by: "anthropic", + name: "Claude Sonnet 4", + description: + "Claude Sonnet 4 significantly improves on Sonnet 3.7's industry-leading capabilities", + context_window: 200000, + max_tokens: 64000, + type: "language", + pricing: { + input: "3.00", + output: "15.00", + input_cache_write: "3.75", + input_cache_read: "0.30", + }, + }, + { + id: "anthropic/claude-3.5-haiku", + object: "model", + created: 1640995200, + owned_by: "anthropic", + name: "Claude 3.5 Haiku", + description: "Claude 3.5 Haiku is fast and lightweight", + context_window: 200000, + max_tokens: 32000, + type: "language", + pricing: { + input: "1.00", + output: "5.00", + input_cache_write: "1.25", + input_cache_read: "0.10", + }, + }, + { + id: "dall-e-3", + object: "model", + created: 1640995200, + owned_by: "openai", + name: "DALL-E 3", + description: "DALL-E 3 image generation model", + context_window: 4000, + max_tokens: 1000, + type: "image", + pricing: { + input: "40.00", + output: "0.00", + }, + }, + ], + }, + } + + it("fetches and parses models correctly", async () => { + mockedAxios.get.mockResolvedValueOnce(mockResponse) + + const models = await getVercelAiGatewayModels() + + expect(mockedAxios.get).toHaveBeenCalledWith("https://ai-gateway.vercel.sh/v1/models") + expect(Object.keys(models)).toHaveLength(2) // Only language models + expect(models["anthropic/claude-sonnet-4"]).toBeDefined() + expect(models["anthropic/claude-3.5-haiku"]).toBeDefined() + }) + + it("handles API errors gracefully", async () => { + const consoleErrorSpy = vitest.spyOn(console, "error").mockImplementation(() => {}) + mockedAxios.get.mockRejectedValueOnce(new Error("Network error")) + + const models = await getVercelAiGatewayModels() + + expect(models).toEqual({}) + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining("Error fetching Vercel AI Gateway models"), + ) + consoleErrorSpy.mockRestore() + }) + + it("handles invalid response schema gracefully", async () => { + const consoleErrorSpy = vitest.spyOn(console, "error").mockImplementation(() => {}) + mockedAxios.get.mockResolvedValueOnce({ + data: { + invalid: "response", + data: "not an array", + }, + }) + + const models = await getVercelAiGatewayModels() + + expect(models).toEqual({}) + expect(consoleErrorSpy).toHaveBeenCalledWith( + "Vercel AI Gateway models response is invalid", + expect.any(Object), + ) + consoleErrorSpy.mockRestore() + }) + + it("continues processing with partially valid schema", async () => { + const consoleErrorSpy = vitest.spyOn(console, "error").mockImplementation(() => {}) + const invalidResponse = { + data: { + invalid_root: "response", + data: [ + { + id: "anthropic/claude-sonnet-4", + object: "model", + created: 1640995200, + owned_by: "anthropic", + name: "Claude Sonnet 4", + description: "Claude Sonnet 4", + context_window: 200000, + max_tokens: 64000, + type: "language", + pricing: { + input: "3.00", + output: "15.00", + }, + }, + ], + }, + } + mockedAxios.get.mockResolvedValueOnce(invalidResponse) + + const models = await getVercelAiGatewayModels() + + expect(consoleErrorSpy).toHaveBeenCalled() + expect(models["anthropic/claude-sonnet-4"]).toBeDefined() + consoleErrorSpy.mockRestore() + }) + }) + + describe("parseVercelAiGatewayModel", () => { + const baseModel = { + id: "test/model", + object: "model", + created: 1640995200, + owned_by: "test", + name: "Test Model", + description: "A test model", + context_window: 100000, + max_tokens: 8000, + type: "language", + pricing: { + input: "2.50", + output: "10.00", + }, + } + + it("parses basic model info correctly", () => { + const result = parseVercelAiGatewayModel({ + id: "test/model", + model: baseModel, + }) + + expect(result).toEqual({ + maxTokens: 8000, + contextWindow: 100000, + supportsImages: false, + supportsComputerUse: false, + supportsPromptCache: false, + inputPrice: 2500000, + outputPrice: 10000000, + cacheWritesPrice: undefined, + cacheReadsPrice: undefined, + description: "A test model", + }) + }) + + it("parses cache pricing when available", () => { + const modelWithCache = { + ...baseModel, + pricing: { + input: "3.00", + output: "15.00", + input_cache_write: "3.75", + input_cache_read: "0.30", + }, + } + + const result = parseVercelAiGatewayModel({ + id: "anthropic/claude-sonnet-4", + model: modelWithCache, + }) + + expect(result).toMatchObject({ + supportsPromptCache: true, + cacheWritesPrice: 3750000, + cacheReadsPrice: 300000, + }) + }) + + it("detects vision-only models", () => { + // claude 3.5 haiku in VERCEL_AI_GATEWAY_VISION_ONLY_MODELS + const visionModel = { + ...baseModel, + id: "anthropic/claude-3.5-haiku", + } + + const result = parseVercelAiGatewayModel({ + id: "anthropic/claude-3.5-haiku", + model: visionModel, + }) + + expect(result.supportsImages).toBe(VERCEL_AI_GATEWAY_VISION_ONLY_MODELS.has("anthropic/claude-3.5-haiku")) + expect(result.supportsComputerUse).toBe(false) + }) + + it("detects vision and tools models", () => { + // 4 sonnet in VERCEL_AI_GATEWAY_VISION_AND_TOOLS_MODELS + const visionToolsModel = { + ...baseModel, + id: "anthropic/claude-sonnet-4", + } + + const result = parseVercelAiGatewayModel({ + id: "anthropic/claude-sonnet-4", + model: visionToolsModel, + }) + + expect(result.supportsImages).toBe( + VERCEL_AI_GATEWAY_VISION_AND_TOOLS_MODELS.has("anthropic/claude-sonnet-4"), + ) + expect(result.supportsComputerUse).toBe( + VERCEL_AI_GATEWAY_VISION_AND_TOOLS_MODELS.has("anthropic/claude-sonnet-4"), + ) + }) + + it("handles missing cache pricing", () => { + const modelNoCachePricing = { + ...baseModel, + pricing: { + input: "2.50", + output: "10.00", + // No cache pricing + }, + } + + const result = parseVercelAiGatewayModel({ + id: "test/model", + model: modelNoCachePricing, + }) + + expect(result.supportsPromptCache).toBe(false) + expect(result.cacheWritesPrice).toBeUndefined() + expect(result.cacheReadsPrice).toBeUndefined() + }) + + it("handles partial cache pricing", () => { + const modelPartialCachePricing = { + ...baseModel, + pricing: { + input: "2.50", + output: "10.00", + input_cache_write: "3.00", + // Missing input_cache_read + }, + } + + const result = parseVercelAiGatewayModel({ + id: "test/model", + model: modelPartialCachePricing, + }) + + expect(result.supportsPromptCache).toBe(false) + expect(result.cacheWritesPrice).toBe(3000000) + expect(result.cacheReadsPrice).toBeUndefined() + }) + + it("validates all vision model categories", () => { + // Test a few models from each category + const visionOnlyModels = ["anthropic/claude-3.5-haiku", "google/gemini-1.5-flash-8b"] + const visionAndToolsModels = ["anthropic/claude-sonnet-4", "openai/gpt-4o"] + + visionOnlyModels.forEach((modelId) => { + if (VERCEL_AI_GATEWAY_VISION_ONLY_MODELS.has(modelId)) { + const result = parseVercelAiGatewayModel({ + id: modelId, + model: { ...baseModel, id: modelId }, + }) + expect(result.supportsImages).toBe(true) + expect(result.supportsComputerUse).toBe(false) + } + }) + + visionAndToolsModels.forEach((modelId) => { + if (VERCEL_AI_GATEWAY_VISION_AND_TOOLS_MODELS.has(modelId)) { + const result = parseVercelAiGatewayModel({ + id: modelId, + model: { ...baseModel, id: modelId }, + }) + expect(result.supportsImages).toBe(true) + expect(result.supportsComputerUse).toBe(true) + } + }) + }) + }) +}) diff --git a/src/api/transform/caching/__tests__/vercel-ai-gateway.spec.ts b/src/api/transform/caching/__tests__/vercel-ai-gateway.spec.ts new file mode 100644 index 0000000000..86dc593f4f --- /dev/null +++ b/src/api/transform/caching/__tests__/vercel-ai-gateway.spec.ts @@ -0,0 +1,233 @@ +// npx vitest run src/api/transform/caching/__tests__/vercel-ai-gateway.spec.ts + +import OpenAI from "openai" +import { addCacheBreakpoints } from "../vercel-ai-gateway" + +describe("Vercel AI Gateway Caching", () => { + describe("addCacheBreakpoints", () => { + it("adds cache control to system message", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { role: "user", content: "Hello" }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + expect(messages[0]).toEqual({ + role: "system", + content: systemPrompt, + cache_control: { type: "ephemeral" }, + }) + }) + + it("adds cache control to last two user messages with string content", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { role: "user", content: "First message" }, + { role: "assistant", content: "First response" }, + { role: "user", content: "Second message" }, + { role: "assistant", content: "Second response" }, + { role: "user", content: "Third message" }, + { role: "assistant", content: "Third response" }, + { role: "user", content: "Fourth message" }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + const lastUserMessage = messages[7] + expect(Array.isArray(lastUserMessage.content)).toBe(true) + if (Array.isArray(lastUserMessage.content)) { + const textPart = lastUserMessage.content.find((part) => part.type === "text") + expect(textPart).toEqual({ + type: "text", + text: "Fourth message", + cache_control: { type: "ephemeral" }, + }) + } + + const secondLastUserMessage = messages[5] + expect(Array.isArray(secondLastUserMessage.content)).toBe(true) + if (Array.isArray(secondLastUserMessage.content)) { + const textPart = secondLastUserMessage.content.find((part) => part.type === "text") + expect(textPart).toEqual({ + type: "text", + text: "Third message", + cache_control: { type: "ephemeral" }, + }) + } + }) + + it("handles messages with existing array content", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { + role: "user", + content: [ + { type: "text", text: "Hello with image" }, + { type: "image_url", image_url: { url: "data:image/png;base64,..." } }, + ], + }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + const userMessage = messages[1] + expect(Array.isArray(userMessage.content)).toBe(true) + if (Array.isArray(userMessage.content)) { + const textPart = userMessage.content.find((part) => part.type === "text") + expect(textPart).toEqual({ + type: "text", + text: "Hello with image", + cache_control: { type: "ephemeral" }, + }) + + const imagePart = userMessage.content.find((part) => part.type === "image_url") + expect(imagePart).toEqual({ + type: "image_url", + image_url: { url: "data:image/png;base64,..." }, + }) + } + }) + + it("handles empty string content gracefully", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { role: "user", content: "" }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + const userMessage = messages[1] + expect(userMessage.content).toBe("") + }) + + it("handles messages with no text parts", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { + role: "user", + content: [{ type: "image_url", image_url: { url: "data:image/png;base64,..." } }], + }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + const userMessage = messages[1] + expect(Array.isArray(userMessage.content)).toBe(true) + if (Array.isArray(userMessage.content)) { + const textPart = userMessage.content.find((part) => part.type === "text") + expect(textPart).toBeUndefined() + + const imagePart = userMessage.content.find((part) => part.type === "image_url") + expect(imagePart).toEqual({ + type: "image_url", + image_url: { url: "data:image/png;base64,..." }, + }) + } + }) + + it("processes only user messages for conversation caching", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { role: "user", content: "First user" }, + { role: "assistant", content: "Assistant response" }, + { role: "user", content: "Second user" }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + expect(messages[2]).toEqual({ + role: "assistant", + content: "Assistant response", + }) + + const firstUser = messages[1] + const secondUser = messages[3] + + expect(Array.isArray(firstUser.content)).toBe(true) + expect(Array.isArray(secondUser.content)).toBe(true) + }) + + it("handles case with only one user message", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { role: "user", content: "Only message" }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + const userMessage = messages[1] + expect(Array.isArray(userMessage.content)).toBe(true) + if (Array.isArray(userMessage.content)) { + const textPart = userMessage.content.find((part) => part.type === "text") + expect(textPart).toEqual({ + type: "text", + text: "Only message", + cache_control: { type: "ephemeral" }, + }) + } + }) + + it("handles case with no user messages", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { role: "assistant", content: "Assistant only" }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + expect(messages[0]).toEqual({ + role: "system", + content: systemPrompt, + cache_control: { type: "ephemeral" }, + }) + + expect(messages[1]).toEqual({ + role: "assistant", + content: "Assistant only", + }) + }) + + it("handles messages with multiple text parts", () => { + const systemPrompt = "You are a helpful assistant." + const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [ + { role: "system", content: systemPrompt }, + { + role: "user", + content: [ + { type: "text", text: "First part" }, + { type: "image_url", image_url: { url: "data:image/png;base64,..." } }, + { type: "text", text: "Second part" }, + ], + }, + ] + + addCacheBreakpoints(systemPrompt, messages) + + const userMessage = messages[1] + if (Array.isArray(userMessage.content)) { + const textParts = userMessage.content.filter((part) => part.type === "text") + expect(textParts).toHaveLength(2) + + expect(textParts[0]).toEqual({ + type: "text", + text: "First part", + }) + + expect(textParts[1]).toEqual({ + type: "text", + text: "Second part", + cache_control: { type: "ephemeral" }, + }) + } + }) + }) +})