MongoDB Vector Agent

A single-agent pipeline using MongoDB for vector search. The agent creates a collection, inserts documents with sentence-transformer embeddings, computes cosine similarity via aggregation (local MongoDB does not support $vectorSearch -- that requires Atlas), evaluates result diversity, and synthesizes an answer with an LLM. Demonstrates SDK primitives across real MongoDB vector operations.

Environment variables

This example requires OPENAI_API_KEY, WAXELL_API_KEY, and WAXELL_API_URL. Requires MongoDB on localhost:27017. Use --dry-run to skip real API calls.

Architecture

Key Code

Collection Creation and Cosine Similarity Search with @tool

The agent inserts documents with embeddings and performs manual cosine similarity search via Python (local MongoDB fallback).

@waxell.tool(tool_type="database", name="mongo_create_collection")
def create_collection(db, docs: list[dict], embeddings) -> dict:
    """Create collection and insert documents with embeddings."""
    db.drop_collection(COLLECTION_NAME)
    collection = db[COLLECTION_NAME]
    records = []
    for doc, emb in zip(docs, embeddings):
        records.append({**doc, "embedding": emb.tolist()})
    result = collection.insert_many(records)
    return {"inserted": len(result.inserted_ids), "collection": COLLECTION_NAME}

@waxell.tool(tool_type="vector_db", name="mongo_vector_search")
def vector_search(db, query_embedding, limit: int = 3) -> dict:
    """Manual cosine similarity search via aggregation pipeline."""
    collection = db[COLLECTION_NAME]
    all_docs = list(collection.find({}, {"embedding": 1, "title": 1, "text": 1}))
    q_norm = q_vec / np.linalg.norm(q_vec)
    scored = []
    for doc in all_docs:
        d_norm = np.array(doc["embedding"]) / np.linalg.norm(np.array(doc["embedding"]))
        score = float(np.dot(q_norm, d_norm))
        scored.append({"title": doc["title"], "text": doc["text"], "score": round(score, 4)})
    scored.sort(key=lambda x: x["score"], reverse=True)
    return {"hits": scored[:limit], "total_scanned": len(all_docs)}

Diversity Assessment with @reasoning

@waxell.reasoning_dec(step="assess_result_diversity")
def assess_result_diversity(results: list[dict]) -> dict:
    categories = set()
    for r in results:
        text = r.get("text", "")
        categories.add(text.split(":")[0].strip() if ":" in text else "unknown")
    return {
        "thought": f"Checking diversity across {len(results)} results",
        "evidence": [f"Unique categories: {len(categories)}", f"Categories: {', '.join(sorted(categories))}"],
        "conclusion": "Diverse results" if len(categories) >= len(results) * 0.6 else "Concentrated results",
    }

What this demonstrates

• @waxell.observe -- single agent with full lifecycle tracing
• @waxell.tool(tool_type="database") -- MongoDB CRUD operations (insert_many, drop_collection) recorded as database tool spans
• @waxell.tool(tool_type="vector_db") -- cosine similarity search via aggregation recorded as vector DB tool span
• @waxell.retrieval(source="mongodb") -- result extraction recorded with MongoDB as the source
• @waxell.decision -- search breadth selection (narrow_3 vs broad_5 based on collection size)
• @waxell.reasoning_dec -- result diversity assessment across categories
• waxell.score() -- search quality and answer quality scores attached to the trace
• waxell.tag() / waxell.metadata() -- vector DB type, database name, collection, and MongoDB host
• Auto-instrumented LLM calls -- OpenAI synthesis captured without extra code
• Real MongoDB operations -- actual pymongo queries with sentence-transformers embeddings

Run it

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

# Live mode (requires MongoDB on localhost:27017)
export OPENAI_API_KEY="sk-..."
python -m app.demos.mongodb_vector_agent

# Custom query
python -m app.demos.mongodb_vector_agent --dry-run --query "Search for AI deployment best practices"

Source

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