Skip to content

feat: add Google AI (Gemini) embeddings for files and passages - #3181

Open
wayne-accelerandos wants to merge 1 commit into
letta-ai:mainfrom
wayne-accelerandos:feature/google-embeddings-files
Open

feat: add Google AI (Gemini) embeddings for files and passages#3181
wayne-accelerandos wants to merge 1 commit into
letta-ai:mainfrom
wayne-accelerandos:feature/google-embeddings-files

Conversation

@wayne-accelerandos

Copy link
Copy Markdown

Problem

When using Google AI embedding models (e.g., google_ai/gemini-embedding-001), two separate code paths were broken:

  1. Passages/archival memoryrequest_embeddings in GoogleVertexClient raised NotImplementedError, causing 500 errors on any archival memory operation with a Google AI embedding config.

  2. Files/sourcessources.py, folders.py, and agent_serialization_manager.py always used OpenAIEmbedder regardless of the source's embedding_config. This silently produced OpenAI embedding vectors for file chunks. Since semantic_search_files queries use the agent's Google AI embedding config, search vectors never matched indexed file vectors, returning empty or incorrect results.

Changes

1. Implement GoogleVertexClient.request_embeddings() (letta/llm_api/google_vertex_client.py)

  • Replaces empty/whitespace strings with single space (Google AI rejects empty strings)
  • Batches at max 100 texts per call (Google AI limit vs OpenAI's 2048)
  • Parallel batch processing via asyncio.gather()

Both GoogleAIClient and GoogleVertexClient inherit this implementation.

2. Add GoogleAIEmbedder (letta/services/file_processor/embedder/google_ai_embedder.py)

New embedder for the file processing pipeline, mirroring the OpenAIEmbedder interface:

  • Global semaphore (3 concurrent requests) to avoid API quota exhaustion
  • Recursive halving retry on batch failure (same strategy as OpenAIEmbedder)
  • OpenTelemetry tracing via log_event

3. Route Google AI through GoogleAIEmbedder (sources.py, folders.py, agent_serialization_manager.py)

Added elif embedding_config.embedding_endpoint_type == "google_ai" before the OpenAI fallback in the existing embedder selection chain: Turbopuffer → Pinecone → Google AI → OpenAI.

Why the embedder must match

The embedding model used to index file chunks must be the same model used to embed search queries. Mixing OpenAI vectors (file index) with Google AI vectors (search query) produces semantically meaningless results — the vector spaces are incompatible.

Tests Added

  • tests/test_google_ai_file_embedder.py — 12 unit tests for GoogleAIEmbedder (all mocked, no API key needed)
  • tests/test_embeddings.py — 14 new unit tests for request_embeddings() on both GoogleAIClient and GoogleVertexClient (all mocked)
  • tests/test_google_ai_file_embeddings_integration.py — 4 integration tests (require GEMINI_API_KEY, skipped otherwise)
  • tests/configs/embedding_model_configs/google_ai_embed.json — config for parametrized embedding tests

Note: There are 3 pre-existing failures in tests/test_file_processor.py unrelated to this PR.

Running the tests

# Unit tests (no API key needed)
uv run pytest tests/test_google_ai_file_embedder.py tests/test_embeddings.py -m "unit and google" -v

# Integration tests (requires GEMINI_API_KEY)
GEMINI_API_KEY=your_key uv run pytest tests/test_google_ai_file_embeddings_integration.py -v

Problem:
- GoogleVertexClient.request_embeddings() raised NotImplementedError,
  causing 500 errors when creating passages with Google AI models
- File processing pipeline always used OpenAIEmbedder regardless of the
  source's embedding_config, silently indexing files with OpenAI vectors.
  Since semantic_search_files queries use the agent's Google AI embedding
  config, search vectors never matched indexed file vectors.

Changes:
- Implement GoogleVertexClient.request_embeddings(): batches at 100 (Google
  AI limit), sanitizes empty strings, uses asyncio.gather() for parallel
  batch processing. GoogleAIClient inherits this implementation.
- Add GoogleAIEmbedder: mirrors OpenAIEmbedder interface, uses global
  semaphore (3 concurrent), recursive halving retry on batch failure
- Route google_ai embedding configs to GoogleAIEmbedder in sources.py,
  folders.py, and agent_serialization_manager.py

Tests:
- tests/test_google_ai_file_embedder.py: 12 unit tests for GoogleAIEmbedder
- tests/test_embeddings.py: +14 unit tests for request_embeddings() on both
  GoogleAIClient and GoogleVertexClient (all mocked, no API key needed)
- tests/test_google_ai_file_embeddings_integration.py: 4 integration tests
  (require GEMINI_API_KEY, skipped otherwise)
- tests/configs/embedding_model_configs/google_ai_embed.json: config for
  parametrized embedding tests
huangyingting pushed a commit to repomesh/letta that referenced this pull request May 17, 2026
Co-authored-by: Jin Peng <jinjpeng@Jins-MacBook-Pro.local>
Co-authored-by: cpacker <packercharles@gmail.com>
huangyingting pushed a commit to repomesh/letta that referenced this pull request May 17, 2026
Co-authored-by: Jin Peng <jinjpeng@Jins-MacBook-Pro.local>
Co-authored-by: cpacker <packercharles@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants