Browse Source

Add cache breakpoints for custom vertex models on Unbound (#4585)

Co-authored-by: Pugazhendhi <[email protected]>
pugazhendhi-m 8 months ago
parent
commit
2cb02d6512
1 changed files with 5 additions and 0 deletions
  1. 5 0
      src/api/providers/unbound.ts

+ 5 - 0
src/api/providers/unbound.ts

@@ -9,6 +9,7 @@ import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
 import { convertToOpenAiMessages } from "../transform/openai-format"
 import { addCacheBreakpoints as addAnthropicCacheBreakpoints } from "../transform/caching/anthropic"
 import { addCacheBreakpoints as addGeminiCacheBreakpoints } from "../transform/caching/gemini"
+import { addCacheBreakpoints as addVertexCacheBreakpoints } from "../transform/caching/vertex"
 
 import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
 import { RouterProvider } from "./router-provider"
@@ -70,6 +71,10 @@ export class UnboundHandler extends RouterProvider implements SingleCompletionHa
 				addAnthropicCacheBreakpoints(systemPrompt, openAiMessages)
 			}
 		}
+		// Custom models from Vertex AI (no configuration) need to be handled differently.
+		if (modelId.startsWith("vertex-ai/google.") || modelId.startsWith("vertex-ai/anthropic.")) {
+			addVertexCacheBreakpoints(messages)
+		}
 
 		// Required by Anthropic; other providers default to max tokens allowed.
 		let maxTokens: number | undefined