sci-rag-kit 0.3.0__py3-none-any.whl

sci_rag/__init__.py

@@ -0,0 +1,78 @@
+"""sci-rag-kit: a DIY GraphRAG factory for scientific domains.
+
+Point it at a folder of papers and reports, and it gives you a grounded,
+citation-backed question-answering system for your field: ingestion with
+structure-aware chunking, a knowledge graph stored natively in Postgres,
+five fused retrieval layers, an honest evaluation harness, and serving
+over REST and MCP.
+
+The pieces are importable directly for notebook and library use::
+
+    from sci_rag import Retriever, AnswerEngine, RetrievalScope
+
+    retriever = Retriever()
+    result = await retriever.retrieve("rice straw availability", profile="deep")
+
+Exports resolve lazily (PEP 562), so ``import sci_rag`` stays cheap and
+circular imports are impossible.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _installed_version
+from typing import Any
+
+try:
+    #: Read from the installed distribution rather than declared here.
+    #: A hand-maintained copy drifts, and this one had: v0.2.0 shipped
+    #: reporting 0.1.0a0. `sci-rag-new` resolves the template tag from this
+    #: same metadata, so the two cannot disagree about what version this is.
+    __version__ = _installed_version("sci-rag-kit")
+except PackageNotFoundError:  # pragma: no cover - only in a bare source tree
+    __version__ = "0.0.0+unknown"
+
+_EXPORTS: dict[str, str] = {
+    # configuration and domain
+    "Settings": "sci_rag.config",
+    "get_settings": "sci_rag.config",
+    "DomainProfile": "sci_rag.domain",
+    "load_domain": "sci_rag.domain",
+    # ingestion
+    "CorpusEntry": "sci_rag.ingest",
+    "discover_folder": "sci_rag.ingest",
+    "load_manifest": "sci_rag.ingest",
+    "ingest_entries": "sci_rag.ingest",
+    # providers
+    "get_embedder": "sci_rag.embed",
+    "get_llm": "sci_rag.llm",
+    # graph
+    "extract_graph": "sci_rag.graph",
+    "build_communities": "sci_rag.graph",
+    # retrieval and answering
+    "Retriever": "sci_rag.retrieve",
+    "RetrievalResult": "sci_rag.retrieve",
+    "RetrievalScope": "sci_rag.retrieve",
+    "AnswerEngine": "sci_rag.answer",
+    # evaluation
+    "load_seed_questions": "sci_rag.evals",
+    "run_retrieval_eval": "sci_rag.evals",
+    "run_answer_eval": "sci_rag.evals",
+    # serving
+    "create_app": "sci_rag.server",
+    "build_mcp_server": "sci_rag.server",
+}
+
+__all__ = ["__version__", *sorted(_EXPORTS)]
+
+
+def __getattr__(name: str) -> Any:
+    module_path = _EXPORTS.get(name)
+    if module_path is None:
+        raise AttributeError(f"module 'sci_rag' has no attribute {name!r}")
+    return getattr(import_module(module_path), name)
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_EXPORTS))

sci_rag/answer/__init__.py

@@ -0,0 +1,3 @@
+from sci_rag.answer.generator import AnswerEngine, AnswerEvent, AnswerResult, SourceCitation
+
+__all__ = ["AnswerEngine", "AnswerEvent", "AnswerResult", "SourceCitation"]

sci_rag/answer/compress.py

@@ -0,0 +1,132 @@
+"""Question-aware compression of retrieved chunks before answer generation.
+
+The model output is an optimization hint, never new evidence. Valid summaries
+retain the original chunk identity and citation metadata. Missing, malformed,
+failed, or over-budget summaries fall back to the complete retrieved text so a
+provider problem cannot silently erase evidence.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, replace
+from typing import Any
+
+import structlog
+from pydantic import BaseModel, Field, ValidationError
+
+from sci_rag.domain import DomainProfile
+from sci_rag.ingest.tokens import count_tokens
+from sci_rag.llm import LLMClient
+from sci_rag.retrieve import RetrievedItem
+
+log = structlog.get_logger(__name__)
+
+
+class _Snippet(BaseModel):
+    index: int = Field(ge=1, strict=True)
+    relevance_score: float = Field(ge=0.0, le=1.0, strict=True)
+    summary: str = Field(min_length=1, strict=True)
+
+
+@dataclass
+class CompressionResult:
+    items: list[RetrievedItem]
+    failure_count: int = 0
+    dropped_count: int = 0
+
+
+class SnippetCompressor:
+    """Compress chunk text in one batched JSON call.
+
+    Community summaries already are compressed, cross-document evidence and
+    therefore pass through untouched. The compressor only has authority over
+    ordinary chunks returned by retrieval.
+    """
+
+    def __init__(self, domain: DomainProfile, llm: LLMClient) -> None:
+        self.domain = domain
+        self.llm = llm
+
+    async def compress(
+        self,
+        query: str,
+        items: list[RetrievedItem],
+        *,
+        relevance_floor: float,
+        max_tokens_per_chunk: int,
+    ) -> CompressionResult:
+        chunks = [
+            (index, item) for index, item in enumerate(items, start=1) if item.kind == "chunk"
+        ]
+        if not chunks:
+            return CompressionResult(items=list(items))
+
+        try:
+            prompt = self.domain.render_prompt(
+                "snippet_compression",
+                QUERY=query,
+                MAX_TOKENS_PER_CHUNK=str(max_tokens_per_chunk),
+                CHUNKS_JSON=json.dumps(
+                    [{"index": index, "text": item.content} for index, item in chunks],
+                    ensure_ascii=False,
+                ),
+            )
+            payload = await self.llm.generate_json(
+                prompt,
+                max_tokens=min(8192, max(512, len(chunks) * max_tokens_per_chunk + 256)),
+            )
+        except Exception as exc:
+            log.warning(
+                "snippet_compression_failed",
+                error=type(exc).__name__,
+                chunk_count=len(chunks),
+            )
+            return CompressionResult(items=list(items), failure_count=len(chunks))
+
+        snippets = _validated_snippets(payload)
+        if snippets is None:
+            return CompressionResult(items=list(items), failure_count=len(chunks))
+
+        by_index: dict[int, _Snippet] = {}
+        duplicate_indexes: set[int] = set()
+        for snippet in snippets:
+            if snippet.index in by_index:
+                duplicate_indexes.add(snippet.index)
+            by_index[snippet.index] = snippet
+
+        output: list[RetrievedItem] = []
+        failures = 0
+        dropped = 0
+        chunk_indexes = {index for index, _item in chunks}
+        for index, item in enumerate(items, start=1):
+            if item.kind != "chunk":
+                output.append(item)
+                continue
+            selected = by_index.get(index)
+            if (
+                selected is None
+                or index in duplicate_indexes
+                or not selected.summary.strip()
+                or count_tokens(selected.summary.strip()) > max_tokens_per_chunk
+            ):
+                failures += 1
+                output.append(item)
+            elif selected.relevance_score < relevance_floor:
+                dropped += 1
+            else:
+                output.append(replace(item, content=selected.summary.strip()))
+
+        unexpected = set(by_index) - chunk_indexes
+        if unexpected:
+            log.warning("snippet_compression_ignored_indexes", count=len(unexpected))
+        return CompressionResult(items=output, failure_count=failures, dropped_count=dropped)
+
+
+def _validated_snippets(payload: Any) -> list[_Snippet] | None:
+    if not isinstance(payload, dict) or not isinstance(payload.get("snippets"), list):
+        return None
+    try:
+        return [_Snippet.model_validate(value) for value in payload["snippets"]]
+    except ValidationError:
+        return None

