feat(phase6): Embeddings & Semantic Search implementation

- Added EmbeddingService with ChromaDB and SentenceTransformers
- Updated SearchAgent to support semantic deduplication and search
- Updated MagenticOrchestrator to use embeddings if available
- Added 100% async-safe implementation with run_in_executor
- Added comprehensive unit tests

Dockerfile

@@ -19,7 +19,14 @@ COPY src/ src/
 COPY README.md .
 
 # Install dependencies
-RUN uv sync --frozen --no-dev
+RUN uv sync --frozen --no-dev --all-extras
+
+# Set cache directory for HuggingFace models
+ENV HF_HOME=/app/.cache
+ENV TRANSFORMERS_CACHE=/app/.cache
+
+# Pre-download the embedding model during build to speed up cold starts
+RUN uv run python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('all-MiniLM-L6-v2')"
 
 # Create non-root user
 RUN useradd --create-home --shell /bin/bash appuser

pyproject.toml

@@ -49,6 +49,10 @@ dev = [
 magentic = [
     "agent-framework-core",
 ]
+embeddings = [
+    "chromadb>=0.4.0",
+    "sentence-transformers>=2.2.0",
+]
 
 [build-system]
 requires = ["hatchling"]

src/agents/search_agent.py

@@ -1,5 +1,5 @@
 from collections.abc import AsyncIterable
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from agent_framework import (
     AgentRunResponse,
@@ -11,7 +11,10 @@ from agent_framework import (
 )
 
 from src.orchestrator import SearchHandlerProtocol
-from src.utils.models import Evidence, SearchResult
+from src.utils.models import Citation, Evidence, SearchResult
+
+if TYPE_CHECKING:
+    from src.services.embeddings import EmbeddingService
 
 
 class SearchAgent(BaseAgent):  # type: ignore[misc]
@@ -21,6 +24,7 @@ class SearchAgent(BaseAgent):  # type: ignore[misc]
         self,
         search_handler: SearchHandlerProtocol,
         evidence_store: dict[str, list[Evidence]],
+        embedding_service: "EmbeddingService | None" = None,
     ) -> None:
         super().__init__(
             name="SearchAgent",
@@ -28,6 +32,7 @@ class SearchAgent(BaseAgent):  # type: ignore[misc]
         )
         self._handler = search_handler
         self._evidence_store = evidence_store
+        self._embeddings = embedding_service
 
     async def run(
         self,
@@ -62,30 +67,92 @@ class SearchAgent(BaseAgent):  # type: ignore[misc]
         result: SearchResult = await self._handler.execute(query, max_results_per_tool=10)
 
         # Update shared evidence store
-        # We append new evidence, deduplicating by URL is handled in Orchestrator usually,
-        # but here we should probably add to the list.
-        # For simplicity in this MVP phase, we just extend the list.
-        # Ideally, we should dedupe.
-        existing_urls = {e.citation.url for e in self._evidence_store["current"]}
-        new_unique = [e for e in result.evidence if e.citation.url not in existing_urls]
-        self._evidence_store["current"].extend(new_unique)
+        if self._embeddings:
+            # Deduplicate by semantic similarity (async-safe)
+            unique_evidence = await self._embeddings.deduplicate(result.evidence)
+
+            # Also search for semantically related evidence (async-safe)
+            related = await self._embeddings.search_similar(query, n_results=5)
+
+            # Merge related evidence not already in results
+            # We need to reconstruct Evidence objects from stored data
+            existing_urls = {e.citation.url for e in unique_evidence}
+
+            # Also check what's already in the global store to avoid re-adding
+            # logic here is a bit complex: deduplicate returned unique from *new search*
+            # but we also want related from *previous searches*
+
+            related_evidence = []
+            for item in related:
+                if item["id"] not in existing_urls:
+                    # Create Evidence from stored metadata
+                    # Check if metadata has required fields
+                    meta = item.get("metadata", {})
+                    # Fallback if date missing
+                    date = meta.get("date") or "n.d."
+                    authors = meta.get("authors")  # Might be list or string depending on how stored
+                    if isinstance(authors, str):
+                        authors = [authors]
+                    if not authors:
+                        authors = ["Unknown"]
+
+                    ev = Evidence(
+                        content=item["content"],
+                        citation=Citation(
+                            title=meta.get("title", "Untitled"),
+                            url=item["id"],
+                            source=meta.get("source", "vector_db"),
+                            date=date,
+                            authors=authors,
+                        ),
+                        relevance=item.get("distance", 0.0),  # Use distance/similarity as proxy
+                    )
+                    related_evidence.append(ev)
+
+            # Combine
+            final_new_evidence = unique_evidence + related_evidence
+
+            # Add to global store (deduping against global store)
+            global_urls = {e.citation.url for e in self._evidence_store["current"]}
+            really_new = [e for e in final_new_evidence if e.citation.url not in global_urls]
+            self._evidence_store["current"].extend(really_new)
+
+            # Update result for reporting
+            total_new = len(really_new)
+
+        else:
+            # Fallback to URL-based deduplication
+            existing_urls = {e.citation.url for e in self._evidence_store["current"]}
+            new_unique = [e for e in result.evidence if e.citation.url not in existing_urls]
+            self._evidence_store["current"].extend(new_unique)
+            total_new = len(new_unique)
 
         # Format response
+        # Get latest N items from store or just the new ones
+        # Let's show what was found in this run + related
+
+        evidence_to_show = (
+            (unique_evidence + related_evidence)
+            if self._embeddings and "unique_evidence" in locals()
+            else result.evidence
+        )
+
         evidence_text = "\n".join(
             [
                 f"- [{e.citation.title}]({e.citation.url}): {e.content[:200]}..."
-                for e in result.evidence[:5]
+                for e in evidence_to_show[:5]
             ]
         )
 
         response_text = (
-            f"Found {result.total_found} sources ({len(new_unique)} new):\n\n{evidence_text}"
+            f"Found {result.total_found} sources ({total_new} new added to context):\n\n"
+            f"{evidence_text}"
         )
 
         return AgentRunResponse(
             messages=[ChatMessage(role=Role.ASSISTANT, text=response_text)],
             response_id=f"search-{result.total_found}",
-            additional_properties={"evidence": [e.model_dump() for e in result.evidence]},
+            additional_properties={"evidence": [e.model_dump() for e in evidence_to_show]},
         )
 
     async def run_stream(

src/orchestrator.py

@@ -263,7 +263,7 @@ class Orchestrator:
 
         citations = "\n".join(
             [
-                f"{i+1}. [{e.citation.title}]({e.citation.url}) "
+                f"{i + 1}. [{e.citation.title}]({e.citation.url}) "
                 f"({e.citation.source.upper()}, {e.citation.date})"
                 for i, e in enumerate(evidence[:10])  # Limit to 10 citations
             ]
@@ -312,7 +312,7 @@ class Orchestrator:
         """
         citations = "\n".join(
             [
-                f"{i+1}. [{e.citation.title}]({e.citation.url}) ({e.citation.source.upper()})"
+                f"{i + 1}. [{e.citation.title}]({e.citation.url}) ({e.citation.source.upper()})"
                 for i, e in enumerate(evidence[:10])
             ]
         )

src/orchestrator_magentic.py

@@ -54,8 +54,22 @@ class MagenticOrchestrator:
             iteration=0,
         )
 
+        # Initialize embedding service (optional)
+        embedding_service = None
+        try:
+            from src.services.embeddings import get_embedding_service
+
+            embedding_service = get_embedding_service()
+            logger.info("Embedding service enabled")
+        except ImportError:
+            logger.info("Embedding service not available (dependencies missing)")
+        except Exception as e:
+            logger.warning("Failed to initialize embedding service", error=str(e))
+
         # Create agent wrappers
-        search_agent = SearchAgent(self._search_handler, self._evidence_store)
+        search_agent = SearchAgent(
+            self._search_handler, self._evidence_store, embedding_service=embedding_service
+        )
         judge_agent = JudgeAgent(self._judge_handler, self._evidence_store)
 
         # Build Magentic workflow
@@ -78,8 +92,17 @@ class MagenticOrchestrator:
         )
 
         # Task instruction for the manager
-        task = f"""Research drug repurposing opportunities for: {query}
+        semantic_note = ""
+        if embedding_service:
+            semantic_note = """
+The system has semantic search enabled. When evidence is found:
+1. Related concepts will be automatically surfaced
+2. Duplicates are removed by meaning, not just URL
+3. Use the surfaced related concepts to refine searches
+"""
 
+        task = f"""Research drug repurposing opportunities for: {query}
+{semantic_note}
 Instructions:
 1. Use SearcherAgent to find evidence. SEND ONLY A SIMPLE KEYWORD QUERY (e.g. "metformin aging")
    as the instruction. Complex queries fail.

src/services/__init__.py

@@ -0,0 +1 @@
+"""Services for DeepCritical."""

src/services/embeddings.py

@@ -0,0 +1,132 @@
+"""Embedding service for semantic search.
+
+IMPORTANT: All public methods are async to avoid blocking the event loop.
+The sentence-transformers model is CPU-bound, so we use run_in_executor().
+"""
+
+import asyncio
+from typing import Any
+
+import chromadb
+from sentence_transformers import SentenceTransformer
+
+from src.utils.models import Evidence
+
+
+class EmbeddingService:
+    """Handles text embedding and vector storage.
+
+    All embedding operations run in a thread pool to avoid blocking
+    the async event loop. See src/tools/websearch.py for the pattern.
+    """
+
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        self._model = SentenceTransformer(model_name)
+        self._client = chromadb.Client()  # In-memory for hackathon
+        self._collection = self._client.create_collection(
+            name="evidence", metadata={"hnsw:space": "cosine"}
+        )
+
+    # ─────────────────────────────────────────────────────────────────
+    # Sync internal methods (run in thread pool)
+    # ─────────────────────────────────────────────────────────────────
+
+    def _sync_embed(self, text: str) -> list[float]:
+        """Synchronous embedding - DO NOT call directly from async code."""
+        return list(self._model.encode(text).tolist())
+
+    def _sync_batch_embed(self, texts: list[str]) -> list[list[float]]:
+        """Batch embedding for efficiency - DO NOT call directly from async code."""
+        return [list(e.tolist()) for e in self._model.encode(texts)]
+
+    # ─────────────────────────────────────────────────────────────────
+    # Async public methods (safe for event loop)
+    # ─────────────────────────────────────────────────────────────────
+
+    async def embed(self, text: str) -> list[float]:
+        """Embed a single text (async-safe).
+
+        Uses run_in_executor to avoid blocking the event loop.
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self._sync_embed, text)
+
+    async def embed_batch(self, texts: list[str]) -> list[list[float]]:
+        """Batch embed multiple texts (async-safe, more efficient)."""
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self._sync_batch_embed, texts)
+
+    async def add_evidence(self, evidence_id: str, content: str, metadata: dict[str, Any]) -> None:
+        """Add evidence to vector store (async-safe)."""
+        embedding = await self.embed(content)
+        # ChromaDB operations are fast, but wrap for consistency
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(
+            None,
+            lambda: self._collection.add(
+                ids=[evidence_id],
+                embeddings=[embedding],  # type: ignore[arg-type]
+                metadatas=[metadata],
+                documents=[content],
+            ),
+        )
+
+    async def search_similar(self, query: str, n_results: int = 5) -> list[dict[str, Any]]:
+        """Find semantically similar evidence (async-safe)."""
+        query_embedding = await self.embed(query)
+
+        loop = asyncio.get_running_loop()
+        results = await loop.run_in_executor(
+            None,
+            lambda: self._collection.query(
+                query_embeddings=[query_embedding],  # type: ignore[arg-type]
+                n_results=n_results,
+            ),
+        )
+
+        # Handle empty results gracefully
+        ids = results.get("ids")
+        docs = results.get("documents")
+        metas = results.get("metadatas")
+        dists = results.get("distances")
+
+        if not ids or not ids[0] or not docs or not metas or not dists:
+            return []
+
+        return [
+            {"id": id, "content": doc, "metadata": meta, "distance": dist}
+            for id, doc, meta, dist in zip(
+                ids[0],
+                docs[0],
+                metas[0],
+                dists[0],
+                strict=False,
+            )
+        ]
+
+    async def deduplicate(
+        self, new_evidence: list[Evidence], threshold: float = 0.9
+    ) -> list[Evidence]:
+        """Remove semantically duplicate evidence (async-safe)."""
+        unique = []
+        for evidence in new_evidence:
+            similar = await self.search_similar(evidence.content, n_results=1)
+            if not similar or similar[0]["distance"] > (1 - threshold):
+                unique.append(evidence)
+                await self.add_evidence(
+                    evidence_id=evidence.citation.url,
+                    content=evidence.content,
+                    metadata={"source": evidence.citation.source},
+                )
+        return unique
+
+
+_embedding_service: EmbeddingService | None = None
+
+
+def get_embedding_service() -> EmbeddingService:
+    """Get singleton instance of EmbeddingService."""
+    global _embedding_service  # noqa: PLW0603
+    if _embedding_service is None:
+        _embedding_service = EmbeddingService()
+    return _embedding_service

tests/unit/agents/test_search_agent.py

@@ -81,5 +81,48 @@ async def test_run_handles_list_input(mock_handler: AsyncMock) -> None:
         ChatMessage(role=Role.USER, text="test query"),
     ]
     await agent.run(messages)
-
     mock_handler.execute.assert_awaited_once_with("test query", max_results_per_tool=10)
+
+
+@pytest.mark.asyncio
+async def test_run_uses_embeddings(mock_handler: AsyncMock) -> None:
+    """Test that run uses embedding service if provided."""
+    store: dict = {"current": []}
+
+    # Mock embedding service
+    mock_embeddings = AsyncMock()
+    # Mock deduplicate to return the evidence as is (or filtered)
+    mock_embeddings.deduplicate.return_value = [
+        Evidence(
+            content="unique content",
+            citation=Citation(source="pubmed", url="u1", title="t1", date="2024"),
+        )
+    ]
+    # Mock search_similar to return related items
+    mock_embeddings.search_similar.return_value = [
+        {
+            "id": "u2",
+            "content": "related content",
+            "metadata": {"source": "web", "title": "related", "date": "2024"},
+            "distance": 0.1,
+        }
+    ]
+
+    agent = SearchAgent(mock_handler, store, embedding_service=mock_embeddings)
+
+    await agent.run("test query")
+
+    # Verify deduplicate called
+    mock_embeddings.deduplicate.assert_awaited_once()
+
+    # Verify semantic search called
+    mock_embeddings.search_similar.assert_awaited_once_with("test query", n_results=5)
+
+    # Verify store contains related evidence (if logic implemented to add it)
+    # Note: logic for adding related evidence needs to be implemented in SearchAgent
+    # The spec says: "Merge related evidence not already in results"
+
+    # Check if u1 (deduplicated result) is in store
+    assert any(e.citation.url == "u1" for e in store["current"])
+    # Check if u2 (related result) is in store
+    assert any(e.citation.url == "u2" for e in store["current"])

tests/unit/services/test_embeddings.py

@@ -0,0 +1,148 @@
+"""Unit tests for EmbeddingService."""
+
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+# Skip if embeddings dependencies are not installed
+pytest.importorskip("chromadb")
+pytest.importorskip("sentence_transformers")
+
+from src.services.embeddings import EmbeddingService
+
+
+class TestEmbeddingService:
+    @pytest.fixture
+    def mock_sentence_transformer(self):
+        with patch("src.services.embeddings.SentenceTransformer") as mock_st_class:
+            mock_model = mock_st_class.return_value
+            # Mock encode to return a numpy array
+            mock_model.encode.return_value = np.array([0.1, 0.2, 0.3])
+            yield mock_model
+
+    @pytest.fixture
+    def mock_chroma_client(self):
+        with patch("src.services.embeddings.chromadb.Client") as mock_client_class:
+            mock_client = mock_client_class.return_value
+            mock_collection = mock_client.create_collection.return_value
+            # Mock query return structure
+            mock_collection.query.return_value = {
+                "ids": [["id1"]],
+                "documents": [["doc1"]],
+                "metadatas": [[{"source": "pubmed"}]],
+                "distances": [[0.1]],
+            }
+            yield mock_client
+
+    @pytest.mark.asyncio
+    async def test_embed_returns_vector(self, mock_sentence_transformer, mock_chroma_client):
+        """Embedding should return a float vector (async check)."""
+        service = EmbeddingService()
+        embedding = await service.embed("metformin diabetes")
+
+        assert isinstance(embedding, list)
+        assert len(embedding) == 3  # noqa: PLR2004
+        assert all(isinstance(x, float) for x in embedding)
+        # Ensure it ran in executor (mock encode called)
+        mock_sentence_transformer.encode.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_batch_embed_efficient(self, mock_sentence_transformer, mock_chroma_client):
+        """Batch embedding should call encode with list."""
+        # Setup mock for batch return (list of arrays)
+        import numpy as np
+
+        mock_sentence_transformer.encode.return_value = np.array([[0.1, 0.2], [0.3, 0.4]])
+
+        service = EmbeddingService()
+        texts = ["text one", "text two"]
+
+        batch_results = await service.embed_batch(texts)
+
+        assert len(batch_results) == 2  # noqa: PLR2004
+        assert isinstance(batch_results[0], list)
+        mock_sentence_transformer.encode.assert_called_with(texts)
+
+    @pytest.mark.asyncio
+    async def test_add_and_search(self, mock_sentence_transformer, mock_chroma_client):
+        """Should be able to add evidence and search for similar."""
+        service = EmbeddingService()
+        await service.add_evidence(
+            evidence_id="test1",
+            content="Metformin activates AMPK pathway",
+            metadata={"source": "pubmed"},
+        )
+
+        # Verify add was called
+        mock_collection = mock_chroma_client.create_collection.return_value
+        mock_collection.add.assert_called_once()
+
+        results = await service.search_similar("AMPK activation drugs", n_results=1)
+
+        # Verify query was called
+        mock_collection.query.assert_called_once()
+        assert len(results) == 1
+        assert results[0]["id"] == "id1"
+
+    @pytest.mark.asyncio
+    async def test_search_similar_empty_collection(
+        self, mock_sentence_transformer, mock_chroma_client
+    ):
+        """Search on empty collection should return empty list, not error."""
+        mock_collection = mock_chroma_client.create_collection.return_value
+        mock_collection.query.return_value = {
+            "ids": [[]],
+            "documents": [[]],
+            "metadatas": [[]],
+            "distances": [[]],
+        }
+
+        service = EmbeddingService()
+        results = await service.search_similar("anything", n_results=5)
+        assert results == []
+
+    @pytest.mark.asyncio
+    async def test_deduplicate(self, mock_sentence_transformer, mock_chroma_client):
+        """Deduplicate should remove similar items."""
+        from src.utils.models import Citation, Evidence
+
+        service = EmbeddingService()
+
+        # Mock search to return a match for the first item (duplicate)
+        # and no match for the second (unique)
+        mock_collection = mock_chroma_client.create_collection.return_value
+
+        # First call returns match (distance 0.05 < threshold)
+        # Second call returns no match or high distance
+        mock_collection.query.side_effect = [
+            {
+                "ids": [["existing_id"]],
+                "documents": [["doc"]],
+                "metadatas": [[{}]],
+                "distances": [[0.05]],  # Very similar
+            },
+            {
+                "ids": [[]],  # No match
+                "documents": [[]],
+                "metadatas": [[]],
+                "distances": [[]],
+            },
+        ]
+
+        evidence = [
+            Evidence(
+                content="Duplicate content",
+                citation=Citation(source="web", url="u1", title="t1", date="2024"),
+            ),
+            Evidence(
+                content="Unique content",
+                citation=Citation(source="web", url="u2", title="t2", date="2024"),
+            ),
+        ]
+
+        unique = await service.deduplicate(evidence, threshold=0.9)
+
+        # Only the unique one should remain
+        assert len(unique) == 1
+        assert unique[0].citation.url == "u2"