mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-05 08:10:14 +00:00
add tests
This commit is contained in:
joshualipman123
daniel-lxs
parent
6abba41abb
commit
226924e2d6
3 changed files with 933 additions and 0 deletions
383
src/api/providers/__tests__/vercel-ai-gateway.spec.ts
Normal file
383
src/api/providers/__tests__/vercel-ai-gateway.spec.ts
Normal file
|
|
@ -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,
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
317
src/api/providers/fetchers/__tests__/vercel-ai-gateway.spec.ts
Normal file
317
src/api/providers/fetchers/__tests__/vercel-ai-gateway.spec.ts
Normal file
|
|
@ -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)
|
||||
}
|
||||
})
|
||||
})
|
||||
})
|
||||
})
|
||||
233
src/api/transform/caching/__tests__/vercel-ai-gateway.spec.ts
Normal file
233
src/api/transform/caching/__tests__/vercel-ai-gateway.spec.ts
Normal file
|
|
@ -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" },
|
||||
})
|
||||
}
|
||||
})
|
||||
})
|
||||
})
|
||||
Loading…
Add table
Reference in a new issue