diff --git a/src/api/providers/__tests__/ark-caching.spec.ts b/src/api/providers/__tests__/ark-caching.spec.ts new file mode 100644 index 0000000000..5205615a65 --- /dev/null +++ b/src/api/providers/__tests__/ark-caching.spec.ts @@ -0,0 +1,520 @@ +// npx vitest run api/providers/__tests__/ark-caching.spec.ts + +import { Anthropic } from "@anthropic-ai/sdk" +import { OpenAiHandler } from "../openai" +import { ApiHandlerOptions } from "../../../shared/api" + +const mockCreate = vitest.fn() + +vitest.mock("openai", () => { + return { + __esModule: true, + default: vitest.fn().mockImplementation(() => ({ + chat: { + completions: { + create: mockCreate, + }, + }, + })), + } +}) + +describe("OpenAiHandler - Ark Context Caching", () => { + let handler: OpenAiHandler + let arkOptions: ApiHandlerOptions + + beforeEach(() => { + arkOptions = { + openAiApiKey: "test-api-key", + openAiModelId: "doubao-pro-4k", + openAiBaseUrl: "https://ark.cn-beijing.volces.com/api/v3", + } + handler = new OpenAiHandler(arkOptions) + mockCreate.mockClear() + }) + + describe("Ark provider detection", () => { + it("should detect Ark provider from volces.com URL", () => { + expect(arkOptions.openAiBaseUrl).toContain(".volces.com") + }) + + it("should not detect Ark for non-volces URLs", () => { + const nonArkHandler = new OpenAiHandler({ + ...arkOptions, + openAiBaseUrl: "https://api.openai.com/v1", + }) + expect(nonArkHandler).toBeInstanceOf(OpenAiHandler) + }) + }) + + describe("context caching with streaming", () => { + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello!", + }, + ] + + beforeEach(() => { + mockCreate.mockImplementation(async (options) => { + if (!options.stream) { + return { + id: "response-123", + choices: [ + { + message: { role: "assistant", content: "Test response" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + prompt_tokens_details: { + cached_tokens: 8, // Ark-specific cached tokens + }, + }, + } + } + + return { + [Symbol.asyncIterator]: async function* () { + yield { + id: "response-123", + choices: [ + { + delta: { content: "Test response" }, + index: 0, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + prompt_tokens_details: { + cached_tokens: 8, // Ark-specific cached tokens + }, + }, + } + }, + } + }) + }) + + it("should add caching parameters for first request (no previous response ID)", async () => { + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "doubao-pro-4k", + caching: { type: "enabled" }, + cache_ttl: 3600, + // Should not have previous_response_id for first request + }), + {}, + ) + + // Should not have previous_response_id in the first call + const callArgs = mockCreate.mock.calls[0][0] + expect(callArgs).not.toHaveProperty("previous_response_id") + }) + + it("should include previous_response_id for subsequent requests", async () => { + // First request + const stream1 = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream1) { + // Consume stream to store response ID + } + + mockCreate.mockClear() + + // Second request should include previous response ID + const stream2 = handler.createMessage(systemPrompt, [ + ...messages, + { + role: "assistant", + content: "Test response", + }, + { + role: "user", + content: "Follow up question", + }, + ]) + for await (const _chunk of stream2) { + // Consume stream + } + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "doubao-pro-4k", + caching: { type: "enabled" }, + cache_ttl: 3600, + previous_response_id: "response-123", + }), + {}, + ) + }) + + it("should process cached tokens in usage metrics", async () => { + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(10) + expect(usageChunk.outputTokens).toBe(5) + expect(usageChunk.cacheReadTokens).toBe(8) // Ark cached tokens + }) + }) + + describe("context caching with non-streaming", () => { + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: "Hello!", + }, + ] + + beforeEach(() => { + handler = new OpenAiHandler({ + ...arkOptions, + openAiStreamingEnabled: false, + }) + + mockCreate.mockImplementation(async (options) => { + return { + id: "response-456", + choices: [ + { + message: { role: "assistant", content: "Non-streaming response" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 15, + completion_tokens: 8, + total_tokens: 23, + prompt_tokens_details: { + cached_tokens: 12, // Ark-specific cached tokens + }, + }, + } + }) + }) + + it("should add caching parameters for non-streaming requests", async () => { + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "doubao-pro-4k", + caching: { type: "enabled" }, + cache_ttl: 3600, + }), + {}, + ) + }) + + it("should process cached tokens in non-streaming usage metrics", async () => { + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(15) + expect(usageChunk.outputTokens).toBe(8) + expect(usageChunk.cacheReadTokens).toBe(12) // Ark cached tokens + }) + + it("should store response ID for future requests", async () => { + // First request + const stream1 = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream1) { + // Consume stream to store response ID + } + + mockCreate.mockClear() + + // Second request should include previous response ID + const stream2 = handler.createMessage(systemPrompt, messages) + for await (const _chunk of stream2) { + // Consume stream + } + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + previous_response_id: "response-456", + }), + {}, + ) + }) + }) + + describe("O3 family models with Ark caching", () => { + beforeEach(() => { + handler = new OpenAiHandler({ + ...arkOptions, + openAiModelId: "o3-mini", // O3 family model + }) + + mockCreate.mockImplementation(async (options) => { + if (!options.stream) { + return { + id: "o3-response-789", + choices: [ + { + message: { role: "assistant", content: "O3 response" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: 15, + }, + }, + } + } + + return { + [Symbol.asyncIterator]: async function* () { + yield { + id: "o3-response-789", + choices: [ + { + delta: { content: "O3 streaming response" }, + index: 0, + }, + ], + usage: { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: 15, + }, + }, + } + }, + } + }) + }) + + it("should add caching parameters for O3 family streaming requests", async () => { + const stream = handler.createMessage("System prompt", [ + { + role: "user", + content: "Hello O3!", + }, + ]) + for await (const _chunk of stream) { + // Consume stream + } + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "o3-mini", + caching: { type: "enabled" }, + cache_ttl: 3600, + reasoning_effort: undefined, // O3 specific + temperature: undefined, // O3 specific + }), + {}, + ) + }) + + it("should add caching parameters for O3 family non-streaming requests", async () => { + handler = new OpenAiHandler({ + ...arkOptions, + openAiModelId: "o3-mini", + openAiStreamingEnabled: false, + }) + + const stream = handler.createMessage("System prompt", [ + { + role: "user", + content: "Hello O3!", + }, + ]) + for await (const _chunk of stream) { + // Consume stream + } + + expect(mockCreate).toHaveBeenCalledWith( + expect.objectContaining({ + model: "o3-mini", + caching: { type: "enabled" }, + cache_ttl: 3600, + reasoning_effort: undefined, + temperature: undefined, + }), + {}, + ) + }) + }) + + describe("edge cases", () => { + it("should handle missing usage data gracefully", async () => { + mockCreate.mockImplementation(async (options) => { + if (!options.stream) { + return { + id: "response-no-usage", + choices: [ + { + message: { role: "assistant", content: "Response without usage" }, + finish_reason: "stop", + index: 0, + }, + ], + // No usage data + } + } + + return { + [Symbol.asyncIterator]: async function* () { + yield { + id: "response-no-usage", + choices: [ + { + delta: { content: "Response without usage" }, + index: 0, + }, + ], + // No usage data + } + }, + } + }) + + const stream = handler.createMessage("System prompt", [ + { + role: "user", + content: "Hello!", + }, + ]) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + // Should not crash and should not yield usage chunk + const usageChunks = chunks.filter((chunk) => chunk.type === "usage") + expect(usageChunks).toHaveLength(0) + }) + + it("should handle missing cached tokens gracefully", async () => { + mockCreate.mockImplementation(async (options) => { + return { + id: "response-no-cache", + choices: [ + { + message: { role: "assistant", content: "Response without cached tokens" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + // No prompt_tokens_details.cached_tokens + }, + } + }) + + handler = new OpenAiHandler({ + ...arkOptions, + openAiStreamingEnabled: false, + }) + + const stream = handler.createMessage("System prompt", [ + { + role: "user", + content: "Hello!", + }, + ]) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + const usageChunk = chunks.find((chunk) => chunk.type === "usage") + expect(usageChunk).toBeDefined() + expect(usageChunk.inputTokens).toBe(10) + expect(usageChunk.outputTokens).toBe(5) + expect(usageChunk.cacheReadTokens).toBeUndefined() // No cached tokens + }) + + it("should handle missing response ID gracefully", async () => { + mockCreate.mockImplementation(async (options) => { + return { + // No id field + choices: [ + { + message: { role: "assistant", content: "Response without ID" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }, + } + }) + + handler = new OpenAiHandler({ + ...arkOptions, + openAiStreamingEnabled: false, + }) + + const stream = handler.createMessage("System prompt", [ + { + role: "user", + content: "Hello!", + }, + ]) + for await (const _chunk of stream) { + // Consume stream + } + + // Should not crash, and subsequent requests should not have previous_response_id + mockCreate.mockClear() + + const stream2 = handler.createMessage("System prompt", [ + { + role: "user", + content: "Follow up", + }, + ]) + for await (const _chunk of stream2) { + // Consume stream + } + + const callArgs = mockCreate.mock.calls[0][0] + expect(callArgs).not.toHaveProperty("previous_response_id") + }) + }) +}) diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index f5e4e4c985..df502c0442 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -19,6 +19,7 @@ import { convertToR1Format } from "../transform/r1-format" import { convertToSimpleMessages } from "../transform/simple-format" import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { getModelParams } from "../transform/model-params" +import { addArkCaching, extractArkResponseId, getArkCachedTokens } from "../transform/caching/ark" import { DEFAULT_HEADERS } from "./constants" import { BaseProvider } from "./base-provider" @@ -30,6 +31,7 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from ". export class OpenAiHandler extends BaseProvider implements SingleCompletionHandler { protected options: ApiHandlerOptions private client: OpenAI + private arkPreviousResponseId?: string constructor(options: ApiHandlerOptions) { super() @@ -161,6 +163,14 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl // Add max_tokens if needed this.addMaxTokensIfNeeded(requestOptions, modelInfo) + // Add Ark context caching if this is an Ark provider + if (ark) { + addArkCaching(requestOptions, { + previousResponseId: this.arkPreviousResponseId, + cacheTtl: 3600, // 1 hour as recommended in the issue + }) + } + const stream = await this.client.chat.completions.create( requestOptions, isAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {}, @@ -176,6 +186,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl ) let lastUsage + let responseId: string | undefined for await (const chunk of stream) { const delta = chunk.choices[0]?.delta ?? {} @@ -195,14 +206,24 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl if (chunk.usage) { lastUsage = chunk.usage } + + // Extract response ID for Ark caching + if (ark && chunk.id) { + responseId = chunk.id + } } for (const chunk of matcher.final()) { yield chunk } + // Store response ID for future Ark caching + if (ark && responseId) { + this.arkPreviousResponseId = responseId + } + if (lastUsage) { - yield this.processUsageMetrics(lastUsage, modelInfo) + yield this.processUsageMetrics(lastUsage, modelInfo, ark) } } else { // o1 for instance doesnt support streaming, non-1 temp, or system prompt @@ -223,28 +244,51 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl // Add max_tokens if needed this.addMaxTokensIfNeeded(requestOptions, modelInfo) + // Add Ark context caching if this is an Ark provider + if (ark) { + addArkCaching(requestOptions, { + previousResponseId: this.arkPreviousResponseId, + cacheTtl: 3600, // 1 hour as recommended in the issue + }) + } + const response = await this.client.chat.completions.create( requestOptions, this._isAzureAiInference(modelUrl) ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {}, ) + // Store response ID for future Ark caching + if (ark && response.id) { + this.arkPreviousResponseId = response.id + } + yield { type: "text", text: response.choices[0]?.message.content || "", } - yield this.processUsageMetrics(response.usage, modelInfo) + yield this.processUsageMetrics(response.usage, modelInfo, ark) } } - protected processUsageMetrics(usage: any, _modelInfo?: ModelInfo): ApiStreamUsageChunk { - return { + protected processUsageMetrics(usage: any, _modelInfo?: ModelInfo, isArk?: boolean): ApiStreamUsageChunk { + const result: ApiStreamUsageChunk = { type: "usage", inputTokens: usage?.prompt_tokens || 0, outputTokens: usage?.completion_tokens || 0, cacheWriteTokens: usage?.cache_creation_input_tokens || undefined, cacheReadTokens: usage?.cache_read_input_tokens || undefined, } + + // Handle Ark-specific cached tokens + if (isArk) { + const arkCachedTokens = getArkCachedTokens(usage) + if (arkCachedTokens > 0) { + result.cacheReadTokens = arkCachedTokens + } + } + + return result } override getModel() { @@ -290,6 +334,8 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl ): ApiStream { const modelInfo = this.getModel().info const methodIsAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl) + const modelUrl = this.options.openAiBaseUrl ?? "" + const ark = modelUrl.includes(".volces.com") if (this.options.openAiStreamingEnabled ?? true) { const isGrokXAI = this._isGrokXAI(this.options.openAiBaseUrl) @@ -314,12 +360,20 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl // This allows O3 models to limit response length when includeMaxTokens is enabled this.addMaxTokensIfNeeded(requestOptions, modelInfo) + // Add Ark context caching if this is an Ark provider + if (ark) { + addArkCaching(requestOptions, { + previousResponseId: this.arkPreviousResponseId, + cacheTtl: 3600, // 1 hour as recommended in the issue + }) + } + const stream = await this.client.chat.completions.create( requestOptions, methodIsAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {}, ) - yield* this.handleStreamResponse(stream) + yield* this.handleStreamResponse(stream, ark) } else { const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = { model: modelId, @@ -339,20 +393,38 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl // This allows O3 models to limit response length when includeMaxTokens is enabled this.addMaxTokensIfNeeded(requestOptions, modelInfo) + // Add Ark context caching if this is an Ark provider + if (ark) { + addArkCaching(requestOptions, { + previousResponseId: this.arkPreviousResponseId, + cacheTtl: 3600, // 1 hour as recommended in the issue + }) + } + const response = await this.client.chat.completions.create( requestOptions, methodIsAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {}, ) + // Store response ID for future Ark caching + if (ark && response.id) { + this.arkPreviousResponseId = response.id + } + yield { type: "text", text: response.choices[0]?.message.content || "", } - yield this.processUsageMetrics(response.usage) + yield this.processUsageMetrics(response.usage, modelInfo, ark) } } - private async *handleStreamResponse(stream: AsyncIterable): ApiStream { + private async *handleStreamResponse( + stream: AsyncIterable, + isArk?: boolean, + ): ApiStream { + let responseId: string | undefined + for await (const chunk of stream) { const delta = chunk.choices[0]?.delta if (delta?.content) { @@ -362,14 +434,28 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl } } + // Extract response ID for Ark caching + if (isArk && chunk.id) { + responseId = chunk.id + } + if (chunk.usage) { yield { type: "usage", inputTokens: chunk.usage.prompt_tokens || 0, outputTokens: chunk.usage.completion_tokens || 0, + ...(isArk && + getArkCachedTokens(chunk.usage) > 0 && { + cacheReadTokens: getArkCachedTokens(chunk.usage), + }), } } } + + // Store response ID for future Ark caching + if (isArk && responseId) { + this.arkPreviousResponseId = responseId + } } private _getUrlHost(baseUrl?: string): string { diff --git a/src/api/transform/caching/__tests__/ark.spec.ts b/src/api/transform/caching/__tests__/ark.spec.ts new file mode 100644 index 0000000000..4a50664463 --- /dev/null +++ b/src/api/transform/caching/__tests__/ark.spec.ts @@ -0,0 +1,414 @@ +// npx vitest run api/transform/caching/__tests__/ark.spec.ts + +import { addArkCaching, extractArkResponseId, getArkCachedTokens, hasArkCachedTokens } from "../ark" +import OpenAI from "openai" + +describe("Ark Context Caching", () => { + describe("addArkCaching", () => { + it("should add basic caching configuration", () => { + const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: "doubao-pro-4k", + messages: [ + { role: "system", content: "You are a helpful assistant." }, + { role: "user", content: "Hello!" }, + ], + stream: true, + } + + addArkCaching(requestOptions) + + expect(requestOptions).toEqual( + expect.objectContaining({ + model: "doubao-pro-4k", + messages: expect.any(Array), + stream: true, + caching: { type: "enabled" }, + }), + ) + }) + + it("should add previous response ID when provided", () => { + const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: "doubao-pro-4k", + messages: [{ role: "user", content: "Follow up question" }], + stream: true, + } + + addArkCaching(requestOptions, { + previousResponseId: "response-123", + }) + + expect(requestOptions).toEqual( + expect.objectContaining({ + caching: { type: "enabled" }, + previous_response_id: "response-123", + }), + ) + }) + + it("should add cache TTL when provided", () => { + const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: "doubao-pro-4k", + messages: [{ role: "user", content: "Hello!" }], + stream: true, + } + + addArkCaching(requestOptions, { + cacheTtl: 7200, // 2 hours + }) + + expect(requestOptions).toEqual( + expect.objectContaining({ + caching: { type: "enabled" }, + cache_ttl: 7200, + }), + ) + }) + + it("should add both previous response ID and cache TTL", () => { + const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: "doubao-pro-4k", + messages: [{ role: "user", content: "Hello!" }], + stream: true, + } + + addArkCaching(requestOptions, { + previousResponseId: "response-456", + cacheTtl: 3600, + }) + + expect(requestOptions).toEqual( + expect.objectContaining({ + caching: { type: "enabled" }, + previous_response_id: "response-456", + cache_ttl: 3600, + }), + ) + }) + + it("should work with non-streaming requests", () => { + const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = { + model: "doubao-pro-4k", + messages: [{ role: "user", content: "Hello!" }], + } + + addArkCaching(requestOptions, { + previousResponseId: "response-789", + cacheTtl: 1800, + }) + + expect(requestOptions).toEqual( + expect.objectContaining({ + caching: { type: "enabled" }, + previous_response_id: "response-789", + cache_ttl: 1800, + }), + ) + }) + + it("should not add optional fields when not provided", () => { + const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: "doubao-pro-4k", + messages: [{ role: "user", content: "Hello!" }], + stream: true, + } + + addArkCaching(requestOptions, {}) + + expect(requestOptions).toEqual( + expect.objectContaining({ + caching: { type: "enabled" }, + }), + ) + expect(requestOptions).not.toHaveProperty("previous_response_id") + expect(requestOptions).not.toHaveProperty("cache_ttl") + }) + }) + + describe("extractArkResponseId", () => { + it("should extract response ID from valid response", () => { + const response = { + id: "response-123", + choices: [ + { + message: { role: "assistant", content: "Hello!" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }, + } + + const responseId = extractArkResponseId(response) + expect(responseId).toBe("response-123") + }) + + it("should return undefined for response without ID", () => { + const response = { + choices: [ + { + message: { role: "assistant", content: "Hello!" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }, + } + + const responseId = extractArkResponseId(response) + expect(responseId).toBeUndefined() + }) + + it("should return undefined for null/undefined response", () => { + expect(extractArkResponseId(null)).toBeUndefined() + expect(extractArkResponseId(undefined)).toBeUndefined() + }) + + it("should return undefined for empty object", () => { + const responseId = extractArkResponseId({}) + expect(responseId).toBeUndefined() + }) + }) + + describe("hasArkCachedTokens", () => { + it("should return true when cached tokens are present and greater than 0", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: 15, + }, + } + + expect(hasArkCachedTokens(usage)).toBe(true) + }) + + it("should return false when cached tokens are 0", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: 0, + }, + } + + expect(hasArkCachedTokens(usage)).toBe(false) + }) + + it("should return false when cached tokens are missing", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: {}, + } + + expect(hasArkCachedTokens(usage)).toBe(false) + }) + + it("should return false when prompt_tokens_details is missing", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + } + + expect(hasArkCachedTokens(usage)).toBe(false) + }) + + it("should return false for null/undefined usage", () => { + expect(hasArkCachedTokens(null)).toBe(false) + expect(hasArkCachedTokens(undefined)).toBe(false) + }) + + it("should return false for empty object", () => { + expect(hasArkCachedTokens({})).toBe(false) + }) + }) + + describe("getArkCachedTokens", () => { + it("should return cached tokens count when present", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: 15, + }, + } + + expect(getArkCachedTokens(usage)).toBe(15) + }) + + it("should return 0 when cached tokens are 0", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: 0, + }, + } + + expect(getArkCachedTokens(usage)).toBe(0) + }) + + it("should return 0 when cached tokens are missing", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: {}, + } + + expect(getArkCachedTokens(usage)).toBe(0) + }) + + it("should return 0 when prompt_tokens_details is missing", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + } + + expect(getArkCachedTokens(usage)).toBe(0) + }) + + it("should return 0 for null/undefined usage", () => { + expect(getArkCachedTokens(null)).toBe(0) + expect(getArkCachedTokens(undefined)).toBe(0) + }) + + it("should return 0 for empty object", () => { + expect(getArkCachedTokens({})).toBe(0) + }) + + it("should handle negative cached tokens gracefully", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: -5, // Invalid but should be handled + }, + } + + expect(getArkCachedTokens(usage)).toBe(-5) + }) + + it("should handle non-numeric cached tokens gracefully", () => { + const usage = { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + prompt_tokens_details: { + cached_tokens: "invalid", // Invalid type + }, + } + + expect(getArkCachedTokens(usage)).toBe("invalid") + }) + }) + + describe("integration scenarios", () => { + it("should handle complete caching workflow", () => { + // First request - no previous response ID + const firstRequest: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: "doubao-pro-4k", + messages: [{ role: "user", content: "Hello!" }], + stream: true, + } + + addArkCaching(firstRequest, { cacheTtl: 3600 }) + + expect(firstRequest).toEqual( + expect.objectContaining({ + caching: { type: "enabled" }, + cache_ttl: 3600, + }), + ) + expect(firstRequest).not.toHaveProperty("previous_response_id") + + // Simulate first response + const firstResponse = { + id: "response-first", + choices: [ + { + message: { role: "assistant", content: "Hi there!" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + prompt_tokens_details: { + cached_tokens: 0, // No cache on first request + }, + }, + } + + const firstResponseId = extractArkResponseId(firstResponse) + expect(firstResponseId).toBe("response-first") + expect(hasArkCachedTokens(firstResponse.usage)).toBe(false) + + // Second request - with previous response ID + const secondRequest: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = { + model: "doubao-pro-4k", + messages: [ + { role: "user", content: "Hello!" }, + { role: "assistant", content: "Hi there!" }, + { role: "user", content: "How are you?" }, + ], + stream: true, + } + + addArkCaching(secondRequest, { + previousResponseId: firstResponseId, + cacheTtl: 3600, + }) + + expect(secondRequest).toEqual( + expect.objectContaining({ + caching: { type: "enabled" }, + previous_response_id: "response-first", + cache_ttl: 3600, + }), + ) + + // Simulate second response with cached tokens + const secondResponse = { + id: "response-second", + choices: [ + { + message: { role: "assistant", content: "I'm doing well, thanks!" }, + finish_reason: "stop", + index: 0, + }, + ], + usage: { + prompt_tokens: 20, + completion_tokens: 8, + total_tokens: 28, + prompt_tokens_details: { + cached_tokens: 15, // Cache hit! + }, + }, + } + + expect(hasArkCachedTokens(secondResponse.usage)).toBe(true) + expect(getArkCachedTokens(secondResponse.usage)).toBe(15) + }) + }) +}) diff --git a/src/api/transform/caching/ark.ts b/src/api/transform/caching/ark.ts new file mode 100644 index 0000000000..14c80c9a56 --- /dev/null +++ b/src/api/transform/caching/ark.ts @@ -0,0 +1,78 @@ +import OpenAI from "openai" + +/** + * Ark/Volcengine context caching implementation using the Responses API + * + * According to Volcengine documentation: + * - Uses `previous_response_id` referring to the `id` of a previous request + * - Requires `"caching": {"type": "enabled"}` in the request body + * - Provides finer-grained control compared to the Context API + * - Supports going back to checkpoints (better for non-linear conversations) + * + * @see https://www.volcengine.com/docs/82379/1602228 + */ + +export interface ArkCacheOptions { + /** Previous response ID to reference for caching */ + previousResponseId?: string + /** Cache TTL in seconds (default: 3600 = 1 hour) */ + cacheTtl?: number +} + +/** + * Add context caching support for Ark/Volcengine using the Responses API + * + * @param requestOptions - The OpenAI request options to modify + * @param cacheOptions - Ark-specific caching options + */ +export function addArkCaching( + requestOptions: + | OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming + | OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming, + cacheOptions: ArkCacheOptions = {}, +): void { + // Enable caching for this request + ;(requestOptions as any).caching = { + type: "enabled", + } + + // If we have a previous response ID, reference it for context continuation + if (cacheOptions.previousResponseId) { + ;(requestOptions as any).previous_response_id = cacheOptions.previousResponseId + } + + // Set cache TTL (default to 1 hour as recommended in the issue) + if (cacheOptions.cacheTtl) { + ;(requestOptions as any).cache_ttl = cacheOptions.cacheTtl + } +} + +/** + * Extract response ID from Ark API response for future caching + * + * @param response - The API response from Ark/Volcengine + * @returns The response ID if available + */ +export function extractArkResponseId(response: any): string | undefined { + return response?.id +} + +/** + * Check if the response contains cached tokens information + * + * @param usage - Usage object from the API response + * @returns True if cached tokens are present + */ +export function hasArkCachedTokens(usage: any): boolean { + return usage?.prompt_tokens_details?.cached_tokens > 0 +} + +/** + * Extract cached tokens count from Ark usage metrics + * + * @param usage - Usage object from the API response + * @returns Number of cached tokens used + */ +export function getArkCachedTokens(usage: any): number { + return usage?.prompt_tokens_details?.cached_tokens || 0 +}