feature(go/plugins/vertexai): add support for context caching in VertexAI and GoogleAI plugins #1566

falonso81 · 2024-12-20T19:39:23Z

The following changes are included in this PR:

Added Context Cache for VertexAI and GeminiAI plugins
Context Cache Samples for both plugins
Unit tests

Checklist (if applicable):

PR title is following https://www.conventionalcommits.org/en/v1.0.0/
Tested (manually, unit tested, etc.)
Docs updated (updated docs or a docs bug required)

hugoaguirre · 2025-03-03T18:28:20Z

go/plugins/googleai/googleai.go


 //copy:stop

-//copy:start vertexai.go translateResponse


VertexAI should not replicate this function anymore since u.CachedContentTokenCount (@ line 558) is not being provided in the response struct (see https://pkg.go.dev/cloud.google.com/go/vertexai/genai#UsageMetadata)

hugoaguirre · 2025-03-04T00:48:18Z

go/plugins/googleai/cache.go

+		return nil, err
+	}
+
+	return client.CreateCachedContent(ctx, cc)


A thing to mention is that in JS we do a cache lookup based on the displayName (https://github.com/google-gemini/generative-ai-js/blob/0c8d8cde00a33ea9da94e5d1bb033dc261d0588f/types/server/caching.ts#L35) of the CachedContent in Go, this field is not available; making us unable to perform such operations.

apascal07

Overall, this follows the JS implementation a little too closely with too much overly complicated logic. Let's keep it simple for now, focusing only on the cache marker + TTL in the metadata logic. If we ever add more fields, we can return a struct. It's all internal functions so there won't be breaking changes.

apascal07 · 2025-03-04T05:01:22Z

go/ai/generate.go


 // WithTextPrompt adds a simple text user prompt to ModelRequest.
-func WithTextPrompt(prompt string) GenerateOption {
+func WithTextPrompt(prompt string, ttl ...int) GenerateOption {


Why is this modified and why is it a variadic arg?

Good catch, this shouldn't be there

go/ai/generate.go

apascal07 · 2025-03-04T05:03:10Z

go/ai/request_helpers.go

+		},
+	}
+	// copy the existing metadata and add cache
+	if m.Metadata != nil {


This does the opposite of what it says here. It will overwrite cache with whatever was in there already. It makes sense for this deliberate call to overwrite any existing metadata.

apascal07 · 2025-03-04T05:05:33Z

go/plugins/googleai/cache.go

+}
+
+var ContextCacheSupportedModels = [...]string{
+	"gemini-1.5-flash-001",


gemini-1.5-flash-002, gemini-1.5-pro-002, and Gemini 2.0 don't support caching?

you are right, added gemini-2.0 stable models as well

apascal07 · 2025-03-04T05:11:21Z

go/plugins/vertexai/cache.go

+		return nil, nil
+	}
+	if endOfCachedContents < 0 || endOfCachedContents >= len(request.Messages) {
+		return nil, fmt.Errorf("invalid endOfCachedContents index")


Suggested change

return nil, fmt.Errorf("invalid endOfCachedContents index")

return nil, fmt.Errorf("end of cached content index %q is invalid", cacheEndIdx)

apascal07 · 2025-03-04T05:25:39Z

go/ai/gen.go

+	Temperature     float64       `json:"temperature,omitempty"`
+	TopK            int           `json:"topK,omitempty"`
+	TopP            float64       `json:"topP,omitempty"`
+	TTL             time.Duration `json:"ttl,omitempty"`


This is a generated file, there's no corresponding TTL field in the source:

genkit/genkit-tools/common/src/types/model.ts

Line 244 in 54d439e

export const GenerationCommonConfigSchema = z.object({

Yeah, after addressing your comments, TTL config should not be in this place but in Message.Metadata instead

apascal07 · 2025-03-04T05:30:41Z

go/ai/gen.go

 	OutputTokens     int                `json:"outputTokens,omitempty"`
 	OutputVideos     float64            `json:"outputVideos,omitempty"`
 	TotalTokens      int                `json:"totalTokens,omitempty"`
+	CachedTokens     int                `json:"cachedTokens,omitempty"`


Same with this. We can probably either add it as a field or add it to Custom.

apascal07 · 2025-03-04T05:44:03Z

go/plugins/googleai/cache.go

+)
+
+// DEFAULT_TTL in seconds (5 minutes)
+const DEFAULT_TTL = 300


There's an inconsistent mix of seconds int and time.Duration throughout.

fixed: ttlSeconds is always int and conversion to time.Duration is done right before detecting the cache mark

go/plugins/googleai/cache.go

hugoaguirre · 2025-03-04T22:36:15Z

Moving efforts to address issue #2159 first since VertexAI plugin does not fully support Context Caching. When completing #2159, this PR should be modified to use the new SDK functions and remove redundant code between VertexAI and GoogleAI plugins.

cc @yesudeep @apascal07

hugoaguirre · 2025-03-07T20:42:39Z

Since we are moving forward on using the unified go-genai SDK in GooleAI and VertexAI plugins, this PR will be closed and moved to #2266

falonso81 added 9 commits December 17, 2024 18:21

feat: WIP add cache context

1fe8ac3

Fix: remove keys

09596fb

WIP add context cache

22a49ae

WIP: add context caching

2b6f573

Feat: add endOfCachedContents

5f560e7

Fix: remove stuff

3901616

Feat: add tests

0ad4740

Feat: add tests 2

f3828f3

Merge branch 'main' into feat/create-context-caching-for-vertexai

18c5cab

falonso81 changed the title ~~[WIP] Feat: create context caching for vertexai~~ Dec 27, 2024

Merge branch 'main' into feat/create-context-caching-for-vertexai

fce1d81

github-actions bot added go config root labels Feb 4, 2025

falonso81 added 3 commits February 4, 2025 17:33

Fix: copyright statement missing

7131a51

Fix: go func params

73b3db1

Fix: var names warnings

be381e8

This was referenced Feb 5, 2025

feat(go/plugins/vertexai): add context caching to vertexai #1478

Closed

feat(go/plugins/googleai): add context caching to googleai #1518

Closed

falonso81 linked an issue Feb 5, 2025 that may be closed by this pull request

[Go] [Vertex AI] [Google AI] Add context caching support #1443

Closed

falonso81 requested a review from yesudeep February 5, 2025 03:26

yesudeep requested a review from hugoaguirre February 7, 2025 19:06

hugoaguirre force-pushed the feat/create-context-caching-for-vertexai branch from 22a2be5 to be381e8 Compare February 20, 2025 20:09

hugoaguirre added 3 commits February 24, 2025 22:08

Merge branch 'main' into feat/create-context-caching-for-vertexai

d4a0d83

docs: update licences

e9d2305

fix: remove unnecessary code

b73a7cf

hugoaguirre self-assigned this Feb 25, 2025

hugoaguirre added 3 commits February 25, 2025 20:23

Merge branch 'main' into feat/create-context-caching-for-vertexai

4c41930

feat: add cache sample

96293e4

fix: vertexai cache sample

636138d

hugoaguirre added 6 commits February 27, 2025 06:20

feat: working context cache

18bf748

fix: refine some code

d8b9e5f

fix: working test cases

fca6f10

feat: gemini context cache handlers

2116440

feat: add context cache in googleai plugin

2e8f308

Merge branch 'main' into feat/create-context-caching-for-vertexai

54e2076

hugoaguirre requested a review from apascal07 March 3, 2025 17:17

remove logs and fix naming

d7e09c7

hugoaguirre reviewed Mar 3, 2025

View reviewed changes

hugoaguirre changed the title ~~feature(go/plugins/vertexai): create context caching for vertexai~~ Mar 3, 2025

hugoaguirre changed the title ~~feature(go/plugins/vertexai): create context caching for VertexAI and GoogleAI plugins~~ Mar 3, 2025

hugoaguirre changed the title ~~feature(go/plugins/vertexai): add support for context caching for VertexAI and GoogleAI plugins~~ Mar 3, 2025

fix: samples and remove dead comments

2a839be

hugoaguirre reviewed Mar 4, 2025

View reviewed changes

apascal07 requested changes Mar 4, 2025

View reviewed changes

hugoaguirre added 2 commits March 4, 2025 20:37

fix: address review comments

2a7d605

fix: add vertexai changes

5aa0fa3

hugoaguirre closed this Mar 7, 2025

github-project-automation bot moved this from In Progress to Done in Genkit Backlog Mar 7, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feature(go/plugins/vertexai): add support for context caching in VertexAI and GoogleAI plugins #1566

feature(go/plugins/vertexai): add support for context caching in VertexAI and GoogleAI plugins #1566

Uh oh!

falonso81 commented Dec 20, 2024 •

edited by hugoaguirre

Loading

hugoaguirre Mar 3, 2025

hugoaguirre Mar 4, 2025 •

edited

Loading

apascal07 left a comment

apascal07 Mar 4, 2025

hugoaguirre Mar 4, 2025

Uh oh!

apascal07 Mar 4, 2025

apascal07 Mar 4, 2025

hugoaguirre Mar 4, 2025

apascal07 Mar 4, 2025

apascal07 Mar 4, 2025

hugoaguirre Mar 4, 2025

apascal07 Mar 4, 2025

apascal07 Mar 4, 2025

hugoaguirre Mar 4, 2025

Uh oh!

Uh oh!

hugoaguirre commented Mar 4, 2025

hugoaguirre commented Mar 7, 2025

Labels

3 participants

	return nil, fmt.Errorf("invalid endOfCachedContents index")
	return nil, fmt.Errorf("end of cached content index %q is invalid", cacheEndIdx)

feature(go/plugins/vertexai): add support for context caching in VertexAI and GoogleAI plugins #1566

feature(go/plugins/vertexai): add support for context caching in VertexAI and GoogleAI plugins #1566

Uh oh!

Conversation

falonso81 commented Dec 20, 2024 • edited by hugoaguirre Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

hugoaguirre Mar 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

apascal07 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hugoaguirre commented Mar 4, 2025

hugoaguirre commented Mar 7, 2025

Labels

3 participants

falonso81 commented Dec 20, 2024 •

edited by hugoaguirre

Loading

hugoaguirre Mar 4, 2025 •

edited

Loading