Skip to content
3 min read · 566 words

RAG Retrieval

One function does all retrieval: embed the query, search the shared collection filtered to this agent, return top-k chunks with scores. Two endpoints expose it — raw search for inspection, and the ask pipeline for answers.

The retrieval function

python
def search_agent_documents(agent_id: int, query: str, top_k: int = 3):
    query_embedding = create_embedding(query)            # same model as ingestion
 
    results = documents_collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        where={"agent_id": agent_id},                    # ← tenant isolation
    )
    # zip Chroma's parallel arrays into dicts:
    # {chunk_text, score, document_id, filename, chunk_index}

Three properties to internalize:

  1. Symmetric embedding. The query is embedded with the same text-embedding-3-small model as the chunks — distances are only meaningful within one embedding space (warning).
  2. Agent scoping is a metadata filter, not separate collections. One agent_documents collection holds everyone's vectors; where={"agent_id": ...} restricts the search. Any new query path must keep this filter.
  3. score is a raw distance — lower is better. It is not normalized similarity; do not render it as a percentage.
Rendering diagram…

The raw search endpoint

POST /agents/{id}/documents/search — gated on the document_search tool (403 otherwise), logs a document_search ToolCall, no generation:

bash
curl -X POST http://localhost:8001/agents/1/documents/search \
  -H "Content-Type: application/json" \
  -d '{"query": "support SLA", "top_k": 3}'
json
{
  "query": "support SLA",
  "results": [
    {
      "chunk_text": "Premium customers have a 4-hour support SLA...",
      "score": 0.31,
      "document_id": 1,
      "filename": "sla.txt",
      "chunk_index": 0
    }
  ]
}

Use it for: debugging why the ask endpoint refused (was the right chunk even retrieved?), tuning top_k, and building custom retrieval UX. The ask pipeline calls the exact same function, so raw search is a faithful window into what the evaluator saw (Corrective RAG).

Choosing top_k

top_kEffect
1–2tight context, cheapest evaluator/answer calls; risks missing multi-chunk answers
3 (default)balanced — matches the chunk size so context stays ≈ 600 tokens
5+better recall for scattered facts; inflates evaluator+answer input tokens and dilutes context

top_k comes from the request body (DocumentSearchRequest / RAGAnswerRequest, default 3) with no server-side cap — an extreme value is the caller's own cost problem (Security).

Failure & edge behavior

CaseBehavior
No documents / no matches for the agentempty results (search) · short-circuit to refusal (ask)
Fewer than top_k vectors existChroma returns what it has
Embedding API failure500; failed document_search ToolCall committed
Cross-agent query attemptimpossible by construction — the filter is applied server-side

Quality levers (when retrieval feels wrong)

In diagnostic order:

  1. Look at raw search results for the failing question — is the needed chunk present at all?
  2. Chunking granularity — facts split across a boundary retrieve poorly; the 100-char overlap mitigates but doesn't eliminate (Ingestion).
  3. Query phrasing — dense retrieval matches meaning, not keywords; restate the question closer to the document's vocabulary.
  4. top_k — raise it when answers span multiple chunks.

Bigger levers (hybrid search, re-ranking, query expansion) are roadmap territory.

PreviousIngestionNextCorrective RAG