ChromaDB Agent

A multi-agent document search pipeline using ChromaDB as the vector store. A parent orchestrator coordinates 3 child agents -- an indexer that creates collections and adds documents, a searcher that queries and ranks results, and a synthesizer that generates answers with LLM-powered quality assessment. The pipeline demonstrates every major SDK primitive across ChromaDB-specific vector operations.

Environment variables

This example requires OPENAI_API_KEY, WAXELL_API_KEY, and WAXELL_API_URL. Use --dry-run to skip real API calls.

Architecture

Key Code

ChromaDB Collection Management with @tool

The indexer child agent uses @tool(tool_type="vector_db") for both collection creation and document ingestion. Each operation is recorded as a separate span in the trace.

@waxell.tool(tool_type="vector_db")
def get_or_create_collection(name: str, metadata: dict) -> dict:
    """Get or create a ChromaDB collection."""
    return {"name": name, "metadata": metadata, "count": 0}

@waxell.tool(tool_type="vector_db")
def add_documents(collection_name: str, ids: list, documents: list,
                  metadatas: list) -> dict:
    """Add documents to a ChromaDB collection."""
    return {"added_count": len(ids), "ids": ids, "collection": collection_name}

@waxell.observe(agent_name="chroma-indexer", workflow_name="document-indexing")
async def run_indexer(documents: list, waxell_ctx=None):
    waxell.tag("vector_db", "chromadb")
    waxell.tag("agent_role", "indexer")

    collection_info = get_or_create_collection(
        name="governance-docs",
        metadata={"hnsw:space": "cosine", "description": "AI governance docs"},
    )
    waxell.metadata("collection_name", collection_info["name"])
    waxell.metadata("distance_metric", "cosine")

    add_result = add_documents(
        collection_name="governance-docs",
        ids=[d["id"] for d in documents],
        documents=[d["document"] for d in documents],
        metadatas=[d["metadata"] for d in documents],
    )
    waxell.metadata("documents_indexed", add_result["added_count"])
    return {"collection": collection_info["name"], "documents_added": add_result["added_count"]}

ChromaDB Query with @tool, @retrieval, and @reasoning

The searcher child agent queries ChromaDB, ranks results with @retrieval, and evaluates quality with @reasoning.

@waxell.tool(tool_type="vector_db")
def query_collection(collection_name: str, query_texts: list,
                     n_results: int = 3, where: dict | None = None) -> dict:
    """Query a ChromaDB collection for similar documents."""
    return {"ids": result_ids, "distances": result_distances, "num_results": len(result_ids)}

@waxell.retrieval(source="chromadb")
def rank_search_results(query: str, ids: list, distances: list,
                        documents: list, metadatas: list) -> list[dict]:
    """Rank and return matched documents from ChromaDB query results."""
    matched = []
    for i, (doc_id, dist) in enumerate(zip(ids, distances)):
        matched.append({
            "id": doc_id, "document": documents[i],
            "score": round(1.0 / (1.0 + dist), 4), "distance": dist,
        })
    return matched

@waxell.reasoning_dec(step="result_quality_evaluation")
async def evaluate_result_quality(results: list, query: str) -> dict:
    avg_score = sum(r.get("score", 0) for r in results) / max(len(results), 1)
    return {
        "thought": f"Retrieved {len(results)} documents with avg score {avg_score:.3f}.",
        "evidence": [f"Source: {r.get('id')} (score: {r.get('score', 0):.3f})" for r in results],
        "conclusion": "Results are relevant" if avg_score > 0.5 else "Results may need broader search",
    }

Synthesis with @reasoning and score()

The synthesizer builds context from retrieved documents, generates an LLM answer, then assesses quality.

@waxell.reasoning_dec(step="answer_quality_assessment")
async def assess_answer_quality(answer: str, sources: list) -> dict:
    coverage = len([s for s in sources if s.lower() in answer.lower()])
    return {
        "thought": f"Generated answer references {coverage}/{len(sources)} source documents.",
        "evidence": [f"Source: {s}" for s in sources],
        "conclusion": "Answer adequately covers source material" if coverage > 0 else "Answer may need more grounding",
    }

# Scores attached to the synthesizer span
waxell.score("answer_quality", 0.88, comment="auto-scored based on source coverage")
waxell.score("retrieval_relevance", 0.91, comment="based on ChromaDB distance scores")

What this demonstrates

• @waxell.observe -- parent-child agent hierarchy (orchestrator + 3 child agents) with automatic lineage via WaxellContext
• @waxell.tool(tool_type="vector_db") -- ChromaDB operations (create collection, add documents, query) recorded as tool spans
• @waxell.retrieval(source="chromadb") -- search-and-rank recorded with ChromaDB as the source
• @waxell.decision -- query classification via OpenAI with constrained options
• waxell.decide() -- manual collection routing decision with categories, reasoning, and confidence
• @waxell.reasoning_dec -- chain-of-thought evaluation for both search quality and answer quality
• @waxell.step_dec -- query preprocessing and context preparation recorded as execution steps
• waxell.score() -- answer quality and retrieval relevance scores attached to the trace
• waxell.tag() / waxell.metadata() -- vector DB type, agent role, collection name, and distance metric
• Auto-instrumented LLM calls -- OpenAI calls in classification and synthesis captured without extra code

Run it

# Dry-run mode (no API key needed)
cd dev/waxell-dev
python -m app.demos.chroma_agent --dry-run

# Live mode
export OPENAI_API_KEY="sk-..."
export WAXELL_API_KEY="your-waxell-api-key"
export WAXELL_API_URL="https://api.waxell.ai"
python -m app.demos.chroma_agent

# Custom query
python -m app.demos.chroma_agent --dry-run --query "Find AI deployment patterns"

Source

dev/waxell-dev/app/demos/chroma_agent.py