diff --git a/docsgpt/retriever/classic_rag.py b/docsgpt/retriever/classic_rag.py index db5e09b9..bcf49038 100644 --- a/docsgpt/retriever/classic_rag.py +++ b/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 diff --git a/docsgpt/retriever/dispatcher.py b/docsgpt/retriever/dispatcher.py index 31cc75e0..7ec23c5c 100644 --- a/docsgpt/retriever/dispatcher.py +++ b/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) diff --git a/docsgpt/retriever/fanout.py b/docsgpt/retriever/fanout.py index cb1ad3a6..fcca623d 100644 --- a/docsgpt/retriever/fanout.py +++ b/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( diff --git a/docsgpt/retriever/graph_rag.py b/docsgpt/retriever/graph_rag.py index 8d191f7b..59ee6163 100644 --- a/docsgpt/retriever/graph_rag.py +++ b/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 diff --git a/docsgpt/retriever/stages/prescreen.py b/docsgpt/retriever/stages/prescreen.py index cae50dea..2dfa5140 100644 --- a/docsgpt/retriever/stages/prescreen.py +++ b/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( diff --git a/docsgpt/tracing/core.py b/docsgpt/tracing/core.py index b0289152..46e9c4a7 100644 --- a/docsgpt/tracing/core.py +++ b/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" diff --git a/docsgpt/tracing/retrieval.py b/docsgpt/tracing/retrieval.py new file mode 100644 index 00000000..f9ce7103 --- /dev/null +++ b/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}, + ) diff --git a/tests/tracing/test_retrieval_spans.py b/tests/tracing/test_retrieval_spans.py new file mode 100644 index 00000000..be45d98c --- /dev/null +++ b/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)