From 9d4e810918ad857da821f9e2d29f3e8323e5236b Mon Sep 17 00:00:00 2001 From: zhanghui Date: Mon, 12 Oct 2026 18:45:18 +0800 Subject: [PATCH] fix(cascade): embed a file's changed entries in one batched call The four daily-log handlers each called the embedding endpoint once per entry: _embed_entries gathered _build_row over the changed entries and every _build_row awaited embed_or_none. embed_batch, which chunks by [embedding].batch_size, was only ever used by the vector backfill, so batch_size had no effect on the live write path. One turn that writes ~20 atomic facts therefore queued ~20 single-input requests on the process-wide embedding semaphore, and the next search's query embedding waited behind all of them -- about five times the endpoint's latency at the previous max_concurrent of 5. Measured against a stub endpoint of 1 s per request, 23 entries: 24 requests and 5.13 s of query wait before, 4 requests and 1.05 s after. Subclasses now declare the texts they embed in _embed_texts; the base sends every text of a file in one embed_batch call and hands each entry its vectors. _build_row became a pure mapping. The provider now rejects a response whose vector count differs from the input count: callers slice the batch back per input, so a short list would silently shift every later vector onto the wrong text. [embedding].max_concurrent now defaults to 10, and the EmbeddingSettings class default (a dead 50, always overridden by the shipped TOML) is aligned to it. Refs: #486 Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 21 ++ config.example.toml | 2 +- docs/configuration.md | 8 +- src/everos/component/embedding/capability.py | 18 +- .../component/embedding/openai_provider.py | 7 +- src/everos/config/default.toml | 2 +- src/everos/config/settings.py | 2 +- .../memory/cascade/handlers/__init__.py | 5 +- .../cascade/handlers/_daily_log_base.py | 68 ++++- .../memory/cascade/handlers/agent_case.py | 20 +- .../memory/cascade/handlers/atomic_fact.py | 20 +- src/everos/memory/cascade/handlers/episode.py | 32 +- .../memory/cascade/handlers/foresight.py | 20 +- .../test_embedding/test_accessor.py | 15 + .../test_embedding/test_openai_provider.py | 18 +- .../test_writers/test_daily_log_writers.py | 6 +- .../test_cascade/test_handler_episode.py | 11 +- .../test_handlers_daily_log_mapping.py | 28 +- .../test_cascade/test_handlers_embed_batch.py | 276 ++++++++++++++++++ .../test_handlers_embed_or_none.py | 34 ++- .../test_strategy_to_handler_contract.py | 8 +- 21 files changed, 506 insertions(+), 115 deletions(-) create mode 100644 tests/unit/test_memory/test_cascade/test_handlers_embed_batch.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ff9416fa..7bd56aad1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed + +- **`[embedding].max_concurrent` defaults to 10 instead of 5.** The + `EmbeddingSettings` class default (50) is aligned to the same value; the + shipped `default.toml` always overrode it, so the effective default moves + from 5 to 10. (#486) + ### Fixed - **The same owner's memories in two projects no longer overwrite each other @@ -44,6 +51,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 nothing pending, waiting for the worker's in-flight batch, and `everos cascade sync` / `rebuild` use the same loop. +- **A cascade pass embeds a file's changed entries in one batched request, + not one request per entry.** The four daily-log handlers (episode, atomic + fact, foresight, agent case) each called the embedding endpoint once per + entry, so a turn that wrote ~20 atomic facts queued ~20 single-input + requests on the process-wide embedding semaphore, and the next search's + query embedding waited behind all of them: about five times the + endpoint's latency at the previous default `max_concurrent = 5`. The + handlers now + send every text of a file through `embed_batch`, which + `[embedding].batch_size` splits into requests (23 entries become 3 + requests at the default 10), so after a single turn's writes the query + embedding waits for one request; a large backlog still queues ahead of + it (#473). (#486) + - **Changing a knowledge document's category keeps it inside `knowledge/`.** `PATCH /knowledge/documents/{doc_id}` with `category_id` set to `..` moved the document's directory up into the project directory, where the markdown diff --git a/config.example.toml b/config.example.toml index 48cff25fd..b411bbf6a 100644 --- a/config.example.toml +++ b/config.example.toml @@ -34,7 +34,7 @@ base_url = "https://api.deepinfra.com/v1/openai" timeout_seconds = 30.0 max_retries = 3 batch_size = 10 -max_concurrent = 5 +max_concurrent = 10 # ── Rerank ──────────────────────────────────────────── [rerank] diff --git a/docs/configuration.md b/docs/configuration.md index 5cf85a198..6035483ca 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -151,7 +151,13 @@ Zilliz Cloud endpoint; a Milvus Lite filesystem path is rejected. | `timeout_seconds` | float | `30.0` | No | Request timeout. | | `max_retries` | int | `3` | No | Retry count on failure. | | `batch_size` | int | `10` | No | Texts per batch request. | -| `max_concurrent` | int | `5` | No | Max parallel batch requests. | +| `max_concurrent` | int | `10` | No | Max parallel batch requests. | + +`timeout_seconds` bounds one request, and one request carries up to +`batch_size` texts (the cascade embeds a file's changed entries in one +batched call). On a slow or serialized local endpoint (CPU Ollama, a single +GPU) queued batch requests can exceed the timeout: lower `batch_size` and +`max_concurrent` together, or raise `timeout_seconds`. ### `[rerank]` diff --git a/src/everos/component/embedding/capability.py b/src/everos/component/embedding/capability.py index 0826bc924..6b21ba447 100644 --- a/src/everos/component/embedding/capability.py +++ b/src/everos/component/embedding/capability.py @@ -3,8 +3,9 @@ The wrapper encapsulates ``Optional[EmbeddingProvider]`` so consumers never see ``None`` directly. Two consumption modes: -- :meth:`embed_or_none` for soft-degrade write paths (cascade handlers - that write ``vector=NULL`` when embedding is unavailable). +- :meth:`embed_or_none` / :meth:`embed_batch_or_none` for soft-degrade + write paths (cascade handlers that write ``vector=NULL`` when embedding + is unavailable). - :meth:`require` for hard-required paths (search VECTOR/HYBRID/AGENTIC) that raise :class:`ProviderNotConfiguredError` -> HTTP 422 when missing. @@ -17,6 +18,7 @@ from __future__ import annotations +from collections.abc import Sequence from dataclasses import dataclass from everos.core.errors import ProviderNotConfiguredError @@ -45,6 +47,18 @@ async def embed_or_none(self, text: str) -> list[float] | None: return None return list(await self.provider.embed(text)) + async def embed_batch_or_none( + self, texts: Sequence[str] + ) -> list[list[float]] | None: + """Embed ``texts`` in one batched call if available, else ``None``. + + Same soft-degrade contract as :meth:`embed_or_none`; the provider + chunks the request by its configured ``batch_size``. + """ + if self.provider is None: + return None + return [list(v) for v in await self.provider.embed_batch(texts)] + def require(self) -> EmbeddingProvider: """Return the provider, or raise :class:`ProviderNotConfiguredError`. diff --git a/src/everos/component/embedding/openai_provider.py b/src/everos/component/embedding/openai_provider.py index b62ec5d83..886ee85bd 100644 --- a/src/everos/component/embedding/openai_provider.py +++ b/src/everos/component/embedding/openai_provider.py @@ -132,9 +132,12 @@ async def _embed_chunk(self, chunk: list[str]) -> list[list[float]]: ) except openai.OpenAIError as exc: raise _classify(exc) from exc - if not response.data: + # Callers slice the batch back per input, so a short (or empty) + # list would silently shift every later vector onto the wrong text. + if len(response.data) != len(chunk): raise EmbeddingServiceError( - f"Embedding API returned empty data for {len(chunk)} inputs" + f"Embedding API returned {len(response.data)} vectors " + f"for {len(chunk)} inputs" ) # Embeddings report only input (prompt) tokens. usage = getattr(response, "usage", None) diff --git a/src/everos/config/default.toml b/src/everos/config/default.toml index 47004496f..ca6322e5d 100644 --- a/src/everos/config/default.toml +++ b/src/everos/config/default.toml @@ -135,7 +135,7 @@ base_url = "https://api.deepinfra.com/v1/openai" timeout_seconds = 30.0 max_retries = 3 batch_size = 10 -max_concurrent = 5 +max_concurrent = 10 [rerank] # Rerank provider. Override via env: diff --git a/src/everos/config/settings.py b/src/everos/config/settings.py index aecbeb39d..f9be62a72 100644 --- a/src/everos/config/settings.py +++ b/src/everos/config/settings.py @@ -377,7 +377,7 @@ class EmbeddingSettings(BaseModel): timeout_seconds: float = Field(default=30.0, gt=0) max_retries: int = Field(default=3, ge=0) batch_size: int = Field(default=10, ge=1) - max_concurrent: int = Field(default=50, ge=1) + max_concurrent: int = Field(default=10, ge=1) class RerankSettings(BaseModel): diff --git a/src/everos/memory/cascade/handlers/__init__.py b/src/everos/memory/cascade/handlers/__init__.py index b9306d74c..69ac162d9 100644 --- a/src/everos/memory/cascade/handlers/__init__.py +++ b/src/everos/memory/cascade/handlers/__init__.py @@ -2,8 +2,9 @@ Four daily-log handlers (episode / atomic_fact / foresight / agent_case) inherit :class:`BaseDailyLogHandler` for the shared -read / diff / upsert / delete loop; the per-kind subclass only -declares its repo binding and ``_build_row`` mapping. ``agent_skill``, +read / diff / embed / upsert / delete loop; the per-kind subclass only +declares its repo binding, the texts to embed (``_embed_texts``) and +the ``_build_row`` mapping. ``agent_skill``, ``user_profile``, and ``knowledge_topic`` stand alone — they're single-file kinds (no entries, no per-entry diff), so they implement :class:`Handler` directly and own their reconcile loop. diff --git a/src/everos/memory/cascade/handlers/_daily_log_base.py b/src/everos/memory/cascade/handlers/_daily_log_base.py index dc59287c5..fd9e97cf8 100644 --- a/src/everos/memory/cascade/handlers/_daily_log_base.py +++ b/src/everos/memory/cascade/handlers/_daily_log_base.py @@ -26,8 +26,10 @@ import abc import asyncio import dataclasses +from collections.abc import Sequence from typing import Any, ClassVar +from everos.component.embedding import get_embedding_capability from everos.core.observability.logging import get_logger from everos.core.persistence import MarkdownReader, StructuredEntry from everos.infra.persistence.index import all_of, eq, one_of @@ -65,9 +67,11 @@ class BaseDailyLogHandler(Handler): subset of inline + section fields whose changes should trigger re-upsert + re-embed. Each key is ``"section:Name"`` or ``"inline:name"``. + - :meth:`_embed_texts` (override) — the texts one entry needs embedded. - :meth:`_build_row` (override) — turn a :class:`ParsedEntry` plus - common context (owner_id / owner_type / md_path) into a typed - LanceDB row. Tokenisation + embedding live in the subclass. + common context (owner_id / owner_type / md_path) and the vectors + for its texts into a typed LanceDB row. Tokenisation lives in the + subclass; embedding is batched here, one call per file. """ kind: ClassVar[str] = "" @@ -170,24 +174,46 @@ async def _embed_entries( project_id: str, md_path: str, ) -> list[Any]: - """Build LanceDB rows for changed entries (embed concurrently).""" + """Build LanceDB rows for changed entries. + + Every text the changed entries need goes out in ONE ``embed_batch`` + call (the provider chunks it by ``batch_size``), not one ``embed`` + call per entry: the per-entry fan-out queued N single-input requests + on the process-wide embedding semaphore ahead of the interactive + query embedding (issue #486). + """ if not to_build: return [] - return list( - await asyncio.gather( - *( - self._build_row( - owner_id=owner_id, - owner_type=owner_type, - app_id=app_id, - project_id=project_id, - md_path=md_path, - entry=entry, - ) - for entry in to_build + texts_per_entry = [self._embed_texts(entry) for entry in to_build] + flat = [text for texts in texts_per_entry for text in texts] + vectors: Sequence[list[float] | None] | None + vectors = await get_embedding_capability().embed_batch_or_none(flat) + if vectors is None: + logger.debug( + "cascade_handler_embed_skipped", + kind=self.kind, + md_path=md_path, + entries=len(to_build), + reason="embedding_capability_unavailable", + ) + vectors = [None] * len(flat) + rows: list[Any] = [] + cursor = 0 + for entry, texts in zip(to_build, texts_per_entry, strict=True): + entry_vectors = vectors[cursor : cursor + len(texts)] + cursor += len(texts) + rows.append( + await self._build_row( + owner_id=owner_id, + owner_type=owner_type, + app_id=app_id, + project_id=project_id, + md_path=md_path, + entry=entry, + vectors=entry_vectors, ) ) - ) + return rows async def _apply_lance_changes( self, @@ -272,6 +298,11 @@ async def _mark_deprecated( exc_info=True, ) + @abc.abstractmethod + def _embed_texts(self, entry: ParsedEntry) -> tuple[str, ...]: + """Subclass: the texts one entry needs embedded, in the order + :meth:`_build_row` receives their ``vectors`` back.""" + @abc.abstractmethod async def _build_row( self, @@ -282,9 +313,14 @@ async def _build_row( project_id: str = "default", md_path: str, entry: ParsedEntry, + vectors: Sequence[list[float] | None] = (), ) -> Any: """Subclass: build the typed LanceDB row for one parsed entry. + ``vectors`` aligns with :meth:`_embed_texts`; an empty sequence or + ``None`` items mean "no vector" (embedding unavailable) and the row + is written with ``vector=None``. + ``app_id`` / ``project_id`` carry the path-derived scope; the base always supplies them (via :func:`resolve_scope`). They default to ``"default"`` so white-box callers exercising only the field mapping diff --git a/src/everos/memory/cascade/handlers/agent_case.py b/src/everos/memory/cascade/handlers/agent_case.py index 5ba0fae3b..d64795b2e 100644 --- a/src/everos/memory/cascade/handlers/agent_case.py +++ b/src/everos/memory/cascade/handlers/agent_case.py @@ -28,15 +28,13 @@ from __future__ import annotations -from everos.component.embedding import get_embedding_capability -from everos.core.observability.logging import get_logger +from collections.abc import Sequence + from everos.infra.persistence.index import AgentCase, ParentType, agent_case_repo from ._common import require_float, require_iso_timestamp from ._daily_log_base import BaseDailyLogHandler, ParsedEntry -logger = get_logger(__name__) - class AgentCaseHandler(BaseDailyLogHandler): """Cascade handler for ``agents//.cases/agent_case-*.md``.""" @@ -52,6 +50,9 @@ class AgentCaseHandler(BaseDailyLogHandler): """Includes quality_score (semantic score, not audit). Audit inline (owner_id / session_id / timestamp / parent_id) is excluded.""" + def _embed_texts(self, entry: ParsedEntry) -> tuple[str, ...]: + return (entry.structured.sections.get("TaskIntent", "").strip(),) + async def _build_row( self, *, @@ -61,6 +62,7 @@ async def _build_row( project_id: str = "default", md_path: str, entry: ParsedEntry, + vectors: Sequence[list[float] | None] = (), ) -> AgentCase: s = entry.structured task_intent = s.sections.get("TaskIntent", "").strip() @@ -68,14 +70,6 @@ async def _build_row( key_insight = (s.sections.get("KeyInsight") or "").strip() or None intent_tokens = self._deps.tokenizer.tokenize(task_intent) approach_tokens = self._deps.tokenizer.tokenize(approach) - vector = await get_embedding_capability().embed_or_none(task_intent) - if vector is None: - logger.debug( - "cascade_handler_embed_skipped", - kind=self.kind, - entry_id=entry.entry_id, - reason="embedding_capability_unavailable", - ) return AgentCase( id=f"{owner_id}_{entry.entry_id}", entry_id=entry.entry_id, @@ -97,5 +91,5 @@ async def _build_row( key_insight=key_insight, md_path=md_path, content_sha256=entry.content_sha256, - vector=vector, + vector=vectors[0] if vectors else None, ) diff --git a/src/everos/memory/cascade/handlers/atomic_fact.py b/src/everos/memory/cascade/handlers/atomic_fact.py index 6dee3b94a..0c060d521 100644 --- a/src/everos/memory/cascade/handlers/atomic_fact.py +++ b/src/everos/memory/cascade/handlers/atomic_fact.py @@ -21,15 +21,13 @@ from __future__ import annotations -from everos.component.embedding import get_embedding_capability -from everos.core.observability.logging import get_logger +from collections.abc import Sequence + from everos.infra.persistence.index import AtomicFact, ParentType, atomic_fact_repo from ._common import parse_inline_list, require_iso_timestamp from ._daily_log_base import BaseDailyLogHandler, ParsedEntry -logger = get_logger(__name__) - class AtomicFactHandler(BaseDailyLogHandler): """Cascade handler for ``users//.atomic_facts/atomic_fact-*.md``.""" @@ -40,6 +38,9 @@ class AtomicFactHandler(BaseDailyLogHandler): """Only ``Fact`` matters — it's both the embedded text AND the BM25 source. Audit inline is excluded.""" + def _embed_texts(self, entry: ParsedEntry) -> tuple[str, ...]: + return (entry.structured.sections.get("Fact", "").strip(),) + async def _build_row( self, *, @@ -49,18 +50,11 @@ async def _build_row( project_id: str = "default", md_path: str, entry: ParsedEntry, + vectors: Sequence[list[float] | None] = (), ) -> AtomicFact: s = entry.structured text = s.sections.get("Fact", "").strip() tokens = self._deps.tokenizer.tokenize(text) - vector = await get_embedding_capability().embed_or_none(text) - if vector is None: - logger.debug( - "cascade_handler_embed_skipped", - kind=self.kind, - entry_id=entry.entry_id, - reason="embedding_capability_unavailable", - ) return AtomicFact( id=f"{owner_id}_{entry.entry_id}", entry_id=entry.entry_id, @@ -77,5 +71,5 @@ async def _build_row( fact_tokens=" ".join(tokens), md_path=md_path, content_sha256=entry.content_sha256, - vector=vector, + vector=vectors[0] if vectors else None, ) diff --git a/src/everos/memory/cascade/handlers/episode.py b/src/everos/memory/cascade/handlers/episode.py index c8980650b..f77864d0e 100644 --- a/src/everos/memory/cascade/handlers/episode.py +++ b/src/everos/memory/cascade/handlers/episode.py @@ -34,17 +34,13 @@ from __future__ import annotations -import asyncio +from collections.abc import Sequence -from everos.component.embedding import get_embedding_capability -from everos.core.observability.logging import get_logger from everos.infra.persistence.index import Episode, ParentType, episode_repo from ._common import parse_inline_list, require_iso_timestamp from ._daily_log_base import BaseDailyLogHandler, ParsedEntry -logger = get_logger(__name__) - class EpisodeHandler(BaseDailyLogHandler): """Cascade handler for ``users//episodes/episode-*.md``.""" @@ -65,6 +61,12 @@ class EpisodeHandler(BaseDailyLogHandler): call on Subject edits) is accepted under the single-hash design (cascade Q2 discussion).""" + def _embed_texts(self, entry: ParsedEntry) -> tuple[str, ...]: + s = entry.structured + text = s.sections.get("Content", "").strip() + subject_text = s.sections.get("Subject", "").strip() + return (text, subject_text) if subject_text else (text,) + async def _build_row( self, *, @@ -74,28 +76,14 @@ async def _build_row( project_id: str = "default", md_path: str, entry: ParsedEntry, + vectors: Sequence[list[float] | None] = (), ) -> Episode: s = entry.structured text = s.sections.get("Content", "").strip() subject_text = s.sections.get("Subject", "").strip() - # Embed content and subject concurrently; skip subject embed when absent. - capability = get_embedding_capability() - if subject_text: - vector, subject_vector = await asyncio.gather( - capability.embed_or_none(text), - capability.embed_or_none(subject_text), - ) - else: - vector = await capability.embed_or_none(text) - subject_vector = None - if vector is None: - logger.debug( - "cascade_handler_embed_skipped", - kind=self.kind, - entry_id=entry.entry_id, - reason="embedding_capability_unavailable", - ) + vector = vectors[0] if vectors else None + subject_vector = vectors[1] if len(vectors) > 1 else None # BM25 tokenization covers both body and subject keywords. tokenize_source = f"{text} {subject_text}" if subject_text else text diff --git a/src/everos/memory/cascade/handlers/foresight.py b/src/everos/memory/cascade/handlers/foresight.py index 2c618df18..9d1184faf 100644 --- a/src/everos/memory/cascade/handlers/foresight.py +++ b/src/everos/memory/cascade/handlers/foresight.py @@ -28,8 +28,8 @@ from __future__ import annotations -from everos.component.embedding import get_embedding_capability -from everos.core.observability.logging import get_logger +from collections.abc import Sequence + from everos.infra.persistence.index import Foresight, ParentType, foresight_repo from ._common import ( @@ -40,8 +40,6 @@ ) from ._daily_log_base import BaseDailyLogHandler, ParsedEntry -logger = get_logger(__name__) - class ForesightHandler(BaseDailyLogHandler): """Cascade handler for ``users//.foresights/foresight-*.md``.""" @@ -60,6 +58,9 @@ class ForesightHandler(BaseDailyLogHandler): (owner_id / session_id / timestamp / parent_id / sender_ids) is excluded — changes there don't propagate.""" + def _embed_texts(self, entry: ParsedEntry) -> tuple[str, ...]: + return (entry.structured.sections.get("Foresight", "").strip(),) + async def _build_row( self, *, @@ -69,19 +70,12 @@ async def _build_row( project_id: str = "default", md_path: str, entry: ParsedEntry, + vectors: Sequence[list[float] | None] = (), ) -> Foresight: s = entry.structured text = s.sections.get("Foresight", "").strip() evidence = (s.sections.get("Evidence") or "").strip() or None tokens = self._deps.tokenizer.tokenize(text) - vector = await get_embedding_capability().embed_or_none(text) - if vector is None: - logger.debug( - "cascade_handler_embed_skipped", - kind=self.kind, - entry_id=entry.entry_id, - reason="embedding_capability_unavailable", - ) evidence_tokens = ( " ".join(self._deps.tokenizer.tokenize(evidence)) if evidence else None ) @@ -106,5 +100,5 @@ async def _build_row( evidence_tokens=evidence_tokens, md_path=md_path, content_sha256=entry.content_sha256, - vector=vector, + vector=vectors[0] if vectors else None, ) diff --git a/tests/unit/test_component/test_embedding/test_accessor.py b/tests/unit/test_component/test_embedding/test_accessor.py index 3c284f704..9f309a6f9 100644 --- a/tests/unit/test_component/test_embedding/test_accessor.py +++ b/tests/unit/test_component/test_embedding/test_accessor.py @@ -20,6 +20,9 @@ class _MockEmbedder: async def embed(self, text: str) -> list[float]: return [0.1, 0.2, 0.3] + async def embed_batch(self, texts) -> list[list[float]]: # type: ignore[no-untyped-def] + return [[float(len(t)), 0.2, 0.3] for t in texts] + def _reset_capability(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(_accessor_mod, "_capability", None, raising=False) @@ -64,6 +67,18 @@ async def test_embed_or_none_returns_vector_when_available() -> None: assert result == [0.1, 0.2, 0.3] +async def test_embed_batch_or_none_returns_none_when_unavailable() -> None: + cap = EmbeddingCapability(provider=None) + assert await cap.embed_batch_or_none(["hello", "world"]) is None + + +async def test_embed_batch_or_none_returns_vectors_in_input_order() -> None: + """Callers slice the result back per input, so order is the contract.""" + cap = EmbeddingCapability(provider=_MockEmbedder()) + result = await cap.embed_batch_or_none(["ab", "cdef"]) + assert result == [[2.0, 0.2, 0.3], [4.0, 0.2, 0.3]] + + def test_require_raises_when_unavailable() -> None: cap = EmbeddingCapability(provider=None) with pytest.raises(ProviderNotConfiguredError) as excinfo: diff --git a/tests/unit/test_component/test_embedding/test_openai_provider.py b/tests/unit/test_component/test_embedding/test_openai_provider.py index abc6c4232..054ce95a5 100644 --- a/tests/unit/test_component/test_embedding/test_openai_provider.py +++ b/tests/unit/test_component/test_embedding/test_openai_provider.py @@ -33,5 +33,21 @@ async def test_empty_response_data_raises_embedding_error() -> None: provider._client = mock.AsyncMock() provider._client.embeddings.create = mock.AsyncMock(return_value=empty_response) - with pytest.raises(EmbeddingServiceError, match="empty data"): + with pytest.raises(EmbeddingServiceError, match="0 vectors for 1 inputs"): await provider.embed("hello") + + +async def test_short_response_data_raises_embedding_error() -> None: + """One vector back for two inputs must raise: a silent short list would + shift every later vector onto the wrong text (cascade slices the batch + back per entry).""" + provider = _make_provider() + short_response = mock.MagicMock() + item = mock.MagicMock() + item.embedding = [0.1, 0.2, 0.3, 0.4] + short_response.data = [item] + provider._client = mock.AsyncMock() + provider._client.embeddings.create = mock.AsyncMock(return_value=short_response) + + with pytest.raises(EmbeddingServiceError, match="1 vectors for 2 inputs"): + await provider.embed_batch(["hello", "world"]) diff --git a/tests/unit/test_infra/test_markdown/test_writers/test_daily_log_writers.py b/tests/unit/test_infra/test_markdown/test_writers/test_daily_log_writers.py index 20372fb37..ffb57f243 100644 --- a/tests/unit/test_infra/test_markdown/test_writers/test_daily_log_writers.py +++ b/tests/unit/test_infra/test_markdown/test_writers/test_daily_log_writers.py @@ -228,7 +228,7 @@ async def embed_batch(self, ts): # type: ignore[no-untyped-def] # AtomicFactHandler fetches the embedder lazily via # ``get_embedding_capability()`` — patch the process-wide singleton - # so ``_build_row`` returns a real (stub) vector. + # so the batch embed path returns a real (stub) vector. monkeypatch.setattr(acc, "_capability", EmbeddingCapability(provider=_E())) today = _dt.date(2026, 5, 15) @@ -253,9 +253,7 @@ async def embed_batch(self, ts): # type: ignore[no-untyped-def] handler = AtomicFactHandler(HandlerDeps(memory_root=memory_root, tokenizer=_T())) structured = entry.as_structured() pe = ParsedEntry(entry.id, structured, handler._content_sha256(structured)) - row = await handler._build_row( - owner_id="u1", owner_type="user", md_path=rel, entry=pe - ) + [row] = await handler._embed_entries([pe], "u1", "user", "default", "default", rel) assert row.id == f"u1_{eid.format()}" assert row.fact == "Alice prefers Italian." assert row.parent_id == "mc_1" diff --git a/tests/unit/test_memory/test_cascade/test_handler_episode.py b/tests/unit/test_memory/test_cascade/test_handler_episode.py index d423aa64d..7e3671b94 100644 --- a/tests/unit/test_memory/test_cascade/test_handler_episode.py +++ b/tests/unit/test_memory/test_cascade/test_handler_episode.py @@ -35,7 +35,8 @@ def tokenize_batch(self, texts): # type: ignore[no-untyped-def] class _StubEmbedder(EmbeddingProvider): - """Returns a fixed 1024-dim vector; records call count.""" + """Returns a fixed 1024-dim vector; counts the texts embedded (the + batch path reaches ``embed`` once per text via ``embed_batch``).""" dim = 1024 @@ -162,7 +163,7 @@ async def test_added_entry_upserts_typed_row( assert outcome.upserted == 1 assert outcome.deleted == 0 assert outcome.skipped == 0 - # Subject "Test" is present → two embed calls (content + subject). + # Subject "Test" is present → two texts embedded (content + subject). assert embedder.calls == 2 assert len(fake_repo.upserts) == 1 row = fake_repo.upserts[0][0] @@ -230,7 +231,7 @@ async def test_modified_entry_reembeds( outcome = await handler.handle_added_or_modified(rel) assert outcome.upserted == 1 assert outcome.skipped == 0 - # Subject "Test" is present → two embed calls (content + subject). + # Subject "Test" is present → two texts embedded (content + subject). assert embedder.calls == 2 @@ -239,7 +240,7 @@ async def test_no_subject_skips_subject_embed( fake_repo: _FakeEpisodeRepo, stub_embedder: _StubEmbedder, ) -> None: - """When Subject is absent, subject_vector is None; only one embed call is made.""" + """When Subject is absent, subject_vector is None; only content is embedded.""" today = _dt.date(2026, 5, 14) writer = EpisodeWriter(memory_root) await writer.append_entry( @@ -262,7 +263,7 @@ async def test_no_subject_skips_subject_embed( outcome = await handler.handle_added_or_modified(rel) assert outcome.upserted == 1 - # No Subject → single embed call for content only. + # No Subject → only the content text is embedded. assert embedder.calls == 1 row = fake_repo.upserts[0][0] assert row.subject_vector is None diff --git a/tests/unit/test_memory/test_cascade/test_handlers_daily_log_mapping.py b/tests/unit/test_memory/test_cascade/test_handlers_daily_log_mapping.py index b3fcb4a46..0afce454e 100644 --- a/tests/unit/test_memory/test_cascade/test_handlers_daily_log_mapping.py +++ b/tests/unit/test_memory/test_cascade/test_handlers_daily_log_mapping.py @@ -90,12 +90,26 @@ def _entry( ) +async def _embed_one(handler, **kw): # type: ignore[no-untyped-def] + """Run one entry through the batch embed path (the production route).""" + [row] = await handler._embed_entries( + [kw["entry"]], + kw["owner_id"], + kw["owner_type"], + kw.get("app_id", "default"), + kw.get("project_id", "default"), + kw["md_path"], + ) + return row + + # ── AtomicFact ─────────────────────────────────────────────────────────── async def test_atomic_fact_build_row_maps_inline_and_section(tmp_path) -> None: # type: ignore[no-untyped-def] handler = AtomicFactHandler(_deps(tmp_path)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="u1", owner_type="user", md_path="users/u1/.atomic_facts/atomic_fact-2026-05-14.md", @@ -141,7 +155,8 @@ async def test_atomic_fact_missing_timestamp_raises(tmp_path) -> None: # type: async def test_foresight_build_row_with_evidence(tmp_path) -> None: # type: ignore[no-untyped-def] handler = ForesightHandler(_deps(tmp_path)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="u1", owner_type="user", md_path="users/u1/.foresights/foresight-2026-05-14.md", @@ -173,7 +188,8 @@ async def test_foresight_build_row_with_evidence(tmp_path) -> None: # type: ign async def test_foresight_optional_evidence_left_none(tmp_path) -> None: # type: ignore[no-untyped-def] handler = ForesightHandler(_deps(tmp_path)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="u1", owner_type="user", md_path="x.md", @@ -200,7 +216,8 @@ async def test_foresight_optional_evidence_left_none(tmp_path) -> None: # type: async def test_agent_case_build_row_maps_intent_approach_insight(tmp_path) -> None: # type: ignore[no-untyped-def] handler = AgentCaseHandler(_deps(tmp_path)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="a1", owner_type="agent", md_path="agents/a1/.cases/agent_case-2026-05-14.md", @@ -230,7 +247,8 @@ async def test_agent_case_build_row_maps_intent_approach_insight(tmp_path) -> No async def test_agent_case_optional_insight_left_none(tmp_path) -> None: # type: ignore[no-untyped-def] handler = AgentCaseHandler(_deps(tmp_path)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="a1", owner_type="agent", md_path="x.md", diff --git a/tests/unit/test_memory/test_cascade/test_handlers_embed_batch.py b/tests/unit/test_memory/test_cascade/test_handlers_embed_batch.py new file mode 100644 index 000000000..acea537a9 --- /dev/null +++ b/tests/unit/test_memory/test_cascade/test_handlers_embed_batch.py @@ -0,0 +1,276 @@ +"""Contract (issue #486): one cascade pass over a daily-log file embeds every +changed entry's texts in a single ``embed_batch`` call, never one ``embed`` +call per entry. + +The base handler collects the texts each kind declares via ``_embed_texts``, +embeds them all at once and hands each entry's vectors back to ``_build_row``. +The per-entry fan-out this replaces put N single-input requests on the +process-wide embedding semaphore and starved the interactive query embedding. +""" + +from __future__ import annotations + +import datetime as _dt +from typing import Any + +import pytest + +from everos.component.embedding import EmbeddingCapability, EmbeddingProvider +from everos.component.tokenizer import Tokenizer +from everos.core.persistence import MemoryRoot, StructuredEntry +from everos.infra.persistence.index import Comparison, Predicate +from everos.infra.persistence.markdown import AtomicFactWriter, EpisodeWriter +from everos.memory.cascade.handlers import ( + AgentCaseHandler, + AtomicFactHandler, + ForesightHandler, + HandlerDeps, +) +from everos.memory.cascade.handlers._daily_log_base import ParsedEntry +from everos.memory.cascade.handlers.episode import EpisodeHandler + +_DIM = 1024 + + +def _vec(text: str) -> list[float]: + """A vector that identifies the text it was made from (by its length).""" + return [float(len(text))] + [0.0] * (_DIM - 1) + + +class _RecordingEmbedder(EmbeddingProvider): + dim = _DIM + + def __init__(self) -> None: + self.single: list[str] = [] + self.batches: list[list[str]] = [] + + async def embed(self, text: str) -> list[float]: + self.single.append(text) + return _vec(text) + + async def embed_batch(self, texts): # type: ignore[no-untyped-def] + texts = list(texts) + self.batches.append(texts) + return [_vec(t) for t in texts] + + +class _StubTokenizer(Tokenizer): + def tokenize(self, text: str) -> list[str]: + return [tok for tok in text.split() if tok] + + def tokenize_batch(self, texts): # type: ignore[no-untyped-def] + return [self.tokenize(t) for t in texts] + + +class _FakeRepo: + def __init__(self) -> None: + self.rows: list[Any] = [] + + async def find_where(self, where: Predicate, *, limit: int = 100) -> list[Any]: + if isinstance(where, Comparison) and where.field == "md_path": + return [r for r in self.rows if r.md_path == where.value] + return [] + + async def upsert(self, rows: list[Any]) -> None: + self.rows.extend(rows) + + async def delete(self, predicate: Predicate) -> None: + pass + + async def delete_by_md_path(self, md_path: str) -> int: + return 0 + + +def _set_capability( + monkeypatch: pytest.MonkeyPatch, provider: EmbeddingProvider | None +) -> None: + import everos.component.embedding.accessor as acc + + monkeypatch.setattr(acc, "_capability", EmbeddingCapability(provider=provider)) + + +def _deps(tmp_path) -> HandlerDeps: # type: ignore[no-untyped-def] + mr = MemoryRoot(tmp_path) + mr.ensure() + return HandlerDeps(memory_root=mr, tokenizer=_StubTokenizer()) + + +_INLINE = { + "owner_id": "u1", + "session_id": "s1", + "timestamp": "2026-05-14T10:00:00+00:00", + "parent_id": "mc_1", + "quality_score": "0.5", +} + + +def _entry(sections: dict[str, str]) -> ParsedEntry: + return ParsedEntry( + entry_id="x_20260514_0001", + structured=StructuredEntry( + id="x_20260514_0001", + body="", + start=0, + end=0, + header=None, + inline=_INLINE, + sections=sections, + ), + content_sha256="f" * 64, + ) + + +@pytest.mark.parametrize( + ("handler_cls", "sections", "expected_texts"), + [ + (AtomicFactHandler, {"Fact": " the fact "}, ("the fact",)), + ( + ForesightHandler, + {"Foresight": "will book lunch", "Evidence": "said so"}, + ("will book lunch",), + ), + ( + AgentCaseHandler, + {"TaskIntent": "scan contract", "Approach": "read pages"}, + ("scan contract",), + ), + ( + EpisodeHandler, + {"Subject": "Sub", "Summary": "s", "Content": "hello world"}, + ("hello world", "Sub"), + ), + (EpisodeHandler, {"Summary": "s", "Content": "hello world"}, ("hello world",)), + ], +) +async def test_embed_texts_and_vectors_round_trip( # type: ignore[no-untyped-def] + tmp_path, handler_cls, sections, expected_texts +) -> None: + """Each kind declares exactly the texts it embeds, and the vectors the + base hands back land on the matching row columns.""" + handler = handler_cls(_deps(tmp_path)) + entry = _entry(sections) + + assert handler._embed_texts(entry) == expected_texts + + vectors = [_vec(t) for t in expected_texts] + row = await handler._build_row( + owner_id="u1", + owner_type="user", + md_path="x.md", + entry=entry, + vectors=vectors, + ) + assert list(row.vector) == vectors[0] + if handler_cls is EpisodeHandler: + subject_vector = vectors[1] if len(vectors) > 1 else None + assert ( + None if row.subject_vector is None else list(row.subject_vector) + ) == subject_vector + + +async def _write_episodes( + memory_root: MemoryRoot, entries: list[tuple[str, str | None]] +) -> str: + writer = EpisodeWriter(memory_root) + for content, subject in entries: + sections = {"Summary": "s", "Content": content} + if subject is not None: + sections["Subject"] = subject + await writer.append_entry( + "u1", + inline={ + "owner_id": "u1", + "session_id": "s1", + "timestamp": "2026-05-14T10:00:00+00:00", + "parent_type": "memcell", + "parent_id": "mc_1", + "sender_ids": ["u1"], + }, + sections=sections, + date=_dt.date(2026, 5, 14), + ) + return "default_app/default_project/users/u1/episodes/episode-2026-05-14.md" + + +async def test_one_file_pass_issues_one_embed_batch( # type: ignore[no-untyped-def] + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + embedder = _RecordingEmbedder() + _set_capability(monkeypatch, embedder) + repo = _FakeRepo() + monkeypatch.setattr(EpisodeHandler, "index_repo", repo) + deps = _deps(tmp_path) + rel = await _write_episodes( + deps.memory_root, [("aa", "S1"), ("bbbb", None), ("cccccc", "S3xyz")] + ) + + outcome = await EpisodeHandler(deps).handle_added_or_modified(rel) + + assert outcome.upserted == 3 + assert embedder.single == [] + assert embedder.batches == [["aa", "S1", "bbbb", "cccccc", "S3xyz"]] + by_text = {row.episode: row for row in repo.rows} + assert by_text["aa"].vector[0] == 2.0 + assert by_text["aa"].subject_vector[0] == 2.0 + assert by_text["bbbb"].vector[0] == 4.0 + assert by_text["bbbb"].subject_vector is None + assert by_text["cccccc"].vector[0] == 6.0 + assert by_text["cccccc"].subject_vector[0] == 5.0 + + +async def test_missing_provider_writes_null_vectors_through_batch_path( # type: ignore[no-untyped-def] + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + _set_capability(monkeypatch, None) + repo = _FakeRepo() + monkeypatch.setattr(EpisodeHandler, "index_repo", repo) + deps = _deps(tmp_path) + rel = await _write_episodes(deps.memory_root, [("aa", "S1"), ("bbbb", None)]) + + outcome = await EpisodeHandler(deps).handle_added_or_modified(rel) + + assert outcome.upserted == 2 + assert all(r.vector is None and r.subject_vector is None for r in repo.rows) + assert {r.episode_tokens for r in repo.rows} == {"aa S1", "bbbb"} + + +async def test_atomic_facts_of_one_file_go_out_in_one_batch( # type: ignore[no-untyped-def] + tmp_path, monkeypatch: pytest.MonkeyPatch +) -> None: + """The issue's shape: ~20 atomic facts per turn, one request chunked by + the provider instead of twenty. Second witness for the batching contract + next to the episode test above.""" + embedder = _RecordingEmbedder() + _set_capability(monkeypatch, embedder) + repo = _FakeRepo() + monkeypatch.setattr(AtomicFactHandler, "index_repo", repo) + deps = _deps(tmp_path) + writer = AtomicFactWriter(deps.memory_root) + facts = [f"fact number {i}" for i in range(20)] + for fact in facts: + await writer.append_entry( + "u1", + inline={ + "owner_id": "u1", + "session_id": "s1", + "timestamp": "2026-05-14T10:00:00+00:00", + "parent_id": "mc_1", + "sender_ids": ["u1"], + }, + sections={"Fact": fact}, + date=_dt.date(2026, 5, 14), + ) + path = ( + deps.memory_root.users_dir() + / "u1" + / ".atomic_facts" + / "atomic_fact-2026-05-14.md" + ) + rel = path.relative_to(deps.memory_root.root).as_posix() + + outcome = await AtomicFactHandler(deps).handle_added_or_modified(rel) + + assert outcome.upserted == 20 + assert embedder.single == [] + assert embedder.batches == [facts] + assert [row.vector[0] for row in repo.rows] == [float(len(f)) for f in facts] diff --git a/tests/unit/test_memory/test_cascade/test_handlers_embed_or_none.py b/tests/unit/test_memory/test_cascade/test_handlers_embed_or_none.py index 299678be3..78afe14fc 100644 --- a/tests/unit/test_memory/test_cascade/test_handlers_embed_or_none.py +++ b/tests/unit/test_memory/test_cascade/test_handlers_embed_or_none.py @@ -1,11 +1,12 @@ """Contract: every memory handler writes ``vector=None`` when embedding is unavailable. -Handlers fetch the embedder lazily via +The base handler fetches the embedder lazily via ``everos.component.embedding.get_embedding_capability()`` rather than -receiving it through :class:`HandlerDeps`. When the process-wide -capability wraps ``provider=None`` (embedding not configured), each -handler must still upsert its row — BM25 tokenization and scalar +receiving it through :class:`HandlerDeps`, and embeds every changed entry +of a file in one batch. When the process-wide capability wraps +``provider=None`` (embedding not configured), each kind must still build +its row — BM25 tokenization and scalar columns are populated as usual — but with ``vector`` (and, for Episode, ``subject_vector`` too) left as ``None`` rather than raising. """ @@ -83,11 +84,25 @@ def _entry( ) +async def _embed_one(handler, **kw): # type: ignore[no-untyped-def] + """Run one entry through the batch embed path (the production route).""" + [row] = await handler._embed_entries( + [kw["entry"]], + kw["owner_id"], + kw["owner_type"], + kw.get("app_id", "default"), + kw.get("project_id", "default"), + kw["md_path"], + ) + return row + + async def test_atomic_fact_handler_writes_null_vector_when_embed_unavailable( memory_root: MemoryRoot, ) -> None: handler = AtomicFactHandler(_deps(memory_root)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="u1", owner_type="user", md_path="x.md", @@ -111,7 +126,8 @@ async def test_foresight_handler_writes_null_vector_when_embed_unavailable( memory_root: MemoryRoot, ) -> None: handler = ForesightHandler(_deps(memory_root)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="u1", owner_type="user", md_path="x.md", @@ -134,7 +150,8 @@ async def test_agent_case_handler_writes_null_vector_when_embed_unavailable( memory_root: MemoryRoot, ) -> None: handler = AgentCaseHandler(_deps(memory_root)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="a1", owner_type="agent", md_path="x.md", @@ -159,7 +176,8 @@ async def test_episode_handler_writes_null_vectors_when_embed_unavailable( ) -> None: """Episode has two vector columns — both must degrade to ``None``.""" handler = EpisodeHandler(_deps(memory_root)) - row = await handler._build_row( + row = await _embed_one( + handler, owner_id="u1", owner_type="user", md_path="x.md", diff --git a/tests/unit/test_memory/test_strategies/test_strategy_to_handler_contract.py b/tests/unit/test_memory/test_strategies/test_strategy_to_handler_contract.py index db323d359..6ba5b3ca1 100644 --- a/tests/unit/test_memory/test_strategies/test_strategy_to_handler_contract.py +++ b/tests/unit/test_memory/test_strategies/test_strategy_to_handler_contract.py @@ -125,12 +125,10 @@ async def _build_row_from_md( structured=structured, content_sha256=handler._content_sha256(structured), ) - return await handler._build_row( - owner_id=owner_id, - owner_type=owner_type, - md_path=rel, - entry=pe, + [row] = await handler._embed_entries( + [pe], owner_id, owner_type, "default", "default", rel ) + return row async def test_atomic_fact_strategy_md_feeds_handler_with_content(