Skip to content

Commit 233b328

Browse files
committed
Update docs-vectorize to the embeddinggemma model
1 parent 577eacf commit 233b328

File tree

4 files changed

+17
-17
lines changed

4 files changed

+17
-17
lines changed

apps/docs-vectorize/wrangler.jsonc

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
"vectorize": [
3232
{
3333
"binding": "VECTORIZE",
34-
"index_name": "docs-bge-base"
34+
"index_name": "docs-embeddinggemma-v1"
3535
}
3636
],
3737
"vars": {
@@ -72,7 +72,7 @@
7272
"vectorize": [
7373
{
7474
"binding": "VECTORIZE",
75-
"index_name": "docs-bge-base"
75+
"index_name": "docs-embeddinggemma-v1"
7676
}
7777
],
7878
"analytics_engine_datasets": [
@@ -103,7 +103,7 @@
103103
"vectorize": [
104104
{
105105
"binding": "VECTORIZE",
106-
"index_name": "docs-bge-base"
106+
"index_name": "docs-embeddinggemma-v1"
107107
}
108108
],
109109
"analytics_engine_datasets": [

apps/workers-bindings/wrangler.jsonc

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@
5151
"vectorize": [
5252
{
5353
"binding": "VECTORIZE",
54-
"index_name": "docs-bge-base"
54+
"index_name": "docs-embeddinggemma-v1"
5555
}
5656
],
5757
"workers_dev": false,
@@ -92,7 +92,7 @@
9292
"vectorize": [
9393
{
9494
"binding": "VECTORIZE",
95-
"index_name": "docs-bge-base"
95+
"index_name": "docs-embeddinggemma-v1"
9696
}
9797
],
9898
"vars": {
@@ -141,7 +141,7 @@
141141
"vectorize": [
142142
{
143143
"binding": "VECTORIZE",
144-
"index_name": "docs-bge-base"
144+
"index_name": "docs-embeddinggemma-v1"
145145
}
146146
],
147147
"analytics_engine_datasets": [

apps/workers-observability/wrangler.jsonc

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
"vectorize": [
4242
{
4343
"binding": "VECTORIZE",
44-
"index_name": "docs-bge-base"
44+
"index_name": "docs-embeddinggemma-v1"
4545
}
4646
],
4747
"vars": {
@@ -98,7 +98,7 @@
9898
"vectorize": [
9999
{
100100
"binding": "VECTORIZE",
101-
"index_name": "docs-bge-base"
101+
"index_name": "docs-embeddinggemma-v1"
102102
}
103103
],
104104
"analytics_engine_datasets": [
@@ -136,7 +136,7 @@
136136
"vectorize": [
137137
{
138138
"binding": "VECTORIZE",
139-
"index_name": "docs-bge-base"
139+
"index_name": "docs-embeddinggemma-v1"
140140
}
141141
],
142142
"vars": {

packages/mcp-common/src/tools/docs-vectorize.tools.ts

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ export function registerDocsTools(agent: CloudflareMcpAgentNoAccount, env: Requi
4242
.map((result) => {
4343
return `<result>
4444
<url>${result.url}</url>
45+
<title>${result.title}</title>
4546
<text>
4647
${result.text}
4748
</text>
@@ -93,10 +94,8 @@ ${result.text}
9394
}
9495

9596
async function queryVectorize(ai: Ai, vectorizeIndex: VectorizeIndex, query: string, topK: number) {
96-
// Recommendation from: https://huggingface.co/BAAI/bge-base-en-v1.5#model-list
97-
const [queryEmbedding] = await getEmbeddings(ai, [
98-
'Represent this sentence for searching relevant passages: ' + query,
99-
])
97+
// Recommendation from: https://ai.google.dev/gemma/docs/embeddinggemma/model_card#prompt_instructions
98+
const [queryEmbedding] = await getEmbeddings(ai, ['task: search result | query: ' + query])
10099

101100
const { matches } = await vectorizeIndex.query(queryEmbedding, {
102101
topK,
@@ -108,6 +107,7 @@ async function queryVectorize(ai: Ai, vectorizeIndex: VectorizeIndex, query: str
108107
similarity: Math.min(match.score, 1),
109108
id: match.id,
110109
url: sourceToUrl(String(match.metadata?.filePath ?? '')),
110+
title: String(match.metadata?.title ?? ''),
111111
text: String(match.metadata?.text ?? ''),
112112
}))
113113
}
@@ -123,15 +123,15 @@ function sourceToUrl(path: string) {
123123
)
124124
}
125125

126-
async function getEmbeddings(ai: Ai, strings: string[]) {
126+
async function getEmbeddings(ai: Ai, strings: string[]): Promise<number[][]> {
127127
const response = await doWithRetries(() =>
128-
ai.run('@cf/baai/bge-base-en-v1.5', {
128+
// @ts-expect-error embeddinggemma not in types yet
129+
ai.run('@cf/google/embeddinggemma-300m', {
129130
text: strings,
130-
// @ts-expect-error pooling not in types yet
131-
pooling: 'cls',
132131
})
133132
)
134133

134+
// @ts-expect-error embeddinggemma not in types yet
135135
return response.data
136136
}
137137

0 commit comments

Comments
 (0)