longprobe 0.1.0__py3-none-any.whl

longprobe/__init__.py

@@ -0,0 +1,122 @@
+"""
+LongProbe — RAG retrieval regression testing harness.
+
+Define Golden Questions once. Run ``longprobe check`` on every commit.
+Get an exact diff of which document chunks were lost — before your users notice.
+"""
+
+from longprobe.core.golden import GoldenQuestion, GoldenSet
+from longprobe.core.scorer import RecallScorer, QuestionResult, ProbeReport
+from longprobe.core.embedder import QueryEmbedder
+from longprobe.core.baseline import BaselineStore
+from longprobe.core.diff import DiffReporter, RegressionDiff
+from longprobe.config import ProbeConfig
+from longprobe.adapters.base import AbstractRetrieverAdapter
+from longprobe.adapters import create_adapter
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "GoldenQuestion",
+    "GoldenSet",
+    "RecallScorer",
+    "QuestionResult",
+    "ProbeReport",
+    "QueryEmbedder",
+    "BaselineStore",
+    "DiffReporter",
+    "RegressionDiff",
+    "ProbeConfig",
+    "AbstractRetrieverAdapter",
+    "create_adapter",
+]
+
+
+class LongProbe:
+    """
+    High-level facade for running RAG regression probes.
+
+    Usage::
+
+        from longprobe import LongProbe, ChromaAdapter
+
+        adapter = ChromaAdapter(collection_name="my_docs", persist_directory="./db")
+        probe = LongProbe(adapter=adapter, goldens_path="goldens.yaml")
+        report = probe.run()
+        print(f"Overall recall: {report.overall_recall:.2f}")
+    """
+
+    def __init__(
+        self,
+        adapter: "AbstractRetrieverAdapter",
+        goldens_path: str = "goldens.yaml",
+        config_path: str = "longprobe.yaml",
+        recall_threshold: float | None = None,
+    ):
+        from pathlib import Path
+
+        self.adapter = adapter
+        self.golden_set = GoldenSet.from_yaml(goldens_path)
+
+        config_file = Path(config_path)
+        if config_file.exists():
+            self.config = ProbeConfig.from_yaml(str(config_file))
+        else:
+            self.config = ProbeConfig.defaults()
+
+        if recall_threshold is not None:
+            self.config.scoring.recall_threshold = recall_threshold
+
+        self.scorer = RecallScorer(
+            recall_threshold=self.config.scoring.recall_threshold
+        )
+        self.baseline_store = BaselineStore(
+            db_path=self.config.baseline.db_path
+        )
+        self.diff_reporter = DiffReporter()
+        self._last_report: ProbeReport | None = None
+
+    def run(self, top_k_override: int | None = None) -> ProbeReport:
+        """Run the probe against the golden set and return a report."""
+        report = self.scorer.score_all(
+            self.golden_set,
+            self.adapter.retrieve,
+            top_k_override=top_k_override,
+        )
+
+        # Auto-compare against baseline
+        if self.config.baseline.auto_compare:
+            baseline = self.baseline_store.load("latest")
+            if baseline is not None:
+                report.baseline_recall = baseline.overall_recall
+                report.recall_delta = report.overall_recall - baseline.overall_recall
+                diff_result = self.baseline_store.diff(report, baseline)
+                report.regression_detected = len(diff_result["regressions"]) > 0
+
+        self._last_report = report
+        return report
+
+    def save_baseline(self, label: str = "latest") -> None:
+        """Save the last run report as a baseline."""
+        if self._last_report is None:
+            raise RuntimeError("No report to save. Run probe.run() first.")
+        self.baseline_store.save(self._last_report, label=label)
+
+    def diff(self, baseline_label: str = "latest") -> dict:
+        """Compare last run against a saved baseline."""
+        if self._last_report is None:
+            raise RuntimeError("No report to diff. Run probe.run() first.")
+        baseline = self.baseline_store.load(baseline_label)
+        if baseline is None:
+            raise ValueError(f"Baseline '{baseline_label}' not found.")
+        return self.baseline_store.diff(self._last_report, baseline)
+
+    def get_missing_chunks(self) -> dict[str, list[str]]:
+        """Return a dict mapping question_id to its missing chunks."""
+        if self._last_report is None:
+            raise RuntimeError("No report available. Run probe.run() first.")
+        return {
+            r.question_id: r.missing_chunks
+            for r in self._last_report.results
+            if r.missing_chunks
+        }

