ChromaDB semantic search

The retrieval branch fires when the orchestrator sees a factual question. The layer turns that question into an embedding, looks up the closest documents in ChromaDB, and passes them to the synthesis step. ChromaDB runs embedded in your process with Sentence Transformers, so you do not need a separate vector service to get started.

layers/retrieval.py
python
try:
    import chromadb
    from chromadb.utils import embedding_functions
    _CHROMA_AVAILABLE = True
except Exception:
    _CHROMA_AVAILABLE = False


SEED_DOCS = [
    ("production-agent", "A production agent is an LLM system with transport, orchestration, tools, memory, retrieval, guardrails, and observability layers."),
    ("guardrails", "Guardrails scrub PII, limit input length, and block unsafe outputs before they reach users."),
    ("observability", "Observability captures per-request traces with timings for each layer so failures can be debugged in minutes."),
    ("orchestration", "Orchestration is a small state machine that decides whether to call a tool, retrieve context, or reply directly."),
]


class RetrievalLayer:
    def __init__(self, collection_name: str = "prod_agent_kb") -> None:
        self.available = _CHROMA_AVAILABLE
        self._collection = None
        if not self.available:
            return

        client = chromadb.Client()
        embed_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name="all-MiniLM-L6-v2"
        )
        self._collection = client.get_or_create_collection(
            name=collection_name,
            embedding_function=embed_fn,
        )
        if self._collection.count() == 0:
            self._collection.add(
                ids=[doc_id for doc_id, _ in SEED_DOCS],
                documents=[text for _, text in SEED_DOCS],
            )

    def search(self, query: str, k: int = 3) -> List[str]:
        if not self.available or self._collection is None:
            return []
        res = self._collection.query(query_texts=[query], n_results=k)
        docs = res.get("documents", [[]])
        return docs[0] if docs else []

Import guard, seed documents, search method. Everything the rest of the pipeline needs from retrieval is behind this small surface. Notice how availability is a first-class attribute: when the vector DB is missing, the agent keeps running with an empty retrieval step.

Flashcards: Flashcards

Loading practice…

Quiz: Quiz

Loading practice…