Trace RAG retrieval, query embeddings and prescreen

The dispatcher and each retriever's search open retrieval spans recording
sources, chunk counts, top scores and a chunk preview. The shared query
embedding, every per-source vector search and the prescreen rerank get
their own spans; pool workers carry the trace so they nest correctly.

docsgpt/retriever/classic_rag.py

@@ -2,6 +2,11 @@ import logging
 from typing import Any, Dict, List, Optional, Tuple
 
 from docsgpt.core.settings import settings
+from docsgpt.tracing.retrieval import (
+    describe_documents,
+    start_retrieval_span,
+    start_source_search_span,
+)
 from docsgpt.llm.llm_creator import LLMCreator
 from docsgpt.retriever.base import BaseRetriever
 from docsgpt.retriever.fanout import fetch_per_source, max_parallel_sources
@@ -266,6 +271,15 @@ class ClassicRAG(BaseRetriever):
         logged and reported as ``None`` so one bad source cannot take the rest
         of the retrieval down with it.
         """
+        span = start_source_search_span(
+            plan["id"],
+            top_k=plan["src_k"],
+            **{
+                "docsgpt.vector_store": settings.VECTOR_STORE,
+                "docsgpt.retriever": type(self).__name__,
+                "docsgpt.shared_query_vector": query_vector is not None,
+            },
+        )
         try:
             if docsearch is None:
                 docsearch = VectorCreator.create_vectorstore(
@@ -279,8 +293,14 @@ class ClassicRAG(BaseRetriever):
                 query_vector=query_vector,
             )
             score_kind = self._score_kind(docsearch) if self.include_scores else None
+            try:
+                candidates = len(docs_temp)
+            except TypeError:
+                candidates = None
+            span.end(attributes={"docsgpt.candidate_count": candidates})
             return docs_temp, score_kind
         except Exception as e:
+            span.end(error=e)
             logger.error(
                 f"Error searching vectorstore {plan['id']}: {e}", exc_info=True
             )
@@ -409,11 +429,20 @@ class ClassicRAG(BaseRetriever):
 
     def search(self, query: str = ""):
         """Search for documents using optional query override"""
-        if query:
-            self.original_question = query
-            # Invalidate the cached rephrase so a per-source path that opts in
-            # rephrases against the new query, not a stale one.
-            self._rephrased_question = None
-            self.question = self._rephrase_query()
-            self._rephrased_question = self.question
-        return self._get_data()
+        with start_retrieval_span(
+            f"retrieval {type(self).__name__}",
+            sources=self.vectorstores,
+            **{"docsgpt.retriever": type(self).__name__, "docsgpt.top_k": self.chunks},
+        ) as span:
+            if query:
+                self.original_question = query
+                # Invalidate the cached rephrase so a per-source path that opts in
+                # rephrases against the new query, not a stale one.
+                self._rephrased_question = None
+                self.question = self._rephrase_query()
+                self._rephrased_question = self.question
+            docs = self._get_data()
+            rephrased = self._rephrased_question
+            span.set(**{"docsgpt.rephrased": bool(rephrased and rephrased != self.original_question)})
+            describe_documents(span, docs, query=rephrased or self.original_question)
+            return docs

docsgpt/retriever/dispatcher.py

@@ -23,6 +23,7 @@ from docsgpt.retriever.stages.prescreen import (
     max_candidate_k,
 )
 from docsgpt.storage.db.source_config import RetrievalConfig
+from docsgpt.tracing.retrieval import describe_documents, start_retrieval_span
 from docsgpt.utils import num_tokens_from_string
 
 logger = logging.getLogger(__name__)
@@ -294,6 +295,17 @@ class Dispatcher(BaseRetriever):
 
     def search(self, query: str = "") -> List[Dict[str, Any]]:
         """Run every group under the shared budget and merge the results."""
+        sources = [s for group in self._groups for s in group.get("doc_ids", [])]
+        with start_retrieval_span(
+            "retrieval",
+            sources=sources,
+            **{"docsgpt.retriever": "Dispatcher", "docsgpt.group_count": len(self._groups)},
+        ) as span:
+            docs = self._search_groups(query)
+            describe_documents(span, docs)
+            return docs
+
+    def _search_groups(self, query: str) -> List[Dict[str, Any]]:
         groups = self._groups
         n_groups = len(groups)
 

docsgpt/retriever/fanout.py

@@ -12,7 +12,9 @@ import logging
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Callable, Dict, Iterable, List, Optional, TypeVar
 
+from docsgpt import tracing
 from docsgpt.core.settings import settings
+from docsgpt.tracing.retrieval import start_embedding_span
 
 logger = logging.getLogger(__name__)
 
@@ -85,16 +87,22 @@ def embed_questions(docsearch, questions: Iterable[str]) -> Dict[str, List[float
     embedder = store_embeddings(docsearch)
     if embedder is None:
         return {}
+    distinct = list(dict.fromkeys(questions))
+    span = start_embedding_span(
+        settings.EMBEDDINGS_NAME,
+        inputs=len(distinct),
+        **{"docsgpt.embeddings_backend": type(embedder).__name__},
+    )
     try:
-        return {
-            question: embedder.embed_query(question)
-            for question in dict.fromkeys(questions)
-        }
+        vectors = {question: embedder.embed_query(question) for question in distinct}
     except Exception as e:
+        span.end(error=e)
         logger.warning(
             "Query embedding failed (%s); each store will embed its own.", e
         )
         return {}
+    span.end()
+    return vectors
 
 
 def run_source_jobs(
@@ -121,10 +129,13 @@ def run_source_jobs(
         workers = max_parallel_sources(len(jobs))
     if workers == 1:
         return [fn(job) for job in jobs]
+    # Pool threads don't inherit context: carry the trace and its current
+    # retrieval span in so per-source spans nest under it.
+    traced_fn = tracing.wrap(fn)
     with ThreadPoolExecutor(
         max_workers=workers, thread_name_prefix="rag-source"
     ) as pool:
-        return list(pool.map(fn, jobs))
+        return list(pool.map(traced_fn, jobs))
 
 
 def fetch_per_source(

docsgpt/retriever/graph_rag.py

@@ -33,6 +33,11 @@ from docsgpt.retriever.base import BaseRetriever
 from docsgpt.retriever.classic_rag import ClassicRAG
 from docsgpt.retriever.labels import labels_from_metadata
 from docsgpt.storage.db.source_config import GraphRetrievalConfig
+from docsgpt.tracing.retrieval import (
+    describe_documents,
+    start_embedding_span,
+    start_retrieval_span,
+)
 from docsgpt.utils import num_tokens_from_string
 from docsgpt.vectorstore.base import get_embeddings
 
@@ -231,7 +236,11 @@ class GraphRAGRetriever(BaseRetriever):
 
     def _embed_query(self, question: str) -> List[float]:
         embedding = get_embeddings()
-        return embedding.embed_query(question)
+        with start_embedding_span(
+            settings.EMBEDDINGS_NAME,
+            **{"docsgpt.embeddings_backend": type(embedding).__name__},
+        ):
+            return embedding.embed_query(question)
 
     def _ppr_scores(self, subgraph, seeds) -> Dict[str, float]:
         """Run Personalized PageRank, then down-weight hub nodes by IDF.
