diff --git a/src/embedding_client.py b/src/embedding_client.py index dc54065f..48ce6924 100644 --- a/src/embedding_client.py +++ b/src/embedding_client.py @@ -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: diff --git a/src/routers/conclusions.py b/src/routers/conclusions.py index eeed2723..89d6207f 100644 --- a/src/routers/conclusions.py +++ b/src/routers/conclusions.py @@ -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( diff --git a/tests/live_llm/README.md b/tests/live_llm/README.md index cef10802..bc638114 100644 --- a/tests/live_llm/README.md +++ b/tests/live_llm/README.md @@ -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 diff --git a/tests/live_llm/conftest.py b/tests/live_llm/conftest.py index d9646383..e016217a 100644 --- a/tests/live_llm/conftest.py +++ b/tests/live_llm/conftest.py @@ -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, diff --git a/tests/live_llm/embedding_matrix.py b/tests/live_llm/embedding_matrix.py new file mode 100644 index 00000000..96d6b552 --- /dev/null +++ b/tests/live_llm/embedding_matrix.py @@ -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 diff --git a/tests/live_llm/test_live_embeddings.py b/tests/live_llm/test_live_embeddings.py new file mode 100644 index 00000000..64814bc6 --- /dev/null +++ b/tests/live_llm/test_live_embeddings.py @@ -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) diff --git a/tests/llm/test_embedding_client.py b/tests/llm/test_embedding_client.py index e8ad30ab..1d912047 100644 --- a/tests/llm/test_embedding_client.py +++ b/tests/llm/test_embedding_client.py @@ -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"