offline_retriever.py

"""In-process offline retriever (BM25 over the local law corpus).

This backend makes the whole experiment runnable without the external
`retrieval-pipeline` HTTP service: it reads the Markdown law files under
``laws/``, splits them into article-level chunks (每一条法条一个 chunk), and
scores queries with Okapi BM25. Retrieval therefore runs fully offline with no
API key and no server; only the LLM answer-generation step (in ``agent.py``)
still needs a provider API.

Chinese text is tokenised with ``jieba`` when available, falling back to a
character uni/bi-gram tokeniser so the module works with only the standard
library installed.
"""

import os
import re
import math
import logging
from pathlib import Path
from typing import Dict, Any, List, Optional
from collections import Counter, defaultdict

logger = logging.getLogger(__name__)


# Article marker at the start of a line, e.g. 第二百三十五条 / 第一百三十三条之一
_ARTICLE_RE = re.compile(r"^第[一二三四五六七八九十百千零两0-9]+条(?:之[一二三四五六七八九十0-9]+)?")


def _tokenize(text: str) -> List[str]:
    """Tokenise mixed Chinese/English text.

    Prefers jieba; otherwise emits ASCII words plus Chinese character uni- and
    bi-grams, which is enough for lexical BM25 matching without extra deps.
    """
    try:
        import jieba  # type: ignore
        return [t for t in jieba.cut(text) if t.strip()]
    except Exception:
        tokens: List[str] = []
        for m in re.findall(r"[a-zA-Z0-9]+|[一-鿿]+", text):
            if m[0].isascii():
                tokens.append(m.lower())
            else:
                tokens.extend(list(m))  # unigrams
                tokens.extend(m[i:i + 2] for i in range(len(m) - 1))  # bigrams
        return tokens


class OfflineRetriever:
    """Okapi BM25 retriever over article-level chunks of the law corpus."""

    def __init__(self,
                 corpus_path: str = "laws",
                 k1: float = 1.5,
                 b: float = 0.75,
                 extensions: Optional[List[str]] = None):
        self.corpus_path = corpus_path
        self.k1 = k1
        self.b = b
        self.extensions = extensions or [".md", ".txt"]

        self.chunks: List[Dict[str, Any]] = []          # {doc_id, chunk_id, title, category, text}
        self.documents: Dict[str, Dict[str, Any]] = {}   # doc_id -> {title, category, file, content}
        self._doc_freqs: List[Counter] = []              # per-chunk term frequencies
        self._doc_lens: List[int] = []
        self._df: Dict[str, int] = defaultdict(int)      # document frequency per term
        self._idf: Dict[str, float] = {}
        self._avg_len: float = 0.0

        self._build_index()

    # ------------------------------------------------------------------ build
    def _iter_files(self):
        root = Path(self.corpus_path)
        if not root.exists():
            logger.warning(f"Offline corpus path not found: {root}")
            return
        for path in sorted(root.rglob("*")):
            if path.is_file() and path.suffix in self.extensions:
                yield path

    def _split_articles(self, content: str) -> List[str]:
        """Split a law document into article-level chunks.

        Falls back to blank-line paragraph grouping when the file has no
        ``第X条`` markers (e.g. non-statute documents).
        """
        lines = content.splitlines()
        articles: List[str] = []
        current: List[str] = []
        seen_article = False

        for line in lines:
            if _ARTICLE_RE.match(line.strip()):
                seen_article = True
                if current:
                    articles.append("\n".join(current).strip())
                current = [line]
            else:
                current.append(line)
        if current:
            articles.append("\n".join(current).strip())

        if not seen_article:
            # No article markers: group by blank lines into ~paragraph chunks.
            articles = [p.strip() for p in content.split("\n\n") if p.strip()]

        return [a for a in articles if a]

    def _build_index(self):
        for path in self._iter_files():
            try:
                content = path.read_text(encoding="utf-8")
            except Exception as e:
                logger.error(f"Error reading {path}: {e}")
                continue

            category = path.parent.name
            title = path.stem
            doc_id = f"{category}/{title}"
            self.documents[doc_id] = {
                "doc_id": doc_id,
                "title": title,
                "category": category,
                "file": str(path),
                "content": content,
            }

            for idx, article in enumerate(self._split_articles(content)):
                if len(article) < 4:
                    continue
                chunk_id = f"{doc_id}_chunk_{idx}"
                self.chunks.append({
                    "doc_id": doc_id,
                    "chunk_id": chunk_id,
                    "title": title,
                    "category": category,
                    "text": article,
                })

        # Build BM25 statistics.
        for chunk in self.chunks:
            tf = Counter(_tokenize(chunk["text"]))
            self._doc_freqs.append(tf)
            self._doc_lens.append(sum(tf.values()))
            for term in tf:
                self._df[term] += 1

        n = len(self.chunks)
        self._avg_len = (sum(self._doc_lens) / n) if n else 0.0
        for term, df in self._df.items():
            # BM25 idf with +1 to stay non-negative.
            self._idf[term] = math.log(1 + (n - df + 0.5) / (df + 0.5))

        logger.info(
            f"OfflineRetriever indexed {n} chunks from {len(self.documents)} "
            f"documents under '{self.corpus_path}'"
        )

    # ----------------------------------------------------------------- search
    def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
        """Return the ``top_k`` article chunks scored by BM25 for ``query``."""
        if not self.chunks:
            return []

        q_terms = _tokenize(query)
        scored: List[tuple] = []
        for i, tf in enumerate(self._doc_freqs):
            dl = self._doc_lens[i]
            score = 0.0
            for term in q_terms:
                f = tf.get(term)
                if not f:
                    continue
                idf = self._idf.get(term, 0.0)
                denom = f + self.k1 * (1 - self.b + self.b * dl / (self._avg_len or 1))
                score += idf * (f * (self.k1 + 1)) / denom
            if score > 0:
                scored.append((score, i))

        scored.sort(reverse=True)
        results: List[Dict[str, Any]] = []
        for score, i in scored[:top_k]:
            chunk = self.chunks[i]
            results.append({
                "doc_id": chunk["doc_id"],
                "chunk_id": chunk["chunk_id"],
                "text": chunk["text"],
                "score": float(score),
                "metadata": {
                    "title": chunk["title"],
                    "category": chunk["category"],
                    "source": "offline",
                },
            })
        return results

    def get_document(self, doc_id: str) -> Dict[str, Any]:
        """Return the full source document for ``doc_id``."""
        doc = self.documents.get(doc_id)
        if not doc:
            return {"error": f"Document {doc_id} not found"}
        return {
            "doc_id": doc_id,
            "content": doc["content"],
            "metadata": {
                "title": doc["title"],
                "category": doc["category"],
                "file": doc["file"],
                "source": "offline",
            },
        }