RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix Gemini batch embedding for gemini-embedding-2* models (#745) 

* Fix Gemini batch embedding for gemini-embedding-2* models

The google-genai SDK treats embed_content(contents=[list_of_strings]) as a
single multi-part document for gemini-embedding-2* models, silently returning
exactly 1 embedding regardless of input count. This caused a zip(...,
strict=True) ValueError in _process_batch.

Wrap each text in genai_types.Content(parts=[genai_types.Part(text=...)])
so the SDK treats each string as a separate content item. This matches the
workaround used by pydantic-ai (#4873) and graphiti (#1474).

Upstream SDK issue: googleapis/python-genai#2523
Fixes plastic-labs/honcho#744

* test(embedding): live embedding coverage for every Gemini and OpenAI model

Adds tests/live_llm/test_live_embeddings.py plus an env-driven embedding
matrix alongside the existing LLM one. Covers single embed, batched embed,
batch-vs-single alignment, chunk-to-id mapping, and the batch-split path.

Only a live call catches the gemini-embedding-2* collapse: the SDK folds a
list of bare strings into one document and returns a single embedding.
Reverting the Content wrapping fails all four batch tests for
gemini-embedding-2-preview and gemini-embedding-2 with the reported
`zip() argument 2 is shorter than argument 1`, while gemini-embedding-001
and text-embedding-3-small stay green.

Also makes the concatenation in the conclusions semantic-search validation
message explicit, so the repo-wide basedpyright pre-push hook passes.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* test(embedding): drop the preview twin from the default embedding matrix

gemini-embedding-2 is the GA release of gemini-embedding-2-preview and
behaves identically, so running both by default doubles the Gemini cost for
no extra coverage. The preview stays reachable through
LIVE_EMBEDDING_GEMINI_MODELS.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

---------

Co-authored-by: Aakash Kattelu <aakash@plasticlabs.ai>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Niyaz Almufti 2026-08-12 00:50:40 +03:00 committed by GitHub
parent 2d174a5fb2
commit a92fb1e078
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 371 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/embedding_client.py
View File

@ -453,7 +453,13 @@ class _EmbeddingClient:
if isinstance(self.client, genai.Client):
response = await self.client.aio.models.embed_content(
model=self.model,
contents=[item.text for item in batch],
# One Content per item: a list of bare strings is folded
# into a single document by gemini-embedding-2*, which
# returns one embedding for the whole batch (#745).
contents=[
genai_types.Content(parts=[genai_types.Part(text=item.text)])
for item in batch
],
config={"output_dimensionality": self.vector_dimensions},
)
if response.embeddings:

6
src/routers/conclusions.py
View File

@ -111,9 +111,9 @@ async def query_conclusions(
if not observer or not observed:
raise ValidationException(
"observer and observed must be specified for semantic search. "
"Pass them inside the 'filters' object, e.g. "
'{"query": "...", "filters": {"observer": "alice", "observed": "bob"}}. '
"Both 'observer'/'observer_id' and 'observed'/'observed_id' are accepted."
+ "Pass them inside the 'filters' object, e.g. "
+ '{"query": "...", "filters": {"observer": "alice", "observed": "bob"}}. '
+ "Both 'observer'/'observer_id' and 'observed'/'observed_id' are accepted."
)
with embedding_call_purpose(

6
tests/live_llm/README.md
View File

@ -24,6 +24,11 @@ Model-family env vars:
- `LIVE_LLM_GEMINI_30_MODELS`
- `LIVE_LLM_GEMINI_31_MODELS`
Embedding-model env vars:
- `LIVE_EMBEDDING_GEMINI_MODELS` (default: `gemini-embedding-001,gemini-embedding-2`; add `gemini-embedding-2-preview` to cover the preview twin)
- `LIVE_EMBEDDING_OPENAI_MODELS` (default: `text-embedding-3-small`)
Each model env var accepts a comma-separated list of bare model ids or provider-qualified ids.
Examples:
@ -57,3 +62,4 @@ Coverage by provider:
- OpenAI transport → OpenRouter non-reasoning models (e.g. `inception/mercury-2`): non-chat / diffusion architectures must stay on `max_tokens`, no `reasoning_effort`, tool-calling parameter-schema compatibility is the canary for exotic OR-served providers
- Gemini 2.5/3.0 classes: structured outputs, cached-content reuse, thought signatures, multi-turn tool replay
- Gemini 3.1 class: thinking and tool replay coverage by default; structured-output/caching coverage should only be added once Google documents support for that path
- Embeddings (`test_live_embeddings.py`): single embed, batched embed, batch-vs-single alignment, and chunk-to-id mapping for every configured embedding model. `gemini-embedding-2*` is the reason this exists — those models collapse a list of bare strings into one document (#745), and only a live call catches it

56
tests/live_llm/conftest.py
View File

@ -6,10 +6,12 @@ from typing import Any
import pytest
from pydantic import BaseModel
from src.config import ModelConfig, settings
from src.config import EmbeddingModelConfig, ModelConfig, settings
from src.embedding_client import _EmbeddingClient # pyright: ignore[reportPrivateUsage]
from src.llm import get_backend
from src.llm.caching import gemini_cache_store
from .embedding_matrix import LiveEmbeddingSpec, selected_embedding_summary_lines
from .model_matrix import LiveModelSpec, selected_model_summary_lines
@ -22,9 +24,12 @@ class StructuredLiveResponse(BaseModel):
def pytest_report_header(config: pytest.Config) -> list[str] | None:
if not config.getoption("--live-llm"):
return None
return ["live llm model matrix:"] + [
f" {line}" for line in selected_model_summary_lines()
]
return (
["live llm model matrix:"]
+ [f" {line}" for line in selected_model_summary_lines()]
+ ["live embedding model matrix:"]
+ [f" {line}" for line in selected_embedding_summary_lines()]
)
@pytest.fixture(autouse=True)
@ -45,6 +50,49 @@ def require_provider_key(model_spec: LiveModelSpec) -> None:
pytest.skip(f"Missing API key for live provider {model_spec.provider}")
def require_embedding_key(spec: LiveEmbeddingSpec) -> str:
key = {
"openai": settings.LLM.OPENAI_API_KEY,
"gemini": settings.LLM.GEMINI_API_KEY,
}[spec.transport]
if not key:
pytest.skip(f"Missing API key for live embedding transport {spec.transport}")
return key
def make_embedding_client(
spec: LiveEmbeddingSpec, **overrides: Any
) -> _EmbeddingClient:
"""Build a live embedding client for one matrix entry.
Bypasses the `EmbeddingClient` singleton so each spec gets its own client
without mutating global settings.
"""
kwargs: dict[str, Any] = {
"vector_dimensions": spec.dimensions,
"max_input_tokens": 2048,
"max_tokens_per_request": 300_000,
# Both providers accept an explicit dimension request for these models.
"send_dimensions": True,
}
kwargs.update(overrides)
return _EmbeddingClient(
EmbeddingModelConfig(
transport=spec.transport,
model=spec.model,
api_key=require_embedding_key(spec),
),
**kwargs,
)
def cosine_similarity(a: list[float], b: list[float]) -> float:
dot = sum(x * y for x, y in zip(a, b, strict=True))
norm_a = sum(x * x for x in a) ** 0.5
norm_b = sum(y * y for y in b) ** 0.5
return dot / (norm_a * norm_b)
def make_model_config(model_spec: LiveModelSpec, **overrides: Any) -> ModelConfig:
return ModelConfig(
model=model_spec.model,

95
tests/live_llm/embedding_matrix.py Normal file
View File

@ -0,0 +1,95 @@
from __future__ import annotations
import os
from dataclasses import dataclass
from src.config import EmbeddingTransport
@dataclass(frozen=True)
class LiveEmbeddingFamily:
transport: EmbeddingTransport
family: str
env_var: str
dimensions: int
default_models: tuple[str, ...] = ()
docs_url: str | None = None
@dataclass(frozen=True)
class LiveEmbeddingSpec:
transport: EmbeddingTransport
family: str
model: str
env_var: str
dimensions: int
docs_url: str | None = None
@property
def id(self) -> str:
return f"{self.transport}:{self.family}:{self.model}"
EMBEDDING_FAMILIES: tuple[LiveEmbeddingFamily, ...] = (
# gemini-embedding-2 is the regression surface for #745: the SDK folds a
# list of bare strings into a single document and returns one embedding for
# the whole batch. Its preview twin behaves identically and is reachable
# through the env var when it needs checking.
LiveEmbeddingFamily(
transport="gemini",
family="gemini_embedding",
env_var="LIVE_EMBEDDING_GEMINI_MODELS",
# Matryoshka dimension supported across the family; keeps vectors small.
dimensions=768,
default_models=(
"gemini-embedding-001",
"gemini-embedding-2",
),
docs_url="https://ai.google.dev/gemini-api/docs/embeddings",
),
LiveEmbeddingFamily(
transport="openai",
family="openai_embedding",
env_var="LIVE_EMBEDDING_OPENAI_MODELS",
dimensions=1536,
default_models=("text-embedding-3-small",),
docs_url="https://platform.openai.com/docs/guides/embeddings",
),
)
def _parse_env_models(value: str | None) -> tuple[str, ...]:
if value is None:
return ()
return tuple(model.strip() for model in value.split(",") if model.strip())
def get_live_embedding_specs(
*, transport: EmbeddingTransport | None = None
) -> tuple[LiveEmbeddingSpec, ...]:
specs: list[LiveEmbeddingSpec] = []
for family in EMBEDDING_FAMILIES:
if transport is not None and family.transport != transport:
continue
models = _parse_env_models(os.getenv(family.env_var)) or family.default_models
for model in models:
specs.append(
LiveEmbeddingSpec(
transport=family.transport,
family=family.family,
model=model,
env_var=family.env_var,
dimensions=family.dimensions,
docs_url=family.docs_url,
)
)
return tuple(specs)
def selected_embedding_summary_lines() -> list[str]:
lines: list[str] = []
for family in EMBEDDING_FAMILIES:
models = _parse_env_models(os.getenv(family.env_var)) or family.default_models
joined = ", ".join(models) if models else "(none configured)"
lines.append(f"{family.env_var} [{family.transport}/{family.family}]: {joined}")
return lines

130
tests/live_llm/test_live_embeddings.py Normal file
View File

@ -0,0 +1,130 @@
from __future__ import annotations
import pytest
from .conftest import cosine_similarity, make_embedding_client
from .embedding_matrix import LiveEmbeddingSpec, get_live_embedding_specs
pytestmark = pytest.mark.live_llm
# Deliberately unrelated topics so a mix-up between them is visible in cosine
# similarity rather than lost in noise.
BATCH_TEXTS: list[str] = [
"The mitochondria generates ATP through oxidative phosphorylation.",
"Barcelona won the treble in the 2014-15 football season.",
"Sourdough starter needs regular feeding with flour and water.",
"Rust's borrow checker enforces ownership rules at compile time.",
]
ALL_SPECS = get_live_embedding_specs()
GEMINI_SPECS = get_live_embedding_specs(transport="gemini")
@pytest.mark.asyncio
@pytest.mark.parametrize("spec", ALL_SPECS, ids=lambda spec: spec.id)
async def test_live_embed_single_returns_configured_dimensions(
spec: LiveEmbeddingSpec,
) -> None:
client = make_embedding_client(spec)
embedding = await client.embed(BATCH_TEXTS[0])
assert len(embedding) == spec.dimensions
assert any(value != 0.0 for value in embedding)
@pytest.mark.asyncio
@pytest.mark.parametrize("spec", ALL_SPECS, ids=lambda spec: spec.id)
async def test_live_batch_embed_returns_one_vector_per_input(
spec: LiveEmbeddingSpec,
) -> None:
"""Regression guard for #745.
`gemini-embedding-2*` treats a list of bare strings as parts of one
document and returns a single embedding, which trips the strict zip in
`_process_batch`. Each input must come back with its own distinct vector.
"""
client = make_embedding_client(spec)
embeddings = await client.simple_batch_embed(BATCH_TEXTS)
assert len(embeddings) == len(BATCH_TEXTS)
assert all(len(embedding) == spec.dimensions for embedding in embeddings)
# A collapsed batch would hand the same vector back for every input.
assert len({tuple(embedding) for embedding in embeddings}) == len(BATCH_TEXTS)
@pytest.mark.asyncio
@pytest.mark.parametrize("spec", ALL_SPECS, ids=lambda spec: spec.id)
async def test_live_batch_embed_aligns_with_single_embed(
spec: LiveEmbeddingSpec,
) -> None:
"""Batched vectors must match the one-at-a-time vectors, position for
position. Catches both a collapsed batch and a silently reordered one."""
client = make_embedding_client(spec)
batched = await client.simple_batch_embed(BATCH_TEXTS)
singles = [await client.embed(text) for text in BATCH_TEXTS]
for index, (batched_vector, single_vector) in enumerate(
zip(batched, singles, strict=True)
):
self_similarity = cosine_similarity(batched_vector, single_vector)
assert self_similarity > 0.95, (
f"{spec.id}: batched vector {index} does not match its own "
f"single embedding (cosine={self_similarity:.3f})"
)
for other_index, other_single in enumerate(singles):
if other_index == index:
continue
assert self_similarity > cosine_similarity(batched_vector, other_single), (
f"{spec.id}: batched vector {index} is closer to text "
f"{other_index} than to its own"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("spec", ALL_SPECS, ids=lambda spec: spec.id)
async def test_live_batch_embed_maps_chunks_to_their_ids(
spec: LiveEmbeddingSpec,
) -> None:
"""`batch_embed` splits oversized inputs, so one request carries chunks
belonging to several ids. Every id must get back exactly its own chunks."""
client = make_embedding_client(spec)
long_text = " ".join(
f"paragraph {index} about photosynthesis" for index in range(900)
)
expected_chunks = {
text_id: len(chunks)
for text_id, chunks in client.prepare_chunks(
{"long": long_text, "short": BATCH_TEXTS[1]}
).items()
}
assert expected_chunks["long"] > 1, "test input must exceed the token limit"
result = await client.batch_embed({"long": long_text, "short": BATCH_TEXTS[1]})
assert {text_id: len(vectors) for text_id, vectors in result.items()} == (
expected_chunks
)
assert all(
len(vector) == spec.dimensions
for vectors in result.values()
for vector in vectors
)
@pytest.mark.asyncio
@pytest.mark.parametrize("spec", GEMINI_SPECS, ids=lambda spec: spec.id)
async def test_live_gemini_batch_embed_survives_batch_split(
spec: LiveEmbeddingSpec,
) -> None:
"""Same fix across the batch boundary: with max_batch_size=2 the four
inputs go out as two separate Gemini requests."""
client = make_embedding_client(spec)
client.max_batch_size = 2
embeddings = await client.simple_batch_embed(BATCH_TEXTS)
assert len(embeddings) == len(BATCH_TEXTS)
assert len({tuple(embedding) for embedding in embeddings}) == len(BATCH_TEXTS)

86
tests/llm/test_embedding_client.py
View File

@ -2,9 +2,18 @@ from types import SimpleNamespace
from typing import Any, cast
import pytest
from google.genai import types as genai_types
from src.config import EmbeddingModelConfig, resolve_embedding_model_config
from src.embedding_client import _EmbeddingClient # pyright: ignore[reportPrivateUsage]
from src.embedding_client import (
BatchItem,
_EmbeddingClient, # pyright: ignore[reportPrivateUsage]
)
def gemini_call_texts(contents: Any) -> list[str]:
"""Unwrap a recorded Gemini `contents` argument back to plain texts."""
return [content.parts[0].text for content in contents]
class FakeOpenAIEmbeddingsAPI:
@ -103,7 +112,7 @@ async def test_gemini_embedding_client_uses_output_dimensionality(
self,
*,
model: str,
contents: str | list[str],
contents: Any,
config: dict[str, Any],
) -> SimpleNamespace:
calls.append(
@ -329,11 +338,11 @@ async def test_gemini_simple_batch_embed_respects_configured_max_batch_size(
self,
*,
model: str,
contents: str | list[str],
contents: Any,
config: dict[str, Any],
) -> SimpleNamespace:
calls.append({"model": model, "contents": contents, "config": config})
n = len(contents) if isinstance(contents, list) else 1
n = len(contents)
return SimpleNamespace(
embeddings=[SimpleNamespace(values=[0.2] * 12) for _ in range(n)]
)
@ -359,7 +368,10 @@ async def test_gemini_simple_batch_embed_respects_configured_max_batch_size(
await client.simple_batch_embed(["a", "b", "c"])
assert [call["contents"] for call in calls] == [["a", "b"], ["c"]]
assert [gemini_call_texts(call["contents"]) for call in calls] == [
["a", "b"],
["c"],
]
@pytest.mark.asyncio
@ -374,11 +386,11 @@ async def test_gemini_simple_batch_embed_defaults_to_100_when_unset(
self,
*,
model: str,
contents: str | list[str],
contents: Any,
config: dict[str, Any],
) -> SimpleNamespace:
calls.append({"model": model, "contents": contents, "config": config})
n = len(contents) if isinstance(contents, list) else 1
n = len(contents)
return SimpleNamespace(
embeddings=[SimpleNamespace(values=[0.2] * 12) for _ in range(n)]
)
@ -404,7 +416,7 @@ async def test_gemini_simple_batch_embed_defaults_to_100_when_unset(
await client.simple_batch_embed(["a", "b", "c"])
assert [call["contents"] for call in calls] == [["a", "b", "c"]]
assert [gemini_call_texts(call["contents"]) for call in calls] == [["a", "b", "c"]]
@pytest.mark.asyncio
@ -625,3 +637,61 @@ def test_embedding_model_config_parses_max_batch_size_from_env(
resolved = resolve_embedding_model_config(s.MODEL_CONFIG)
assert resolved.max_batch_size == 10
@pytest.mark.asyncio
async def test_gemini_process_batch_wraps_contents_as_content_part(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Each batch item must be its own Content so gemini-embedding-2* returns
one embedding per item instead of merging them into one document."""
calls: list[dict[str, Any]] = []
class FakeGeminiModels:
async def embed_content(
self,
*,
model: str,
contents: Any,
config: dict[str, Any],
) -> SimpleNamespace:
calls.append({"model": model, "contents": contents, "config": config})
embeddings = [SimpleNamespace(values=[0.3] * 8) for _ in contents]
return SimpleNamespace(embeddings=embeddings)
class FakeGeminiClient:
def __init__(self, *, api_key: str | None, http_options: Any) -> None:
self.api_key: str | None = api_key
self.http_options: Any = http_options
self.aio: Any = SimpleNamespace(models=FakeGeminiModels())
monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient)
client = _EmbeddingClient(
EmbeddingModelConfig(
transport="gemini",
model="gemini-embedding-2",
api_key="gemini-key",
base_url=None,
),
vector_dimensions=8,
max_input_tokens=4096,
max_tokens_per_request=300_000,
send_dimensions=False,
)
batch = [
BatchItem("hello", "id1", 0, 1),
BatchItem("world", "id2", 0, 1),
]
result = await client._process_batch(batch) # pyright: ignore[reportPrivateUsage]
assert result["id1"][0] == [0.3] * 8
assert result["id2"][0] == [0.3] * 8
assert len(calls) == 1
contents = calls[0]["contents"]
assert len(contents) == 2
assert all(isinstance(c, genai_types.Content) for c in contents)
assert contents[0].parts[0].text == "hello"
assert contents[1].parts[0].text == "world"