@@ -730,10 +739,19 @@ class GraphRAGRetriever(BaseRetriever):
                 logging.debug("Error closing GraphRAG store: %s", e)
 
     def search(self, query: str = "") -> List[Dict[str, Any]]:
-        if query:
-            self.original_question = query
-            self._classic.original_question = query
-            self._classic._rephrased_question = None
-            self._classic.question = self._classic._rephrase_query()
-            self._classic._rephrased_question = self._classic.question
-        return self._get_data()
+        with start_retrieval_span(
+            f"retrieval {type(self).__name__}",
+            sources=self.vectorstores,
+            **{"docsgpt.retriever": type(self).__name__, "docsgpt.top_k": self.chunks},
+        ) as span:
+            if query:
+                self.original_question = query
+                self._classic.original_question = query
+                self._classic._rephrased_question = None
+                self._classic.question = self._classic._rephrase_query()
+                self._classic._rephrased_question = self._classic.question
+            docs = self._get_data()
+            describe_documents(
+                span, docs, query=self._classic._rephrased_question or self.original_question
+            )
+            return docs

docsgpt/retriever/stages/prescreen.py

@@ -18,6 +18,7 @@ import re
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Callable, Dict, List, Optional
 
+from docsgpt import tracing
 from docsgpt.llm.llm_creator import LLMCreator
 from docsgpt.storage.db.source_config import PreScreenConfig
 
