feat: batch embedding pipeline (#292)

Author: alexeyqu

context_use/cli/base.py

@@ -30,6 +30,7 @@ def _batch_detail_from_state(state: State | None) -> str:
         MemoryEmbedCompleteState,
         MemoryGenerateCompleteState,
     )
+    from context_use.thread_embedding.states import ThreadEmbedCompleteState
 
     if isinstance(state, FailedState):
         message = state.error_message.strip()
@@ -50,6 +51,9 @@ def _batch_detail_from_state(state: State | None) -> str:
     if isinstance(state, DescGenerateCompleteState):
         return f"{state.descriptions_count} descriptions generated"
 
+    if isinstance(state, ThreadEmbedCompleteState):
+        return f"{state.embedded_count} threads embedded"
+
     return ""
 
 

context_use/cli/commands/__init__.py

@@ -4,6 +4,7 @@
 from context_use.cli.commands.agent import AgentGroup
 from context_use.cli.commands.config import ConfigGroup
 from context_use.cli.commands.describe import DescribeCommand
+from context_use.cli.commands.embed import EmbedCommand
 from context_use.cli.commands.ingest import IngestCommand
 from context_use.cli.commands.memories import MemoriesGroup
 from context_use.cli.commands.pipeline import PipelineCommand
@@ -15,6 +16,7 @@
     PipelineCommand,
     IngestCommand,
     DescribeCommand,
+    EmbedCommand,
     ResetCommand,
 ]
 

context_use/cli/commands/embed.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import argparse
+from datetime import UTC, datetime, timedelta
+from typing import TYPE_CHECKING
+
+from context_use.cli import output as out
+from context_use.cli.base import ApiCommand, create_batch_reporter, run_batches
+from context_use.cli.config import Config
+
+if TYPE_CHECKING:
+    from context_use import ContextUse
+
+
+class EmbedCommand(ApiCommand):
+    name = "embed"
+    help = "Embed thread content for semantic search"
+    description = (
+        "Generate vector embeddings for all unprocessed threads. "
+        "Asset threads require descriptions first (run 'describe' beforehand). "
+        "Use --last-days or --since to limit the date range."
+    )
+    llm_mode = "batch"
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        parser.add_argument(
+            "--last-days",
+            type=int,
+            default=None,
+            help="Only process threads from the last N days",
+        )
+        parser.add_argument(
+            "--since",
+            type=str,
+            default=None,
+            help="Only process threads after this date (YYYY-MM-DD)",
+        )
+
+    async def run(
+        self,
+        cfg: Config,
+        ctx: ContextUse,
+        args: argparse.Namespace,
+    ) -> None:
+        since = self._resolve_since(args)
+
+        out.header("Embedding threads")
+        out.info("Embeds all threads that have not been embedded yet.")
+        if since:
+            out.kv("Since", since.strftime("%Y-%m-%d"))
+        print()
+
+        batches = await ctx.create_thread_embedding_batches(since=since)
+
+        if not batches:
+            out.info("No threads to embed.")
+            return
+
+        await run_batches(
+            ctx,
+            batches,
+            reporter_factory=create_batch_reporter,
+        )
+
+        out.success("Thread embeddings generated")
+        out.kv("Batches", len(batches))
+        print()
+
+    def _resolve_since(self, args: argparse.Namespace) -> datetime | None:
+        if args.since:
+            return datetime.fromisoformat(args.since).replace(tzinfo=UTC)
+        if args.last_days is not None:
+            return datetime.now(UTC) - timedelta(days=args.last_days)
+        return None

context_use/core.py