longprobe/adapters/__init__.py

@@ -0,0 +1,56 @@
+from .base import AbstractRetrieverAdapter
+from .http import HttpAdapter
+from .langchain import LangChainRetrieverAdapter
+from .llamaindex import LlamaIndexRetrieverAdapter
+from .pinecone import PineconeAdapter
+from .qdrant import QdrantAdapter
+from .chroma import ChromaAdapter
+
+
+def create_adapter(adapter_type: str, **kwargs) -> AbstractRetrieverAdapter:
+    """
+    Factory function to create a retriever adapter.
+
+    Args:
+        adapter_type: One of "langchain", "llamaindex", "pinecone", "qdrant",
+            "chroma", "http"
+        **kwargs: Passed to the adapter constructor.
+                  For "langchain" and "llamaindex": pass "retriever" kwarg
+                  For "http": pass "config" kwarg (HttpRetrieverConfig)
+                  For others: pass connection config kwargs
+
+    Returns:
+        AbstractRetrieverAdapter instance
+
+    Raises:
+        ValueError: If adapter_type is unknown
+    """
+    adapters = {
+        "langchain": LangChainRetrieverAdapter,
+        "llamaindex": LlamaIndexRetrieverAdapter,
+        "pinecone": PineconeAdapter,
+        "qdrant": QdrantAdapter,
+        "chroma": ChromaAdapter,
+        "http": HttpAdapter,
+    }
+
+    adapter_cls = adapters.get(adapter_type)
+    if adapter_cls is None:
+        raise ValueError(
+            f"Unknown adapter type '{adapter_type}'. "
+            f"Must be one of: {', '.join(adapters.keys())}"
+        )
+
+    return adapter_cls(**kwargs)
+
+
+__all__ = [
+    "AbstractRetrieverAdapter",
+    "HttpAdapter",
+    "LangChainRetrieverAdapter",
+    "LlamaIndexRetrieverAdapter",
+    "PineconeAdapter",
+    "QdrantAdapter",
+    "ChromaAdapter",
+    "create_adapter",
+]

longprobe/adapters/base.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class AbstractRetrieverAdapter(ABC):
+    """Base class for all retriever adapters."""
+
+    @abstractmethod
+    def retrieve(self, query: str, top_k: int = 5) -> list[dict[str, Any]]:
+        """
+        Retrieve documents for a query.
+
+        Returns list of dicts with keys:
+        - "id": str - document/chunk identifier
+        - "text": str - document text content
+        - "score": float - relevance score
+        - "metadata": dict - additional metadata
+        """
+        ...
+
+    def health_check(self) -> bool:
+        """Check if the retriever is reachable. Override for custom health checks."""
+        return True