sci_rag/answer/generator.py

@@ -0,0 +1,341 @@
+"""Grounded answer generation with inline citations.
+
+The contract with users: every claim in an answer is backed by a numbered
+source the reader can check, and when the corpus does not contain the
+answer, the answer says so instead of improvising. The prompt enforces it,
+and the evaluation harness's blind judge checks it after the fact.
+
+Both entry points share the same preparation (retrieve, build the numbered
+source block, render the domain's answer prompt):
+
+* :meth:`AnswerEngine.answer` returns the complete result.
+* :meth:`AnswerEngine.answer_stream` yields typed events (retrieval
+  progress, text deltas, citations, done) that the REST SSE endpoint and
+  MCP tools translate directly.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+from typing import Any
+
+import structlog
+
+from sci_rag.config import Settings, get_settings
+from sci_rag.domain import DomainProfile
+from sci_rag.ingest.tokens import count_tokens
+from sci_rag.llm import LLMClient, get_llm
+from sci_rag.retrieve import RetrievalResult, RetrievalScope, Retriever
+
+from .compress import SnippetCompressor
+
+log = structlog.get_logger(__name__)
+
+_CITATION_RE = re.compile(r"\[(\d+)\]")
+
+
+@dataclass
+class SourceCitation:
+    index: int
+    kind: str  # "chunk" | "community"
+    title: str
+    citation: str | None
+    license_class: str
+    document_id: str | None
+    chunk_id: str | None
+    section_path: str | None
+    cited: bool = False
+
+
+@dataclass
+class AnswerEvent:
+    type: str  # retrieval_started | retrieval_done | compression_done | generation_started | delta | citations | done | error
+    data: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class AnswerResult:
+    text: str
+    sources: list[SourceCitation]
+    retrieval: RetrievalResult
+    model: str
+    prompt_retrieval: RetrievalResult | None = None
+    prompt_tokens_before: int = 0
+    prompt_tokens_after: int = 0
+    compression_failure_count: int = 0
+    compression_dropped_count: int = 0
+
+    @property
+    def cited_sources(self) -> list[SourceCitation]:
+        return [s for s in self.sources if s.cited]
+
+
+class AnswerEngine:
+    def __init__(
+        self,
+        *,
+        settings: Settings | None = None,
+        retriever: Retriever | None = None,
+        llm: LLMClient | None = None,
+    ) -> None:
+        self.settings = settings or get_settings()
+        self.retriever = retriever or Retriever(settings=self.settings)
+        self.domain: DomainProfile = self.retriever.domain
+        self._llm = llm
+
+    def _resolve_llm(self, api_key_override: str | None = None) -> LLMClient:
+        if api_key_override:
+            # Bring-your-own-key: a per-request client that is never stored.
+            return get_llm(self.settings, api_key_override=api_key_override)
+        if self._llm is None:
+            self._llm = get_llm(self.settings)
+        return self._llm
+
+    async def answer(
+        self,
+        query: str,
+        *,
+        profile: str = "deep",
+        limit: int = 8,
+        scope: RetrievalScope | None = None,
+        api_key_override: str | None = None,
+        max_tokens: int = 2048,
+        include_compression: bool | None = None,
+    ) -> AnswerResult:
+        events = self.answer_stream(
+            query,
+            profile=profile,
+            limit=limit,
+            scope=scope,
+            api_key_override=api_key_override,
+            max_tokens=max_tokens,
+            include_compression=include_compression,
+        )
+        text_parts: list[str] = []
+        sources: list[SourceCitation] = []
+        retrieval: RetrievalResult | None = None
+        prompt_retrieval: RetrievalResult | None = None
+        model = str(self.settings.model_spec_for("answer"))
+        prompt_tokens_before = 0
+        prompt_tokens_after = 0
+        compression_failure_count = 0
+        compression_dropped_count = 0
+        async for event in events:
+            if event.type == "delta":
+                text_parts.append(event.data["text"])
+            elif event.type == "retrieval_done":
+                retrieval = event.data["_result"]
+            elif event.type == "generation_started":
+                model = event.data["model"]
+            elif event.type == "compression_done":
+                prompt_retrieval = event.data["_result"]
+                prompt_tokens_before = event.data["prompt_tokens_before"]
+                prompt_tokens_after = event.data["prompt_tokens_after"]
+                compression_failure_count = event.data["failure_count"]
+                compression_dropped_count = event.data["dropped_count"]
+            elif event.type == "citations":
+                sources = event.data["_sources"]
+            elif event.type == "error":
+                raise RuntimeError(event.data.get("message", "answer generation failed"))
+        assert retrieval is not None
+        if prompt_retrieval is None:
+            prompt_retrieval = retrieval
+        return AnswerResult(
+            text="".join(text_parts),
+            sources=sources,
+            retrieval=retrieval,
+            model=model,
+            prompt_retrieval=prompt_retrieval,
+            prompt_tokens_before=prompt_tokens_before,
+            prompt_tokens_after=prompt_tokens_after,
+            compression_failure_count=compression_failure_count,
+            compression_dropped_count=compression_dropped_count,
+        )
+
+    async def answer_stream(
+        self,
+        query: str,
+        *,
+        profile: str = "deep",
+        limit: int = 8,
+        scope: RetrievalScope | None = None,
+        api_key_override: str | None = None,
+        max_tokens: int = 2048,
+        include_compression: bool | None = None,
+    ) -> AsyncIterator[AnswerEvent]:
+        yield AnswerEvent(type="retrieval_started", data={"profile": profile})
+        if scope is None:
+            scope = RetrievalScope(exclude_retracted=True)
+        retrieval = await self.retriever.retrieve(query, profile=profile, limit=limit, scope=scope)
+        yield AnswerEvent(
+            type="retrieval_done",
+            data={
+                "item_count": len(retrieval.items),
+                "degraded_stages": retrieval.degraded_stages,
+                "traces": [
+                    {
+                        "stage": t.stage,
+                        "status": t.status,
+                        "duration_ms": t.duration_ms,
+                        "candidates": t.candidate_count,
+                    }
+                    for t in retrieval.traces
+                ],
+                "_result": retrieval,
+            },
+        )
+
+        if not retrieval.items:
+            text = (
+                "The knowledge base has no material matching this question within "
+                "the allowed scope, so I cannot give a grounded answer."
+            )
+            yield AnswerEvent(type="delta", data={"text": text})
+            yield AnswerEvent(type="citations", data={"citations": [], "_sources": []})
+            yield AnswerEvent(type="done", data={"finish_reason": "no_sources"})
+            return
+
+        try:
+            llm = self._resolve_llm(api_key_override)
+        except Exception as exc:
+            yield AnswerEvent(type="error", data={"code": "llm_unavailable", "message": str(exc)})
+            return
+
+        raw_prompt = self._render_answer_prompt(query, retrieval)
+        prompt_retrieval = retrieval
+        failures = 0
+        dropped = 0
+        compression_on = (
+            include_compression
+            if include_compression is not None
+            else self.domain.config.compression.enabled
+        )
+        if compression_on:
+            tuning = self.domain.config.compression
+            compressed = await SnippetCompressor(self.domain, llm).compress(
+                query,
+                retrieval.items,
+                relevance_floor=tuning.relevance_floor,
+                max_tokens_per_chunk=tuning.max_tokens_per_chunk,
+            )
+            prompt_retrieval = RetrievalResult(
+                items=compressed.items,
+                traces=retrieval.traces,
+                profile=retrieval.profile,
+            )
+            failures = compressed.failure_count
+            dropped = compressed.dropped_count
+
+        prompt = self._render_answer_prompt(query, prompt_retrieval)
+        prompt_tokens_before = count_tokens(raw_prompt)
+        prompt_tokens_after = count_tokens(prompt)
+        yield AnswerEvent(
+            type="compression_done",
+            data={
+                "enabled": compression_on,
+                "source_count_before": len(retrieval.items),
+                "source_count_after": len(prompt_retrieval.items),
+                "failure_count": failures,
+                "dropped_count": dropped,
+                "prompt_tokens_before": prompt_tokens_before,
+                "prompt_tokens_after": prompt_tokens_after,
+                "_result": prompt_retrieval,
+            },
+        )
+
+        if not prompt_retrieval.items:
+            text = (
+                "The retrieved material did not contain evidence relevant enough "
+                "to this question, so I cannot give a grounded answer."
+            )
+            yield AnswerEvent(type="delta", data={"text": text})
+            yield AnswerEvent(type="citations", data={"citations": [], "_sources": []})
+            yield AnswerEvent(type="done", data={"finish_reason": "no_relevant_sources"})
+            return
+
+        sources = [
+            SourceCitation(
+                index=i,
+                kind=item.kind,
+                title=item.title,
+                citation=item.citation,
+                license_class=item.license_class,
+                document_id=item.document_id,
+                chunk_id=item.id if item.kind == "chunk" else None,
+                section_path=item.section_path,
+            )
+            for i, item in enumerate(prompt_retrieval.items, start=1)
+        ]
+        model = llm.describe()
+        yield AnswerEvent(type="generation_started", data={"model": model})
+
+        collected: list[str] = []
+        try:
+            async for delta in llm.stream(prompt, max_tokens=max_tokens):
+                collected.append(delta)
+                yield AnswerEvent(type="delta", data={"text": delta})
+        except Exception as exc:
+            # Full detail goes to the server log; the client gets the class
+            # name only, since provider exception text can echo request
+            # internals.
+            log.warning("answer_generation_failed", error=type(exc).__name__, detail=str(exc)[:500])
+            yield AnswerEvent(
+                type="error",
+                data={
+                    "code": "generation_failed",
+                    "message": f"{type(exc).__name__}: generation failed (see server logs)",
+                },
+            )
+            return
+
+        text = "".join(collected)
+        cited_indices = {int(m) for m in _CITATION_RE.findall(text)}
+        for source in sources:
+            source.cited = source.index in cited_indices
+        yield AnswerEvent(
+            type="citations",
+            data={
+                "citations": [
+                    {
+                        "index": s.index,
+                        "kind": s.kind,
+                        "title": s.title,
+                        "citation": s.citation,
+                        "license_class": s.license_class,
+                        "document_id": s.document_id,
+                        "chunk_id": s.chunk_id,
+                        "section_path": s.section_path,
+                        "cited": s.cited,
+                    }
+                    for s in sources
+                ],
+                "_sources": sources,
+            },
+        )
+        yield AnswerEvent(type="done", data={"finish_reason": "stop"})
+
+    def _render_answer_prompt(self, query: str, retrieval: RetrievalResult) -> str:
+        return self.domain.render_prompt(
+            "answer",
+            DOMAIN_NAME=self.domain.name,
+            QUERY=query,
+            SOURCES=format_sources(retrieval),
+        )
+
+
+def format_sources(retrieval: RetrievalResult) -> str:
+    """The numbered source block shown to the answer model (and to the blind
+    grounding judge, which must see exactly what the assistant saw)."""
+    blocks: list[str] = []
+    for i, item in enumerate(retrieval.items, start=1):
+        header = f"[{i}] {item.title}"
+        if item.section_path:
+            header += f" ({item.section_path})"
+        if item.kind == "community":
+            header += " [knowledge graph overview]"
+        if item.citation:
+            header += f"\nCitation: {item.citation}"
+        blocks.append(f"{header}\n{item.content}")
+    return "\n\n---\n\n".join(blocks)

sci_rag/campaigns/__init__.py

@@ -0,0 +1,50 @@
+"""Build, screen, and report reproducible scientific-document campaigns."""
+
+from sci_rag.campaigns.build import (
+    CampaignBuildReport,
+    build_campaign,
+    load_discovered_candidates,
+)
+from sci_rag.campaigns.discovery import (
+    CandidateWork,
+    DiscoveryReport,
+    discover_by_dois,
+    discover_by_topic,
+    normalize_doi,
+)
+from sci_rag.campaigns.download import DownloadOutcome, download_pdf, pdf_filename
+from sci_rag.campaigns.licensing_map import license_class_for
+from sci_rag.campaigns.manifest import ManifestItem, write_campaign_manifest
+from sci_rag.campaigns.resolve import OaResolution, resolve_unpaywall
+from sci_rag.campaigns.screen import (
+    ScreeningDecision,
+    ScreeningReport,
+    apply_human_review,
+    screen_campaign,
+)
+from sci_rag.campaigns.state import CampaignRecord, CampaignState
+
+__all__ = [
+    "CampaignBuildReport",
+    "CampaignRecord",
+    "CampaignState",
+    "CandidateWork",
+    "DiscoveryReport",
+    "DownloadOutcome",
+    "ManifestItem",
+    "OaResolution",
+    "ScreeningDecision",
+    "ScreeningReport",
+    "apply_human_review",
+    "build_campaign",
+    "discover_by_dois",
+    "discover_by_topic",
+    "download_pdf",
+    "license_class_for",
+    "load_discovered_candidates",
+    "normalize_doi",
+    "pdf_filename",
+    "resolve_unpaywall",
+    "screen_campaign",
+    "write_campaign_manifest",
+]