scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-09-06 08:18:39 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

vertex prompt caching (#2026)

Browse source 
* vertex prompt caching

* duplicate declare

* changeset
This commit is contained in:
pashpashpash 2025-02-28 17:30:25 -08:00 • committed by GitHub
parent cee959eeec
commit b4149b5a18
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 118 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
.changeset/six-years-sip.md Normal file
View file

@ -0,0 +1,5 @@
---
"claude-dev": minor
---
re-added prompt caching to vertex

122
src/api/providers/vertex.ts
View file

@ -21,14 +21,113 @@ export class VertexHandler implements ApiHandler {
@withRetry()
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const stream = await this.client.messages.create({
model: this.getModel().id,
max_tokens: this.getModel().info.maxTokens || 8192,
temperature: 0,
system: systemPrompt,
messages,
stream: true,
})
const model = this.getModel()
const modelId = model.id
let stream
switch (modelId) {
case "claude-3-7-sonnet@20250219":
case "claude-3-5-sonnet-v2@20241022":
case "claude-3-5-sonnet@20240620":
case "claude-3-5-haiku@20241022":
case "claude-3-opus@20240229":
case "claude-3-haiku@20240307": {
// Find indices of user messages for cache control
const userMsgIndices = messages.reduce(
(acc, msg, index) => (msg.role === "user" ? [...acc, index] : acc),
[] as number[],
)
const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
stream = await this.client.beta.messages.create(
{
model: modelId,
max_tokens: model.info.maxTokens || 8192,
temperature: 0,
system: [
{
text: systemPrompt,
type: "text",
cache_control: { type: "ephemeral" },
},
],
messages: messages.map((message, index) => {
if (index === lastUserMsgIndex || index === secondLastMsgUserIndex) {
return {
...message,
content:
typeof message.content === "string"
? [
{
type: "text",
text: message.content,
cache_control: {
type: "ephemeral",
},
},
]
: message.content.map((content, contentIndex) =>
contentIndex === message.content.length - 1
? {
...content,
cache_control: {
type: "ephemeral",
},
}
: content,
),
}
}
return {
...message,
content:
typeof message.content === "string"
? [
{
type: "text",
text: message.content,
},
]
: message.content,
}
}),
stream: true,
},
{
headers: {},
},
)
break
}
default: {
stream = await this.client.beta.messages.create({
model: modelId,
max_tokens: model.info.maxTokens || 8192,
temperature: 0,
system: [
{
text: systemPrompt,
type: "text",
},
],
messages: messages.map((message) => ({
...message,
content:
typeof message.content === "string"
? [
{
type: "text",
text: message.content,
},
]
: message.content,
})),
stream: true,
})
break
}
}
for await (const chunk of stream) {
switch (chunk.type) {
case "message_start":
@ -37,6 +136,8 @@ export class VertexHandler implements ApiHandler {
type: "usage",
inputTokens: usage.input_tokens || 0,
outputTokens: usage.output_tokens || 0,
cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
cacheReadTokens: usage.cache_read_input_tokens || undefined,
}
break
case "message_delta":
@ -46,7 +147,8 @@ export class VertexHandler implements ApiHandler {
outputTokens: chunk.usage.output_tokens || 0,
}
break
case "message_stop":
break
case "content_block_start":
switch (chunk.content_block.type) {
case "text":
@ -73,6 +175,8 @@ export class VertexHandler implements ApiHandler {
break
}
break
case "content_block_stop":
break
}
}
}