longprobe/adapters/chroma.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from .base import AbstractRetrieverAdapter
+
+logger = logging.getLogger(__name__)
+
+
+class ChromaAdapter(AbstractRetrieverAdapter):
+    """Adapter for Chroma vector databases.
+
+    Supports both text-based and embedding-based queries.  The Chroma
+    client SDK is imported lazily and cached after the first access.
+    """
+
+    def __init__(
+        self,
+        collection_name: str = "default",
+        persist_directory: str = "",
+    ) -> None:
+        self.collection_name = collection_name
+        self.persist_directory = persist_directory
+        self._client: Any = None
+
+    def _get_client(self) -> Any:
+        """Lazily create and cache a Chroma client."""
+        if self._client is not None:
+            return self._client
+
+        try:
+            import chromadb
+        except ImportError:
+            raise ImportError(
+                "chromadb is required for ChromaAdapter. "
+                "Install it with:  pip install chromadb"
+            )
+
+        if self.persist_directory:
+            self._client = chromadb.PersistentClient(path=self.persist_directory)
+        else:
+            self._client = chromadb.Client()
+
+        return self._client
+
+    def retrieve(
+        self,
+        query: str | list[float],
+        top_k: int = 5,
+    ) -> list[dict[str, Any]]:
+        """Query a Chroma collection.
+
+        Args:
+            query: Either a text string or a pre-computed embedding vector.
+            top_k: Number of results to return.
+
+        Returns:
+            List of result dicts normalised to the LongProbe format.
+            Scores are converted from Chroma distances (``1 - distance``).
+        """
+        client = self._get_client()
+        collection = client.get_collection(self.collection_name)
+
+        if isinstance(query, str):
+            results = collection.query(
+                query_texts=[query],
+                n_results=top_k,
+                include=["documents", "metadatas", "distances"],
+            )
+        elif isinstance(query, list):
+            results = collection.query(
+                query_embeddings=[query],
+                n_results=top_k,
+                include=["documents", "metadatas", "distances"],
+            )
+        else:
+            raise TypeError(
+                f"Query must be a str or list[float], got {type(query).__name__}"
+            )
+
+        ids = results.get("ids", [[]])[0]
+        documents = results.get("documents", [[]])[0]
+        distances = results.get("distances", [[]])[0]
+        metadatas = results.get("metadatas", [[]])[0]
+
+        mapped: list[dict[str, Any]] = []
+        for i in range(len(ids)):
+            mapped.append(
+                {
+                    "id": ids[i],
+                    "text": documents[i] if documents and i < len(documents) else "",
+                    "score": 1.0 - (distances[i] if distances and i < len(distances) else 0.0),
+                    "metadata": (
+                        metadatas[i]
+                        if metadatas and i < len(metadatas)
+                        else {}
+                    ),
+                }
+            )
+
+        return mapped
+
+    def health_check(self) -> bool:
+        """Check connectivity by listing collections."""
+        try:
+            client = self._get_client()
+            client.list_collections()
+            return True
+        except Exception:
+            logger.debug("Chroma health check failed", exc_info=True)
+            return False

longprobe/adapters/http.py

