From 4eb8219bdd6f799f2deecb95af44ce890cc507c1 Mon Sep 17 00:00:00 2001 From: vswaroop04 Date: Thu, 7 May 2026 19:56:02 +0530 Subject: [PATCH 1/3] feat(semantic-cache): add Google AI (Gemini) embedding provider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds packages/semantic-cache/src/embed/google.ts and the matching packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py, following the existing cohere/voyage patterns. Both providers target the Google AI (Gemini) embedContent REST API: - Default model: text-embedding-004 (768-dim) - Configurable taskType (RETRIEVAL_QUERY default, RETRIEVAL_DOCUMENT for storage) - Optional title (improves retrieval quality for document task type) - Optional outputDimensionality for text-embedding-004+ truncation - Uses native fetch (TS) / httpx (Python) — no Google SDK required - API key from GOOGLE_API_KEY env var or explicit option --- .../betterdb_semantic_cache/embed/google.py | 90 +++++++++++++++++ packages/semantic-cache/src/embed/google.ts | 97 +++++++++++++++++++ 2 files changed, 187 insertions(+) create mode 100644 packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py create mode 100644 packages/semantic-cache/src/embed/google.ts diff --git a/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py b/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py new file mode 100644 index 00000000..def6124d --- /dev/null +++ b/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py @@ -0,0 +1,90 @@ +"""Google AI (Gemini) embedding helper for betterdb-semantic-cache. + +Uses the Google AI REST API directly via httpx. +Requires the 'httpx' extra: pip install betterdb-semantic-cache[httpx] + +Usage:: + + from betterdb_semantic_cache.embed.google import create_google_embed + embed = create_google_embed(model="text-embedding-004") + cache = SemanticCache(SemanticCacheOptions(client=client, embed_fn=embed)) +""" +from __future__ import annotations + +import os +from typing import Any, Literal + +from betterdb_semantic_cache.types import EmbedFn + +GoogleEmbedTaskType = Literal[ + "RETRIEVAL_QUERY", + "RETRIEVAL_DOCUMENT", + "SEMANTIC_SIMILARITY", + "CLASSIFICATION", + "CLUSTERING", +] + + +def create_google_embed( + *, + model: str = "text-embedding-004", + api_key: str | None = None, + base_url: str = "https://generativelanguage.googleapis.com/v1beta", + task_type: GoogleEmbedTaskType = "RETRIEVAL_QUERY", + title: str | None = None, + output_dimensionality: int | None = None, +) -> EmbedFn: + """Create an EmbedFn backed by the Google AI (Gemini) Embeddings API. + + Args: + model: Google AI embedding model. Default: 'text-embedding-004' (768-dim). + Other options: 'text-multilingual-embedding-002', 'embedding-001'. + api_key: Google AI API key. Default: GOOGLE_API_KEY env var. + base_url: API base URL. + task_type: Task type hint. Default: 'RETRIEVAL_QUERY'. + Use 'RETRIEVAL_DOCUMENT' when storing documents. + title: Optional document title. Only used with task_type='RETRIEVAL_DOCUMENT'. + output_dimensionality: Optional output dimensionality (truncation). + Supported by text-embedding-004+. + """ + _client: list[Any] = [] + + async def _get_client() -> Any: + if not _client: + try: + import httpx + except ImportError: + raise ImportError( + 'betterdb-semantic-cache embed/google requires the "httpx" package. ' + "Install it: pip install betterdb-semantic-cache[httpx]" + ) + _client.append(httpx.AsyncClient(timeout=30)) + return _client[0] + + async def embed(text: str) -> list[float]: + key = api_key or os.environ.get("GOOGLE_API_KEY") + if not key: + raise ValueError( + "Google API key is required. Set GOOGLE_API_KEY env var or pass api_key." + ) + client = await _get_client() + body: dict[str, Any] = { + "model": f"models/{model}", + "content": {"parts": [{"text": text}]}, + "taskType": task_type, + } + if title is not None: + body["title"] = title + if output_dimensionality is not None: + body["outputDimensionality"] = output_dimensionality + + resp = await client.post( + f"{base_url}/models/{model}:embedContent", + params={"key": key}, + headers={"Content-Type": "application/json"}, + json=body, + ) + resp.raise_for_status() + return resp.json().get("embedding", {}).get("values") or [] + + return embed diff --git a/packages/semantic-cache/src/embed/google.ts b/packages/semantic-cache/src/embed/google.ts new file mode 100644 index 00000000..08e8c4b4 --- /dev/null +++ b/packages/semantic-cache/src/embed/google.ts @@ -0,0 +1,97 @@ +/** + * Google AI (Gemini) embedding helper for @betterdb/semantic-cache. + * + * Supports text-embedding-004 and other Gemini embedding models via the + * Google AI REST API. Uses native fetch - no SDK required. + * + * Usage: + * import { createGoogleEmbed } from '@betterdb/semantic-cache/embed/google'; + * const embed = createGoogleEmbed({ model: 'text-embedding-004' }); + * const cache = new SemanticCache({ client, embedFn: embed }); + */ +import type { EmbedFn } from '../types'; + +export type GoogleEmbedTaskType = + | 'RETRIEVAL_QUERY' + | 'RETRIEVAL_DOCUMENT' + | 'SEMANTIC_SIMILARITY' + | 'CLASSIFICATION' + | 'CLUSTERING' + | (string & {}); + +export interface GoogleEmbedOptions { + /** + * Google AI embedding model. + * Default: 'text-embedding-004' (768 dimensions). + * Other options: 'text-multilingual-embedding-002', 'embedding-001'. + */ + model?: string; + /** Google AI (Gemini) API key. Default: GOOGLE_API_KEY env var. */ + apiKey?: string; + /** API base URL. Default: 'https://generativelanguage.googleapis.com/v1beta'. */ + baseUrl?: string; + /** + * Task type hint for the embedding. + * Default: 'RETRIEVAL_QUERY'. Use 'RETRIEVAL_DOCUMENT' when storing. + */ + taskType?: GoogleEmbedTaskType; + /** + * Optional document title, used only with taskType 'RETRIEVAL_DOCUMENT'. + * Improves retrieval quality when provided alongside the document body. + */ + title?: string; + /** + * Optional output dimensionality (truncation). Supported by text-embedding-004+. + * When omitted, the model's full dimensionality is returned. + */ + outputDimensionality?: number; +} + +/** + * Create an EmbedFn backed by the Google AI (Gemini) Embeddings API. + * Uses native fetch - no SDK required. + */ +export function createGoogleEmbed(opts?: GoogleEmbedOptions): EmbedFn { + const model = opts?.model ?? 'text-embedding-004'; + const baseUrl = opts?.baseUrl ?? 'https://generativelanguage.googleapis.com/v1beta'; + const taskType = opts?.taskType ?? 'RETRIEVAL_QUERY'; + + return async (text: string): Promise => { + const apiKey = opts?.apiKey ?? process.env.GOOGLE_API_KEY; + if (!apiKey) { + throw new Error( + 'Google API key is required. Set GOOGLE_API_KEY env var or pass apiKey in options.', + ); + } + + const requestBody: Record = { + model: `models/${model}`, + content: { parts: [{ text }] }, + taskType, + }; + + if (opts?.title) { + requestBody.title = opts.title; + } + if (opts?.outputDimensionality !== undefined) { + requestBody.outputDimensionality = opts.outputDimensionality; + } + + const res = await fetch( + `${baseUrl}/models/${model}:embedContent?key=${encodeURIComponent(apiKey)}`, + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(requestBody), + }, + ); + + if (!res.ok) { + const body = await res.text().catch(() => ''); + throw new Error(`Google AI API error: ${res.status} ${body}`); + } + + const json = (await res.json()) as { embedding: { values: number[] } }; + return json.embedding?.values ?? []; + }; +} From c083b2b3d156d9a72a731b77cb0bf48e162c83c6 Mon Sep 17 00:00:00 2001 From: vswaroop04 Date: Thu, 14 May 2026 00:29:35 +0530 Subject: [PATCH 2/3] fix(google-embed): address PR review issues - Move API key from URL query param to x-goog-api-key header (security) - Add close() helper to Python embed fn to clean up httpx client - Fix title check in TS to use !== undefined (consistent with Python is not None) --- .../betterdb_semantic_cache/embed/google.py | 9 +++++++-- packages/semantic-cache/src/embed/google.ts | 16 ++++++++-------- 2 files changed, 15 insertions(+), 10 deletions(-) diff --git a/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py b/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py index def6124d..2cf45f8c 100644 --- a/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py +++ b/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py @@ -80,11 +80,16 @@ async def embed(text: str) -> list[float]: resp = await client.post( f"{base_url}/models/{model}:embedContent", - params={"key": key}, - headers={"Content-Type": "application/json"}, + headers={"Content-Type": "application/json", "x-goog-api-key": key}, json=body, ) resp.raise_for_status() return resp.json().get("embedding", {}).get("values") or [] + async def close() -> None: + if _client: + await _client[0].aclose() + _client.clear() + + embed.close = close # type: ignore[attr-defined] return embed diff --git a/packages/semantic-cache/src/embed/google.ts b/packages/semantic-cache/src/embed/google.ts index 08e8c4b4..4d029e01 100644 --- a/packages/semantic-cache/src/embed/google.ts +++ b/packages/semantic-cache/src/embed/google.ts @@ -70,21 +70,21 @@ export function createGoogleEmbed(opts?: GoogleEmbedOptions): EmbedFn { taskType, }; - if (opts?.title) { + if (opts?.title !== undefined) { requestBody.title = opts.title; } if (opts?.outputDimensionality !== undefined) { requestBody.outputDimensionality = opts.outputDimensionality; } - const res = await fetch( - `${baseUrl}/models/${model}:embedContent?key=${encodeURIComponent(apiKey)}`, - { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(requestBody), + const res = await fetch(`${baseUrl}/models/${model}:embedContent`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-goog-api-key': apiKey, }, - ); + body: JSON.stringify(requestBody), + }); if (!res.ok) { const body = await res.text().catch(() => ''); From 4a49295054c9c3a65d075f093aaf719338f1b231 Mon Sep 17 00:00:00 2001 From: vswaroop04 Date: Fri, 15 May 2026 01:18:47 +0530 Subject: [PATCH 3/3] docs(google-embed): document close() helper in create_google_embed docstring --- .../semantic-cache-py/betterdb_semantic_cache/embed/google.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py b/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py index 2cf45f8c..7eda1bd8 100644 --- a/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py +++ b/packages/semantic-cache-py/betterdb_semantic_cache/embed/google.py @@ -46,6 +46,10 @@ def create_google_embed( title: Optional document title. Only used with task_type='RETRIEVAL_DOCUMENT'. output_dimensionality: Optional output dimensionality (truncation). Supported by text-embedding-004+. + + When finished, release the connection pool:: + + await embed.close() """ _client: list[Any] = []