@@ -235,6 +235,37 @@ async def create_asset_description_batches(
         groups = [ThreadGroup(threads=[t], group_id=t.id) for t in asset_threads]
         return await AssetDescriptionBatchFactory.create_batches(groups, self._store)
 
+    # ── Thread embedding batches ─────────────────────────────────────
+
+    async def create_thread_embedding_batches(
+        self,
+        *,
+        task_id: str | None = None,
+        since: datetime | None = None,
+        before: datetime | None = None,
+    ) -> list[Batch]:
+        """Create batches for thread embedding.
+
+        Every thread with embeddable content is included.  Asset threads
+        without a description yet are silently skipped (their
+        ``get_embeddable_content()`` returns ``None``).
+        """
+        from context_use.thread_embedding.factory import ThreadEmbeddingBatchFactory
+
+        threads = await self._store.get_unprocessed_threads(
+            batch_category=BatchCategory.thread_embedding.value,
+            task_id=task_id,
+            since=since,
+            before=before,
+        )
+
+        embeddable = [t for t in threads if t.get_embeddable_content() is not None]
+        if not embeddable:
+            return []
+
+        groups = [ThreadGroup(threads=[t], group_id=t.id) for t in embeddable]
+        return await ThreadEmbeddingBatchFactory.create_batches(groups, self._store)
+
     # ── Memory batches ────────────────────────────────────────────────
 
     async def create_memory_batches(
@@ -452,3 +483,4 @@ def _ensure_managers_registered() -> None:
     """Import manager modules to trigger their @register_batch_manager decorators."""
     import context_use.asset_description.manager  # noqa: F401
     import context_use.memories.manager  # noqa: F401
+    import context_use.thread_embedding.manager  # noqa: F401

context_use/models/batch.py

@@ -20,6 +20,7 @@ class BatchCategory(enum.StrEnum):
 
     memories = "memories"
     asset_description = "asset_description"
+    thread_embedding = "thread_embedding"
 
 
 @dataclass

context_use/models/thread.py

@@ -82,6 +82,18 @@ def get_raw_content(self) -> str:
         """Return semantic content from the payload, ignoring any enriched content."""
         return self._parsed_payload.get_content() or ""
 
+    def get_embeddable_content(self) -> str | None:
+        """Return text suitable for embedding, or ``None`` to skip.
+
+        Asset threads use the enriched ``content`` (set by the describe
+        pipeline).  If no description exists yet the thread is not ready
+        for embedding.  Non-asset threads use the raw payload content.
+        """
+        if self.is_asset:
+            return self.content
+        raw = self.get_raw_content()
+        return raw or None
+
     def get_participant_label(self) -> str:
         return self._parsed_payload.get_participant_label()
 

context_use/store/base.py

@@ -256,6 +256,15 @@ async def search_memories(
         """Search memories by semantic similarity, optionally filtered by date range."""
         ...
 
+    # ── Thread Embeddings ─────────────────────────────────────────────
+
+    @abstractmethod
+    async def upsert_thread_embedding(
+        self, thread_id: str, embedding: list[float]
+    ) -> None:
+        """Insert or replace the embedding vector for a thread."""
+        ...
+
     # ── Memory Facets ────────────────────────────────────────────────
 
     @abstractmethod

context_use/store/sqlite/schema.py

@@ -381,6 +381,27 @@ def from_row(row: Row) -> MemoryFacet:
         )
 
 
+class VecThreadRow:
+    table = "vec_threads"
+
+    @classmethod
+    def ddl(cls, embedding_dimensions: int) -> str:
+        return (
+            "CREATE VIRTUAL TABLE IF NOT EXISTS vec_threads "
+            f"USING vec0(\n"
+            f"    thread_id TEXT PRIMARY KEY,\n"
+            f"    embedding float[{embedding_dimensions}] "
+            f"distance_metric=cosine\n"
+            f")"
+        )
+
+    @staticmethod
+    def serialize(embedding: list[float]) -> bytes:
+        from sqlite_vec import serialize_float32
+
+        return serialize_float32(embedding)
+
+
 class VecFacetRow:
     table = "vec_facets"
 
@@ -420,5 +441,6 @@ def all_ddl_statements(embedding_dimensions: int) -> list[str]:
         stmts.append(model.ddl())
         stmts.extend(model.indices())
     stmts.append(VecMemoryRow.ddl(embedding_dimensions))
+    stmts.append(VecThreadRow.ddl(embedding_dimensions))
     stmts.append(VecFacetRow.ddl(embedding_dimensions))
     return stmts

context_use/store/sqlite/store.py

@@ -33,6 +33,7 @@
     ThreadRow,
     VecFacetRow,
     VecMemoryRow,
+    VecThreadRow,
     all_ddl_statements,
     now_utc_iso,
     parse_dt,
@@ -650,6 +651,20 @@ async def search_memories(
             top_k=top_k,
         )
 
+    async def upsert_thread_embedding(
+        self, thread_id: str, embedding: list[float]
+    ) -> None:
+        db = await self._conn()
+        await db.execute(
+            "DELETE FROM vec_threads WHERE thread_id = ?",
+            (thread_id,),
+        )
+        await db.execute(
+            "INSERT INTO vec_threads (thread_id, embedding) VALUES (?, ?)",
+            (thread_id, VecThreadRow.serialize(embedding)),
+        )
+        await self._commit_unless_atomic()
+
     async def create_memory_facet(self, facet: MemoryFacet) -> MemoryFacet:
         db = await self._conn()
         await db.execute(