@@ -0,0 +1,234 @@
+"""Generic HTTP retriever adapter for LongProbe.
+
+Sends queries to any HTTP-based RAG API endpoint and maps the JSON
+response back to LongProbe's standard format.  This adapter lets LongProbe
+test the *entire* retrieval pipeline (including any reranking, prompt
+rewriting, or filtering) rather than just the raw vector database.
+
+Typical configuration in ``longprobe.yaml``::
+
+    retriever:
+      type: http
+      http:
+        url: "http://localhost:8000/api/retrieve"
+        method: POST
+        body_template: '{"query": "{question}", "top_k": {top_k}}'
+        headers:
+          Authorization: "Bearer ${API_KEY}"
+        response_mapping:
+          results_path: "data.chunks"
+          id_field: "chunk_id"
+          text_field: "content"
+          score_field: "similarity"
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+import requests
+
+from .base import AbstractRetrieverAdapter
+from ..config import HttpRetrieverConfig
+
+logger = logging.getLogger(__name__)
+
+
+class HttpAdapter(AbstractRetrieverAdapter):
+    """Adapter that queries an HTTP-based RAG API endpoint.
+
+    The request body is built from ``body_template`` by replacing
+    ``{question}`` and ``{top_k}`` placeholders.  The JSON response is
+    parsed using ``response_mapping`` to locate the results array and
+    extract per-result fields.
+
+    Args:
+        config: HTTP-specific configuration (URL, method, headers, etc.).
+    """
+
+    def __init__(self, config: HttpRetrieverConfig) -> None:
+        self._config = config
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def retrieve(self, query: str, top_k: int = 5) -> list[dict[str, Any]]:
+        """Send *query* to the configured HTTP endpoint and return normalised
+        results.
+
+        Args:
+            query: Natural-language question to send.
+            top_k: Number of results to request.
+
+        Returns:
+            List of dicts with ``id``, ``text``, ``score``, and ``metadata``
+            keys.
+
+        Raises:
+            RuntimeError: If the endpoint returns a non-2xx status.
+            ValueError: If the body template produces invalid JSON after
+                placeholder substitution.
+            requests.Timeout: If the request exceeds the configured timeout.
+        """
+        cfg = self._config
+
+        # 1. Build request body from template.
+        body = self._build_body(cfg.body_template, query, top_k)
+
+        # 2. Make HTTP request.
+        try:
+            response = requests.request(
+                method=cfg.method,
+                url=cfg.url,
+                json=body,
+                headers=cfg.headers,
+                timeout=cfg.timeout,
+            )
+        except requests.Timeout as exc:
+            raise requests.Timeout(
+                f"HTTP adapter request timed out after {cfg.timeout}s "
+                f"for URL '{cfg.url}'."
+            ) from exc
+
+        # 3. Check status.
+        if response.status_code >= 300:
+            snippet = response.text[:500] if response.text else "(empty body)"
+            raise RuntimeError(
+                f"HTTP adapter received status {response.status_code} from "
+                f"'{cfg.url}': {snippet}"
+            )
+
+        # 4. Parse response.
+        try:
+            data = response.json()
+        except (json.JSONDecodeError, ValueError) as exc:
+            raise RuntimeError(
+                f"HTTP adapter received non-JSON response from '{cfg.url}': "
+                f"{response.text[:200]}"
+            ) from exc
+
+        # 5. Resolve results path.
+        mapping = cfg.response_mapping
+        items = self._resolve_path(data, mapping.results_path)
+
+        if not isinstance(items, list):
+            logger.warning(
+                "HTTP adapter: path '%s' did not resolve to a list "
+                "(got %s). Returning empty results.",
+                mapping.results_path,
+                type(items).__name__,
+            )
+            return []
+
+        # 6. Map each item to LongProbe format.
+        results: list[dict[str, Any]] = []
+        for item in items:
+            if not isinstance(item, dict):
+                logger.warning(
+                    "HTTP adapter: skipping non-dict item in results array."
+                )
+                continue
+            results.append(
+                {
+                    "id": str(item.get(mapping.id_field, "")),
+                    "text": str(item.get(mapping.text_field, "")),
+                    "score": float(item.get(mapping.score_field, 0.0)),
+                    "metadata": {
+                        k: v
+                        for k, v in item.items()
+                        if k not in (mapping.id_field, mapping.text_field, mapping.score_field)
+                    },
+                }
+            )
+
+        return results
+
+    def health_check(self) -> bool:
+        """Check if the configured endpoint is reachable.
+
+        Sends a HEAD (or GET for endpoints that don't support HEAD) to the
+        URL and returns ``True`` for any response with status < 400.
+        Timeouts and connection errors return ``False``.
+        """
+        try:
+            # Try HEAD first (lighter).
+            resp = requests.head(
+                self._config.url,
+                headers=self._config.headers,
+                timeout=min(self._config.timeout, 10),
+            )
+            if resp.status_code == 405:
+                # Fallback to GET if HEAD is not allowed.
+                resp = requests.get(
+                    self._config.url,
+                    headers=self._config.headers,
+                    timeout=min(self._config.timeout, 10),
+                )
+            return resp.status_code < 400
+        except (requests.RequestException, OSError):
+            logger.debug("HTTP adapter health check failed", exc_info=True)
+            return False
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _build_body(
+        template: str,
+        query: str,
+        top_k: int,
+    ) -> Any:
+        """Substitute ``{question}`` and ``{top_k}`` in *template* and
+        parse the result as JSON.
+
+        ``{question}`` is replaced with a JSON-safe escaped string.
+        ``{top_k}`` is replaced with the integer value.
+
+        Returns:
+            The parsed JSON object (usually a dict).
+
+        Raises:
+            ValueError: If the substituted template is not valid JSON.
+        """
+        # JSON-escape the question string, then strip surrounding quotes
+        # so it fits naturally inside the template's own quoting.
+        escaped_question = json.dumps(query)[1:-1]
+        body_str = template.replace("{question}", escaped_question)
+        body_str = body_str.replace("{top_k}", str(top_k))
+
+        try:
+            return json.loads(body_str)
+        except json.JSONDecodeError as exc:
+            raise ValueError(
+                f"HTTP adapter body template produced invalid JSON after "
+                f"substitution: {exc}\n"
+                f"Template: {template}\n"
+                f"Result:   {body_str[:200]}"
+            ) from exc
+
+    @staticmethod
+    def _resolve_path(data: Any, path: str) -> Any:
+        """Traverse *data* using a dot-notation *path*.
+
+        ``\"data.chunks\"`` resolves to ``data[\"data\"][\"chunks\"]``.
+        If any segment is missing, an empty list is returned.
+
+        Args:
+            data: The parsed JSON response.
+            path: Dot-separated key path.
+
+        Returns:
+            The value at the path, or an empty list if not found.
+        """
+        current = data
+        parts = path.split(".") if path else []
+        for part in parts:
+            if isinstance(current, dict) and part in current:
+                current = current[part]
+            else:
+                return []
+        return current

longprobe/adapters/langchain.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from .base import AbstractRetrieverAdapter
+
+logger = logging.getLogger(__name__)
+
+
+class LangChainRetrieverAdapter(AbstractRetrieverAdapter):
+    """Adapter for LangChain BaseRetriever objects.
+
+    Accepts any LangChain retriever (duck-typed) and normalises its
+    output into the LongProbe result format.
+    """
+
+    def __init__(self, retriever: Any) -> None:
+        self.retriever = retriever
+
+    def retrieve(self, query: str, top_k: int = 5) -> list[dict[str, Any]]:
+        """Retrieve documents using a LangChain retriever.
+
+        Tries the modern ``invoke`` API first, then falls back to the
+        legacy ``get_relevant_documents`` method.
+        """
+        try:
+            from langchain_core.documents import Document  # noqa: F401 – ensure importable
+        except ImportError:
+            raise ImportError(
+                "langchain-core is required for LangChainRetrieverAdapter. "
+                "Install it with:  pip install langchain-core"
+            )
+
+        # Invoke the retriever
+        documents = []
+        try:
+            documents = self.retriever.invoke(query)
+        except AttributeError:
+            # Legacy LangChain API
+            documents = self.retriever.get_relevant_documents(query)
+
+        results: list[dict[str, Any]] = []
+        for i, d in enumerate(documents):
+            doc_id = d.metadata.get("chunk_id", d.metadata.get("source", f"doc_{i}"))
+            results.append(
+                {
+                    "id": doc_id,
+                    "text": d.page_content,
+                    "score": float(d.metadata.get("score", 0.0)),
+                    "metadata": dict(d.metadata),
+                }
+            )
+
+        return results[:top_k]
+
+    def health_check(self) -> bool:
+        """Check that the underlying retriever is callable."""
+        try:
+            return (
+                hasattr(self.retriever, "invoke")
+                or hasattr(self.retriever, "get_relevant_documents")
+            )
+        except Exception:
+            return False

longprobe/adapters/llamaindex.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from .base import AbstractRetrieverAdapter
+
+logger = logging.getLogger(__name__)
+
+
+class LlamaIndexRetrieverAdapter(AbstractRetrieverAdapter):
+    """Adapter for LlamaIndex retriever objects.
+
+    Accepts any LlamaIndex retriever and normalises its
+    ``NodeWithScore`` results into the LongProbe format.
+    """
+
+    def __init__(self, retriever: Any) -> None:
+        self.retriever = retriever
+
+    def retrieve(self, query: str, top_k: int = 5) -> list[dict[str, Any]]:
+        """Retrieve documents using a LlamaIndex retriever."""
+        try:
+            from llama_index.core import QueryBundle  # noqa: F401
+        except ImportError:
+            raise ImportError(
+                "llama-index-core is required for LlamaIndexRetrieverAdapter. "
+                "Install it with:  pip install llama-index-core"
+            )
+
+        nodes_with_scores = self.retriever.retrieve(query)
+
+        results: list[dict[str, Any]] = []
+        for node in nodes_with_scores:
+            results.append(
+                {
+                    "id": node.node.node_id,
+                    "text": node.node.get_content(),
+                    "score": float(node.score or 0.0),
+                    "metadata": dict(node.node.metadata),
+                }
+            )
+
+        return results[:top_k]
+
+    def health_check(self) -> bool:
+        """Check that the underlying retriever has a ``retrieve`` method."""
+        return hasattr(self.retriever, "retrieve")