ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/agentic-rag/offline_retriever.py

@@ -0,0 +1,213 @@
+"""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",
+            },
+        }