@@ -164,20 +165,28 @@ class PreScreenStage:
         ]
         llm = self._build_llm()
 
-        max_workers = min(_MAX_WORKERS, len(batches))
-        kept: List[Dict[str, Any]] = []
-        if max_workers <= 1:
-            for batch in batches:
-                kept.extend(self._screen_batch(llm, query, batch))
-        else:
-            with ThreadPoolExecutor(max_workers=max_workers) as pool:
-                results = pool.map(
-                    lambda b: self._screen_batch(llm, query, b), batches
-                )
-                for batch_result in results:
-                    kept.extend(batch_result)
-
-        return kept[: self.config.max_keep]
+        with tracing.span(
+            tracing.KIND_RERANK,
+            "rerank prescreen",
+            attributes={
+                "gen_ai.operation.name": "rerank",
+                "docsgpt.candidate_count": len(docs),
+                "docsgpt.batch_count": len(batches),
+            },
+        ) as span:
+            max_workers = min(_MAX_WORKERS, len(batches))
+            kept: List[Dict[str, Any]] = []
+            if max_workers <= 1:
+                for batch in batches:
+                    kept.extend(self._screen_batch(llm, query, batch))
+            else:
+                screen = tracing.wrap(lambda b: self._screen_batch(llm, query, b))
+                with ThreadPoolExecutor(max_workers=max_workers) as pool:
+                    for batch_result in pool.map(screen, batches):
+                        kept.extend(batch_result)
+            survivors = kept[: self.config.max_keep]
+            span.set(**{"docsgpt.kept_count": len(survivors)})
+            return survivors
 
 
 def build_prescreen_stages(

docsgpt/tracing/core.py

@@ -40,7 +40,7 @@ KIND_GUARDRAIL = "guardrail"
 KIND_STEP = "step"
 
 #: Kinds that become the implicit parent of spans started while they are open.
-CONTAINER_KINDS = frozenset({KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_STEP})
+CONTAINER_KINDS = frozenset({KIND_AGENT, KIND_TOOL, KIND_RETRIEVAL, KIND_RERANK, KIND_STEP})
 
 STATUS_OK = "ok"
 STATUS_ERROR = "error"

docsgpt/tracing/retrieval.py

@@ -0,0 +1,94 @@
+"""Spans for RAG retrieval: the search itself, the query embedding, each source."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Iterable, List, Optional
+
+from docsgpt.tracing import core
+
+#: Characters of each retrieved chunk kept in the stored preview.
+_SNIPPET_CHARS = 240
+#: Chunks listed in the preview; the count attribute still reports all of them.
+_MAX_PREVIEW_CHUNKS = 20
+
+
+def start_retrieval_span(name: str, *, sources: Optional[Iterable[Any]] = None, **attributes: Any):
+    """Open a ``retrieval`` span (a container: embeddings and searches nest under it)."""
+    source_ids = [str(s) for s in (sources or []) if s]
+    base = {
+        "gen_ai.operation.name": "retrieval",
+        "docsgpt.source_ids": source_ids or None,
+        "gen_ai.data_source.id": source_ids[0] if len(source_ids) == 1 else None,
+    }
+    base.update(attributes)
+    return core.start_span(
+        core.KIND_RETRIEVAL,
+        name,
+        attributes={k: v for k, v in base.items() if v is not None},
+    )
+
+
+def describe_documents(span: Any, docs: Optional[List[Dict[str, Any]]], *, query: Optional[str] = None) -> None:
+    """Record what a retrieval returned: counts, sources, top scores and a chunk preview."""
+    if not span:
+        return
+    docs = docs or []
+    scores = [d.get("score") for d in docs if isinstance(d, dict) and isinstance(d.get("score"), (int, float))]
+    score_kinds = {d.get("score_kind") for d in docs if isinstance(d, dict) and d.get("score_kind")}
+    span.set(
+        **{
+            "docsgpt.chunk_count": len(docs),
+            "docsgpt.top_score": max(scores) if scores else None,
+            "docsgpt.score_kind": score_kinds.pop() if len(score_kinds) == 1 else None,
+        }
+    )
+    if query:
+        span.preview("query", query)
+    if docs:
+        span.preview(
+            "chunks",
+            [
+                {
+                    k: v
+                    for k, v in {
+                        "title": d.get("title") or d.get("filename"),
+                        "source": d.get("source"),
+                        "score": d.get("score"),
+                        "text": str(d.get("text") or "")[:_SNIPPET_CHARS],
+                    }.items()
+                    if v not in (None, "")
+                }
+                for d in docs[:_MAX_PREVIEW_CHUNKS]
+                if isinstance(d, dict)
+            ],
+        )
+
+
+def start_embedding_span(model: Optional[str], *, inputs: int = 1, **attributes: Any):
+    """Open an ``embeddings`` span for a query embedding (leaf)."""
+    base = {
+        "gen_ai.operation.name": "embeddings",
+        "gen_ai.request.model": model,
+        "docsgpt.input_count": inputs,
+    }
+    base.update(attributes)
+    return core.start_span(
+        core.KIND_EMBEDDING,
+        f"embeddings {model}" if model else "embeddings",
+        attributes={k: v for k, v in base.items() if v is not None},
+    )
+
+
+def start_source_search_span(source_id: Any, *, top_k: Optional[int] = None, **attributes: Any):
+    """Open a per-source vector ``search`` span (leaf)."""
+    base = {
+        "gen_ai.operation.name": "search",
+        "gen_ai.data_source.id": str(source_id) if source_id else None,
+        "docsgpt.top_k": top_k,
+    }
+    base.update(attributes)
+    return core.start_span(
+        core.KIND_SEARCH,
+        f"search {source_id}" if source_id else "search",
+        attributes={k: v for k, v in base.items() if v is not None},
+    )

tests/tracing/test_retrieval_spans.py

@@ -0,0 +1,153 @@
+"""RAG retrieval becomes retrieval / embeddings / search / rerank spans."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from docsgpt import tracing
+from docsgpt.core.settings import settings
+from docsgpt.retriever import fanout
+
+
+@pytest.fixture(autouse=True)
+def _env(monkeypatch):
+    monkeypatch.setattr(settings, "TRACES_ENABLED", True)
+    monkeypatch.setattr(settings, "TRACES_CAPTURE_CONTENT", True)
+
+
+@pytest.fixture()
+def trace():
+    t = tracing.start_trace(source="stream", capture_otel_context=False)
+    with tracing.activate(t):
+        yield t
+
+
+class TestFanout:
+    def test_embedding_span(self, trace):
+        embedder = MagicMock()
+        embedder.embed_query.side_effect = lambda q: [0.1, 0.2]
+        store = MagicMock(_embedding=embedder)
+        vectors = fanout.embed_questions(store, ["q", "q", "r"])
+        assert set(vectors) == {"q", "r"}
+        (span,) = trace.spans
+        assert span.kind == tracing.KIND_EMBEDDING
+        assert span.attributes["gen_ai.operation.name"] == "embeddings"
+        assert span.attributes["docsgpt.input_count"] == 2
+
+    def test_embedding_failure_is_recorded(self, trace):
+        embedder = MagicMock()
+        embedder.embed_query.side_effect = RuntimeError("model down")
+        assert fanout.embed_questions(MagicMock(_embedding=embedder), ["q"]) == {}
+        assert trace.spans[0].status == "error"
+
+    def test_pool_jobs_nest_under_the_open_retrieval(self, trace, monkeypatch):
+        monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 4)
+
+        def job(i):
+            with tracing.span(tracing.KIND_SEARCH, f"search {i}"):
+                return i
+
+        with tracing.span(tracing.KIND_RETRIEVAL, "retrieval") as outer:
+            assert fanout.run_source_jobs(job, [1, 2, 3]) == [1, 2, 3]
+        searches = [s for s in trace.spans if s.kind == tracing.KIND_SEARCH]
+        assert len(searches) == 3
+        assert all(s.parent_id == outer.id for s in searches)
+
+
+class TestClassicRAG:
+    def _rag(self, stores):
+        from docsgpt.retriever.classic_rag import ClassicRAG
+
+        with patch("docsgpt.retriever.classic_rag.LLMCreator.create_llm", return_value=MagicMock()):
+            rag = ClassicRAG(
+                source={"question": "what is x", "active_docs": list(stores)},
+                chunks=4,
+                decoded_token={"sub": "u1"},
+            )
+        return rag
+
+    def test_search_records_retrieval_tree(self, trace, monkeypatch):
+        monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 4)
+        embedder = MagicMock()
+        embedder.embed_query.return_value = [0.1]
+
+        def make_store(*_args, **_kwargs):
+            store = MagicMock(_embedding=embedder)
+            store.search.return_value = [
+                {"text": "chunk body", "metadata": {"title": "Doc", "source": "doc.md"}}
+            ]
+            return store
+
+        rag = self._rag(["s1", "s2"])
+        with patch(
+            "docsgpt.retriever.classic_rag.VectorCreator.create_vectorstore",
+            side_effect=make_store,
+        ):
+            docs = rag.search()
+        assert len(docs) == 2
+        retrieval = trace.spans[0]
+        assert retrieval.kind == tracing.KIND_RETRIEVAL
+        assert retrieval.attributes["docsgpt.source_ids"] == ["s1", "s2"]
+        assert retrieval.attributes["docsgpt.chunk_count"] == 2
+        assert retrieval.previews["query"] == "what is x"
+        assert retrieval.previews["chunks"][0]["text"] == "chunk body"
+        children = [s for s in trace.spans[1:]]
+        assert {s.kind for s in children} == {tracing.KIND_EMBEDDING, tracing.KIND_SEARCH}
+        assert all(s.parent_id == retrieval.id for s in children)
+        searches = [s for s in children if s.kind == tracing.KIND_SEARCH]
+        assert sorted(s.attributes["gen_ai.data_source.id"] for s in searches) == ["s1", "s2"]
+        assert all(s.attributes["docsgpt.candidate_count"] == 1 for s in searches)
+
+    def test_failed_source_search_is_an_error_span(self, trace, monkeypatch):
+        monkeypatch.setattr(settings, "RETRIEVAL_MAX_PARALLEL_SOURCES", 1)
+        rag = self._rag(["s1", "s2"])
+
+        def make_store(_kind, source_id, _key):
+            if source_id == "s2":
+                raise RuntimeError("no store")
+            store = MagicMock(_embedding=None, _embeddings=None, embeddings=None)
+            store._get_embeddings.side_effect = RuntimeError("no embedder")
+            store.search.return_value = []
+            return store
+
+        with patch(
+            "docsgpt.retriever.classic_rag.VectorCreator.create_vectorstore",
+            side_effect=make_store,
+        ):
+            assert rag.search() == []
+        search = {s.attributes["gen_ai.data_source.id"]: s for s in trace.spans if s.kind == tracing.KIND_SEARCH}
+        assert search["s1"].status == "ok"
+        assert search["s2"].status == "error"
+        assert trace.spans[0].status == "ok"
+
+
+class TestPrescreen:
+    def test_rerank_span_contains_screening_calls(self, trace):
+        from docsgpt.retriever.stages.prescreen import PreScreenStage
+        from docsgpt.storage.db.source_config import PreScreenConfig
+
+        stage = PreScreenStage(
+            PreScreenConfig(candidate_k=10, max_keep=2, batch_size=1),
+            llm_name="openai",
+            api_key=None,
+            model_id="m",
+        )
+        llm = MagicMock()
+
+        def fake_gen(**_kwargs):
+            with tracing.span(tracing.KIND_LLM, "chat m"):
+                return '{"keep": [0]}'
+
+        llm.gen.side_effect = fake_gen
+        stage._build_llm = lambda: llm
+        docs = [{"text": f"d{i}"} for i in range(3)]
+        kept = stage(docs, {"query": "q"})
+        assert len(kept) == 2
+        rerank = trace.spans[0]
+        assert rerank.kind == tracing.KIND_RERANK
+        assert rerank.attributes["docsgpt.kept_count"] == 2
+        llm_spans = [s for s in trace.spans if s.kind == tracing.KIND_LLM]
+        assert len(llm_spans) == 3
+        assert all(s.parent_id == rerank.id for s in llm_spans)