diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index b9f0ff60..2291d133 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -33,6 +33,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation, resolve_citations +from haiku.rag.telemetry import logfire from haiku.rag.tools.search import build_image_content_from_results CITATION_GRACE_REQUESTS = 2 @@ -350,9 +351,13 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): if self.rag is None: async with self.resource_lock: if self.rag is None: - rag = HaikuRAG(self.db_path, config=self.config, read_only=True) - await rag.__aenter__() - self.rag = rag + # Opening the store is lazy, so the run's first tool call + # pays it. Spanned separately or it reads as "the first + # search was slow" with nothing to point at. + with logfire.span("rag.client.open", db_path=str(self.db_path)): + rag = HaikuRAG(self.db_path, config=self.config, read_only=True) + await rag.__aenter__() + self.rag = rag return self.rag async def get_picture_bytes(self, document_id: str, self_ref: str) -> bytes | None: @@ -430,19 +435,34 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): "Search limit reached. Answer the question using " "the results you already have." ) - async with self.rag_lock: - formatted, results = await search_corpus( - await self._ensure_rag(), - query, - limit=limit, - document_filter=getattr(self.state, "document_filter", None), - ) - state = cast(Any, self.state) - state.searches[query] = results - self._note_evidence() - if self.vision and (parts := build_image_content_from_results(results)): - return ToolReturn(return_value=formatted, content=parts) - return formatted + # Opened after the budget check so the span means "a search ran" — + # a refused call is already visible as the tool span's exception. + # `search_index` is what shows whether cost grows across a run or + # one search was pathological; the tool span alone cannot tell them + # apart, and the first search also carries the one-off store open. + with logfire.span( + "ask.tool.search", + namespace=self.state_namespace, + search_index=self.search_count, + max_searches=self._max_searches, + limit=limit, + ) as span: + async with self.rag_lock: + formatted, results = await search_corpus( + await self._ensure_rag(), + query, + limit=limit, + document_filter=getattr(self.state, "document_filter", None), + ) + span.set_attribute("results", len(results)) + span.set_attribute("formatted_chars", len(formatted)) + state = cast(Any, self.state) + state.searches[query] = results + self._note_evidence() + if self.vision and (parts := build_image_content_from_results(results)): + span.set_attribute("image_parts", len(parts)) + return ToolReturn(return_value=formatted, content=parts) + return formatted async def _cite(self, chunk_ids: list[str]) -> str: """Register the evidence behind this answer, or declare there is none. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 46752242..0e720d84 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -4,11 +4,34 @@ from typing import TYPE_CHECKING from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX +from haiku.rag.telemetry import logfire if TYPE_CHECKING: from PIL import Image as PILImage from haiku.rag.client import HaikuRAG + from haiku.rag.reranking.base import RerankerBase + + +def _get_reranker(client: "HaikuRAG") -> "RerankerBase | None": + """Materialize the client's reranker, timing the first touch. + + ``HaikuRAG.reranker`` is a ``cached_property``, and the local rerankers + load model weights in their constructor — seconds of synchronous work on + the event loop that otherwise lands unattributed inside whichever search + happened to be first. Every later search hits the cache and emits + nothing, so a ``search.reranker.load`` span in a trace means a cold + process, not a slow reranker. + """ + if "reranker" in client.__dict__: + return client.reranker + model = client._config.reranking.model + with logfire.span( + "search.reranker.load", + provider=model.provider if model else None, + model=model.name if model else None, + ): + return client.reranker async def search( @@ -42,7 +65,7 @@ async def search( if search_type is None: search_type = "hybrid" - reranker = client.reranker + reranker = _get_reranker(client) if reranker is None: chunk_results = await client.chunk_repository.search( @@ -149,51 +172,64 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) if r.document_id and r.doc_item_refs: by_doc.setdefault(r.document_id, []).append(r) - for doc_id, doc_results in by_doc.items(): - all_refs = {ref for r in doc_results for ref in r.doc_item_refs} - caption_to_picture = await repo.get_caption_picture_refs(doc_id, list(all_refs)) + # Up to three reads per document plus base64 encoding of every blob, so + # `documents` is the round-trip multiplier and `bytes` the encoding load. + # Both are counted raw (pre-base64); the encoded payload is ~4/3 of it. + with logfire.span("search.images", documents=len(by_doc)) as span: + picture_count = 0 + picture_bytes = 0 + for doc_id, doc_results in by_doc.items(): + all_refs = {ref for r in doc_results for ref in r.doc_item_refs} + caption_to_picture = await repo.get_caption_picture_refs( + doc_id, list(all_refs) + ) - result_pictures: list[tuple[SearchResult, list[str]]] = [] - wanted: list[str] = [] - seen: set[str] = set() - for r in doc_results: - pictures: list[str] = [] - for ref in r.doc_item_refs: - picture = ( - ref - if ref.startswith(PICTURE_REF_PREFIX) - else caption_to_picture.get(ref) - ) - if picture and picture not in pictures: - pictures.append(picture) - if pictures: - result_pictures.append((r, pictures)) - for picture in pictures: - if picture not in seen: - wanted.append(picture) - seen.add(picture) - if not wanted: - continue - bytes_by_ref = await repo.get_pictures_for_chunk(doc_id, wanted) - if not bytes_by_ref: - continue - captions_by_ref = await repo.get_text_for_refs( - doc_id, list(bytes_by_ref.keys()) - ) - for r, pictures in result_pictures: - attached: dict[str, str] = {} - captions: dict[str, str] = {} - for ref in pictures: - blob = bytes_by_ref.get(ref) - if blob: - attached[ref] = base64.b64encode(blob).decode("ascii") - caption = captions_by_ref.get(ref) - if caption: - captions[ref] = caption - if attached: - r.image_data = attached - if captions: - r.picture_captions = captions + result_pictures: list[tuple[SearchResult, list[str]]] = [] + wanted: list[str] = [] + seen: set[str] = set() + for r in doc_results: + pictures: list[str] = [] + for ref in r.doc_item_refs: + picture = ( + ref + if ref.startswith(PICTURE_REF_PREFIX) + else caption_to_picture.get(ref) + ) + if picture and picture not in pictures: + pictures.append(picture) + if pictures: + result_pictures.append((r, pictures)) + for picture in pictures: + if picture not in seen: + wanted.append(picture) + seen.add(picture) + if not wanted: + continue + bytes_by_ref = await repo.get_pictures_for_chunk(doc_id, wanted) + if not bytes_by_ref: + continue + captions_by_ref = await repo.get_text_for_refs( + doc_id, list(bytes_by_ref.keys()) + ) + for r, pictures in result_pictures: + attached: dict[str, str] = {} + captions: dict[str, str] = {} + for ref in pictures: + blob = bytes_by_ref.get(ref) + if blob: + attached[ref] = base64.b64encode(blob).decode("ascii") + picture_count += 1 + picture_bytes += len(blob) + caption = captions_by_ref.get(ref) + if caption: + captions[ref] = caption + if attached: + r.image_data = attached + if captions: + r.picture_captions = captions + + span.set_attribute("pictures", picture_count) + span.set_attribute("bytes", picture_bytes) async def expand_context( @@ -223,23 +259,38 @@ async def expand_context( expanded_results = [] - for doc_id, doc_results in document_groups.items(): - if doc_id is None: - expanded_results.extend(doc_results) - continue + # Each expanded document costs two round trips (resolve_refs, then + # get_items_in_range), so `documents` is the multiplier on this stage. + # `context_chars` is the payload the model is about to be handed, which + # ties retrieval time to the next request's prompt size. + with logfire.span( + "search.expand", + documents=len(document_groups), + max_chars=max_chars, + results_in=len(search_results), + ) as span: + for doc_id, doc_results in document_groups.items(): + if doc_id is None: + expanded_results.extend(doc_results) + continue - has_refs = any(r.doc_item_refs for r in doc_results) - if not has_refs: - expanded_results.extend(doc_results) - continue + has_refs = any(r.doc_item_refs for r in doc_results) + if not has_refs: + expanded_results.extend(doc_results) + continue - expanded = await expand_with_items( - client.document_item_repository, - doc_id, - doc_results, - max_chars, + expanded = await expand_with_items( + client.document_item_repository, + doc_id, + doc_results, + max_chars, + ) + expanded_results.extend(expanded) + + span.set_attribute("results_out", len(expanded_results)) + span.set_attribute( + "context_chars", sum(len(r.content) for r in expanded_results) ) - expanded_results.extend(expanded) expanded_results.sort(key=lambda r: r.score, reverse=True) # image_data and picture_captions are preserved through expansion by diff --git a/haiku_rag_slim/haiku/rag/reranking/base.py b/haiku_rag_slim/haiku/rag/reranking/base.py index 703ed912..831c707c 100644 --- a/haiku_rag_slim/haiku/rag/reranking/base.py +++ b/haiku_rag_slim/haiku/rag/reranking/base.py @@ -1,5 +1,6 @@ from haiku.rag.config import Config from haiku.rag.store.models.chunk import Chunk +from haiku.rag.telemetry import logfire class RerankerBase: @@ -10,7 +11,27 @@ class RerankerBase: ) -> list[tuple[Chunk, float]]: if not chunks: return [] - return await self._rerank(query, chunks, top_n) + # Instrumented here rather than in each subclass: this is the one + # choke point every provider passes through, so the span covers + # local weights and remote APIs alike. `candidates` is the search's + # limit*10 fan-out, the main lever on how long scoring takes. + with logfire.span( + "search.rerank", + provider=type(self).__name__, + model=self._model, + candidates=len(chunks), + top_n=top_n, + ) as span: + results = await self._rerank(query, chunks, top_n) + scores = [score for _, score in results] + span.set_attribute("results", len(results)) + # Spread separates a reranker that discriminated from one whose + # scores saturated onto the same value, leaving order to the sort. + span.set_attribute("top_score", max(scores) if scores else None) + span.set_attribute( + "score_spread", max(scores) - min(scores) if scores else None + ) + return results async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 09cc3498..49fe9030 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -11,6 +11,7 @@ from lancedb.rerankers import RRFReranker from haiku.rag.store.engine import Store, query_to_pydantic from haiku.rag.store.models.chunk import Chunk, SearchType +from haiku.rag.telemetry import logfire from haiku.rag.utils import escape_sql_string logger = logging.getLogger(__name__) @@ -205,6 +206,23 @@ class ChunkRepository: await self.store.chunks_table.delete(f"document_id = '{document_id}'") return True + async def _embed_query(self, query: str) -> list[float]: + """Embed a search query, timed as its own stage. + + Instrumented here rather than in ``EmbedderWrapper.embed_query`` + because the multimodal embedders (vllm, cohere, voyageai) override + that method — the very providers whose embedding is a remote round + trip. Wrapping the call site covers every embedder uniformly. + """ + with logfire.span( + "search.embed", + provider=self.store._config.embeddings.model.provider, + model=self.store._config.embeddings.model.name, + ) as span: + embedding = await self.embedder.embed_query(query) + span.set_attribute("dim", len(embedding)) + return embedding + async def search( self, query: str = "", @@ -255,7 +273,7 @@ class ChunkRepository: .refine_factor(self.store._config.search.vector_refine_factor) ) elif search_type == "vector": - query_embedding = await self.embedder.embed_query(query) + query_embedding = await self._embed_query(query) results = ( self.store.chunks_table.query() .nearest_to(query_embedding) @@ -267,7 +285,7 @@ class ChunkRepository: query, columns="content_fts" ) else: # hybrid (default) - query_embedding = await self.embedder.embed_query(query) + query_embedding = await self._embed_query(query) reranker = RRFReranker() results = ( self.store.chunks_table.query() @@ -402,55 +420,70 @@ class ChunkRepository: else: raise ValueError("Unknown search result format, cannot extract scores") - df = await query_result.to_pandas() + # The query builder is lazy — `nearest_to`/`nearest_to_text`/`rerank` + # only describe the query. This await is where LanceDB actually runs + # the ANN and FTS passes, so it is the stage worth timing. + with logfire.span("search.execute") as span: + df = await query_result.to_pandas() + span.set_attribute("rows", len(df)) - # Extract scores - scores = extract_scores(df) + # Turning the frame into Chunks costs one more LanceDB read (the + # document metadata batch) plus a per-row json.loads — at the + # limit*10 candidate fan-out a reranked search uses, that is not + # free. Timed apart from `search.execute` so a slow search can be + # blamed on LanceDB or on this, but not ambiguously on both. + with logfire.span("search.hydrate") as span: + # Extract scores + scores = extract_scores(df) - # Convert DataFrame rows to ChunkRecords - pydantic_results = [ - self.store.ChunkRecord( - id=str(row["id"]), - document_id=str(row["document_id"]), - content=str(row["content"]), - content_fts=str(row.get("content_fts", "")), - metadata=str(row["metadata"]), - order=int(row["order"]) if "order" in row else 0, - ) - for _, row in df.iterrows() - ] + # Convert DataFrame rows to ChunkRecords + pydantic_results = [ + self.store.ChunkRecord( + id=str(row["id"]), + document_id=str(row["document_id"]), + content=str(row["content"]), + content_fts=str(row.get("content_fts", "")), + metadata=str(row["metadata"]), + order=int(row["order"]) if "order" in row else 0, + ) + for _, row in df.iterrows() + ] - # Collect all unique document IDs for batch lookup - document_ids = list(set(chunk.document_id for chunk in pydantic_results)) + # Collect all unique document IDs for batch lookup + document_ids = list(set(chunk.document_id for chunk in pydantic_results)) + span.set_attribute("documents", len(document_ids)) - # Batch fetch document metadata (skip content/docling blobs) - documents_map: dict[str, dict] = {} - if document_ids: - id_list = "', '".join(document_ids) - where_clause = f"id IN ('{id_list}')" - doc_rows = await ( - self.store.document_meta_table.query() - .select(["id", "uri", "title", "metadata"]) - .where(where_clause) - .to_list() - ) - documents_map = {str(row["id"]): row for row in doc_rows} + # Batch fetch document metadata (skip content/docling blobs) + documents_map: dict[str, dict] = {} + if document_ids: + id_list = "', '".join(document_ids) + where_clause = f"id IN ('{id_list}')" + doc_rows = await ( + self.store.document_meta_table.query() + .select(["id", "uri", "title", "metadata"]) + .where(where_clause) + .to_list() + ) + documents_map = {str(row["id"]): row for row in doc_rows} - # Build final results with document info - chunks_with_scores = [] - for i, chunk_record in enumerate(pydantic_results): - doc = documents_map.get(chunk_record.document_id) - chunk = Chunk( - id=chunk_record.id, - document_id=chunk_record.document_id, - content=chunk_record.content, - metadata=json.loads(chunk_record.metadata), - order=chunk_record.order, - document_uri=doc["uri"] if doc else None, - document_title=doc["title"] if doc else None, - document_meta=json.loads(doc.get("metadata", "{}") if doc else "{}"), - ) - score = scores[i] if i < len(scores) else 1.0 - chunks_with_scores.append((chunk, score)) + # Build final results with document info + chunks_with_scores = [] + for i, chunk_record in enumerate(pydantic_results): + doc = documents_map.get(chunk_record.document_id) + chunk = Chunk( + id=chunk_record.id, + document_id=chunk_record.document_id, + content=chunk_record.content, + metadata=json.loads(chunk_record.metadata), + order=chunk_record.order, + document_uri=doc["uri"] if doc else None, + document_title=doc["title"] if doc else None, + document_meta=json.loads( + doc.get("metadata", "{}") if doc else "{}" + ), + ) + score = scores[i] if i < len(scores) else 1.0 + chunks_with_scores.append((chunk, score)) - return chunks_with_scores + span.set_attribute("chunks", len(chunks_with_scores)) + return chunks_with_scores diff --git a/tests/test_search_telemetry.py b/tests/test_search_telemetry.py new file mode 100644 index 00000000..34cac3ed --- /dev/null +++ b/tests/test_search_telemetry.py @@ -0,0 +1,416 @@ +"""Span coverage for the search pathway. + +The pydantic-ai instrumentation already times the ``rag_search`` tool call as +a whole; these spans decompose that duration into embed / execute / hydrate / +rerank / expand / images so a slow search can be attributed to a stage, plus +the two one-off costs (store open, reranker weight load) that the run's first +tool call would otherwise absorb silently. The tests assert the span names and +the attributes queries are written against — renaming either breaks saved +Logfire views, so they are pinned here. + +Spans are captured by monkeypatching each module's ``logfire`` object, the +same approach ``tests/ingester/test_workers.py`` uses for breaker events. +""" + +from contextlib import contextmanager +from dataclasses import dataclass, field +from functools import cached_property +from types import SimpleNamespace + +import pandas as pd +import pytest + +from haiku.rag.capabilities import _base as capabilities_base +from haiku.rag.client import search as search_module +from haiku.rag.config import Config +from haiku.rag.reranking import base as reranking_base +from haiku.rag.reranking.base import RerankerBase +from haiku.rag.store.models.chunk import Chunk, SearchResult +from haiku.rag.store.repositories import chunk as chunk_repo_module +from haiku.rag.store.repositories.chunk import ChunkRepository + + +@dataclass +class RecordedSpan: + name: str + attributes: dict = field(default_factory=dict) + + def set_attribute(self, key, value): + self.attributes[key] = value + + +class SpanRecorder: + """Stand-in for the scoped ``logfire`` instance that records spans.""" + + def __init__(self): + self.spans: list[RecordedSpan] = [] + + @contextmanager + def span(self, name, **attributes): + recorded = RecordedSpan(name, dict(attributes)) + self.spans.append(recorded) + yield recorded + + def by_name(self, name: str) -> RecordedSpan: + matches = [span for span in self.spans if span.name == name] + assert matches, f"no {name!r} span; recorded {[s.name for s in self.spans]}" + assert len(matches) == 1, f"expected one {name!r} span, got {len(matches)}" + return matches[0] + + def names(self) -> list[str]: + return [span.name for span in self.spans] + + +@pytest.fixture +def recorder(monkeypatch): + """Record spans from every module on the search path.""" + recorder = SpanRecorder() + for module in ( + reranking_base, + chunk_repo_module, + search_module, + capabilities_base, + ): + monkeypatch.setattr(module, "logfire", recorder) + return recorder + + +class StubReranker(RerankerBase): + """Reranker that returns fixed scores, so the span is tested without a + provider (local weights or remote API).""" + + _model = "stub-reranker" + + def __init__(self, scores: list[float]): + self._scores = scores + + async def _rerank(self, query, chunks, top_n=10): + return list(zip(chunks, self._scores))[:top_n] + + +def _chunks(count: int) -> list[Chunk]: + return [Chunk(id=f"c{i}", content=f"chunk {i}") for i in range(count)] + + +async def test_rerank_span_records_fan_out_and_scores(recorder): + reranker = StubReranker([0.9, 0.5, 0.2]) + + await reranker.rerank("query", _chunks(3), top_n=2) + + span = recorder.by_name("search.rerank") + assert span.attributes["provider"] == "StubReranker" + assert span.attributes["model"] == "stub-reranker" + # candidates is the limit*10 fan-out the search applied, the main lever + # on rerank latency; top_n is what the caller asked for. + assert span.attributes["candidates"] == 3 + assert span.attributes["top_n"] == 2 + assert span.attributes["results"] == 2 + assert span.attributes["top_score"] == 0.9 + assert span.attributes["score_spread"] == pytest.approx(0.4) + + +async def test_rerank_span_not_emitted_for_empty_candidates(recorder): + assert await StubReranker([]).rerank("query", [], top_n=5) == [] + + assert recorder.spans == [] + + +async def test_rerank_span_tolerates_empty_results(recorder): + """A reranker that drops every candidate still closes its span; the score + attributes go None rather than raising on max() of an empty sequence.""" + await StubReranker([]).rerank("query", _chunks(2), top_n=2) + + span = recorder.by_name("search.rerank") + assert span.attributes["results"] == 0 + assert span.attributes["top_score"] is None + assert span.attributes["score_spread"] is None + + +class StubEmbedder: + async def embed_query(self, text): + return [0.1, 0.2, 0.3] + + +async def test_embed_span_records_provider_and_dim(recorder): + store = SimpleNamespace(embedder=StubEmbedder(), _config=Config) + repository = ChunkRepository(store) # type: ignore[arg-type] + + assert await repository._embed_query("query") == [0.1, 0.2, 0.3] + + span = recorder.by_name("search.embed") + assert span.attributes["provider"] == Config.embeddings.model.provider + assert span.attributes["model"] == Config.embeddings.model.name + assert span.attributes["dim"] == 3 + + +class StubQueryResult: + """LanceDB query builder stub: the awaited ``to_pandas`` is the point at + which the real builder stops being lazy and the search actually runs.""" + + def __init__(self, df: pd.DataFrame): + self._df = df + + async def to_pandas(self): + return self._df + + +class StubTableQuery: + def __init__(self, rows: list[dict]): + self._rows = rows + + def select(self, _columns): + return self + + def where(self, _clause): + return self + + async def to_list(self): + return self._rows + + +async def test_execute_span_records_row_count(recorder): + df = pd.DataFrame( + [ + { + "id": "c0", + "document_id": "d0", + "content": "hello", + "metadata": "{}", + "order": 0, + "_relevance_score": 0.75, + } + ] + ) + doc_rows = [{"id": "d0", "uri": "file://d0", "title": None, "metadata": "{}"}] + store = SimpleNamespace( + embedder=StubEmbedder(), + _config=Config, + ChunkRecord=_chunk_record_type(), + document_meta_table=SimpleNamespace(query=lambda: StubTableQuery(doc_rows)), + ) + repository = ChunkRepository(store) # type: ignore[arg-type] + + results = await repository._process_search_results(StubQueryResult(df)) # type: ignore[arg-type] + + assert len(results) == 1 + # rows is the candidate count LanceDB returned, before rerank narrows it. + assert recorder.by_name("search.execute").attributes["rows"] == 1 + # Turning that frame into Chunks is a separate stage: one more LanceDB + # read for document metadata plus a per-row json.loads. + hydrate = recorder.by_name("search.hydrate") + assert hydrate.attributes["documents"] == 1 + assert hydrate.attributes["chunks"] == 1 + # Ordering matters for reading a trace: execution precedes hydration. + assert recorder.names() == ["search.execute", "search.hydrate"] + + +def _chunk_record_type(): + """The record model Store builds per embedding dimension; only the fields + _process_search_results populates matter here.""" + from pydantic import BaseModel + + class ChunkRecord(BaseModel): + id: str + document_id: str + content: str + content_fts: str = "" + metadata: str = "{}" + order: int = 0 + + return ChunkRecord + + +async def test_expand_span_records_context_size(recorder): + """Results without a document_id pass through unexpanded, so the span is + exercised without a document_items table behind it.""" + client = SimpleNamespace(_config=Config, document_item_repository=None) + results = [ + SearchResult(content="a" * 10, score=0.9), + SearchResult(content="b" * 5, score=0.4), + ] + + expanded = await search_module.expand_context(client, results) # type: ignore[arg-type] + + assert len(expanded) == 2 + span = recorder.by_name("search.expand") + assert span.attributes["documents"] == 1 + assert span.attributes["max_chars"] == Config.search.max_context_chars + assert span.attributes["results_in"] == 2 + assert span.attributes["results_out"] == 2 + # context_chars is the payload about to be handed to the model, which ties + # this stage's cost to the next model request's prompt size. + assert span.attributes["context_chars"] == 15 + + +class StubItemRepository: + """Document-item repository serving one picture for one document.""" + + def __init__(self, picture_ref: str, blob: bytes): + self._picture_ref = picture_ref + self._blob = blob + + async def get_caption_picture_refs(self, _document_id, _refs): + return {} + + async def get_pictures_for_chunk(self, _document_id, _refs): + return {self._picture_ref: self._blob} + + async def get_text_for_refs(self, _document_id, _refs): + return {self._picture_ref: "a caption"} + + +async def test_images_span_records_documents_and_bytes(recorder): + picture_ref = "#/pictures/0" + blob = b"\x89PNG" + b"x" * 60 + client = SimpleNamespace( + _config=Config, + document_item_repository=StubItemRepository(picture_ref, blob), + ) + results = [ + SearchResult( + content="figure", + score=0.9, + document_id="d0", + doc_item_refs=[picture_ref], + ) + ] + + await search_module._populate_image_data(client, results) # type: ignore[arg-type] + + assert results[0].image_data is not None + span = recorder.by_name("search.images") + # documents is the round-trip multiplier — up to three reads each. + assert span.attributes["documents"] == 1 + assert span.attributes["pictures"] == 1 + # Raw blob bytes, not the ~4/3 larger base64 that gets attached. + assert span.attributes["bytes"] == len(blob) + + +async def test_images_span_records_zero_when_nothing_attaches(recorder): + """A result set with no pictures still opens the span, so the absence of + image work is visible rather than inferred from a missing span.""" + client = SimpleNamespace(_config=Config, document_item_repository=None) + + await search_module._populate_image_data( + client, # type: ignore[arg-type] + [SearchResult(content="text only", score=0.5)], + ) + + span = recorder.by_name("search.images") + assert span.attributes["documents"] == 0 + assert span.attributes["pictures"] == 0 + assert span.attributes["bytes"] == 0 + + +class StubRerankerClient: + """Client whose ``reranker`` is a cached_property, like ``HaikuRAG``.""" + + def __init__(self, config): + self._config = config + self.loads = 0 + + @cached_property + def reranker(self): + self.loads += 1 + return StubReranker([1.0]) + + +async def test_reranker_load_span_emitted_once_per_process(recorder): + client = StubRerankerClient(Config) + + first = search_module._get_reranker(client) # type: ignore[arg-type] + second = search_module._get_reranker(client) # type: ignore[arg-type] + + # The cached_property loaded weights exactly once; the span marks that + # cold touch, so a search.reranker.load in a trace means a cold process. + assert first is second + assert client.loads == 1 + assert recorder.names() == ["search.reranker.load"] + + +class StubRag: + """Stands in for HaikuRAG so _ensure_rag opens no real store.""" + + opened = 0 + + def __init__(self, *_args, **_kwargs): + pass + + async def __aenter__(self): + type(self).opened += 1 + return self + + async def __aexit__(self, *_exc): + return False + + +@pytest.fixture +def capability(tmp_path, monkeypatch): + from haiku.rag.capabilities.rag import RAGState, create_capability + + monkeypatch.setattr(capabilities_base, "HaikuRAG", StubRag) + StubRag.opened = 0 + built = create_capability(db_path=tmp_path / "test.lancedb", config=Config) + built.state = RAGState() + return built + + +async def test_client_open_span_emitted_once_per_run(recorder, capability): + first = await capability._ensure_rag() + second = await capability._ensure_rag() + + # Store open is lazy and would otherwise be charged to whichever search + # ran first; the double-checked lock means only one span per run. + assert first is second + assert StubRag.opened == 1 + span = recorder.by_name("rag.client.open") + assert span.attributes["db_path"] == str(capability.db_path) + + +async def test_tool_search_span_records_position_in_run( + recorder, capability, monkeypatch +): + results = [SearchResult(content="evidence", score=0.9, chunk_id="c0")] + + async def _search_corpus(_rag, _query, limit=None, document_filter=None): + return "formatted evidence", results + + monkeypatch.setattr(capabilities_base, "search_corpus", _search_corpus) + + await capability._search("first query", limit=3) + await capability._search("second query", limit=3) + + searches = [s for s in recorder.spans if s.name == "ask.tool.search"] + assert [s.attributes["search_index"] for s in searches] == [1, 2] + assert searches[0].attributes["namespace"] == "rag" + assert searches[0].attributes["max_searches"] == Config.qa.max_searches + assert searches[0].attributes["limit"] == 3 + assert searches[0].attributes["results"] == 1 + assert searches[0].attributes["formatted_chars"] == len("formatted evidence") + # Spans are recorded as they open, so this order is the nesting: the + # one-off store open happens inside the FIRST search and never again, + # which is what makes search #1 legitimately slower than the rest. + assert recorder.names() == [ + "ask.tool.search", + "rag.client.open", + "ask.tool.search", + ] + + +async def test_tool_search_span_not_emitted_when_budget_spent( + recorder, capability, monkeypatch +): + """A refused call emits no span: the ToolFailed is already on the tool + span pydantic-ai opened, and search.* spans should mean work happened.""" + from pydantic_ai import ToolFailed + + async def _search_corpus(_rag, _query, limit=None, document_filter=None): + return "", [] + + monkeypatch.setattr(capabilities_base, "search_corpus", _search_corpus) + capability.search_count = capability._max_searches + + with pytest.raises(ToolFailed): + await capability._search("one too many", limit=3) + + assert "ask.tool.search" not in recorder.names()