feat: Chinese retrieval enhancements - pluggable BM25 tokenizers, per-doc diversity cap, Gradio Web UI

.gitignore

@@ -17,3 +17,5 @@ generated/
 dist/
 build/
 *.egg-info/
+
+.venv/

config.example.toml

@@ -30,9 +30,18 @@ search = 5        # number of hits per search
 [retrieval]
 hybrid = true     # fuse vector search with BM25 (Reciprocal Rank Fusion)
 rrf_k = 60        # RRF smoothing constant; higher = flatter fusion weights
+max_per_doc = 2   # max chunks one document may contribute per search (0 = no cap)
+
+[bm25]
+tokenizer = "default"   # "default" = CJK single-char (no deps) | "jieba" = word-level
+                        # Chinese segmentation (pip install 'loci-rag[zh]' or jieba)
 
 [watch]
 interval = 30     # seconds between scans in `watch` mode
 
+[webui]
+host = "127.0.0.1"   # `loci webui` bind address (Gradio UI, pip install 'loci-rag[ui]')
+port = 7860
+
 [store]
 path = "./chroma_db"   # where the vector store lives on disk

pyproject.toml

@@ -34,6 +34,8 @@ pdf = ["pymupdf4llm>=0.0.21", "pypdf>=4"]
 ocr = ["rapidocr-onnxruntime>=1.3"]
 docx = ["python-docx>=1.1"]
 rerank = ["sentence-transformers>=3"]
+zh = ["jieba>=0.42"]
+ui = ["gradio>=4"]
 dev = ["pytest>=8"]
 
 [project.scripts]

src/loci/bm25.py

@@ -1,28 +1,33 @@
-"""Native BM25 ranking with a CJK-aware tokenizer. No dependencies."""
+"""Native BM25 ranking with pluggable tokenizers. No dependencies by default."""
 from __future__ import annotations
 
 import math
-import re
 
-# Latin words stay whole; CJK text becomes single characters (strong enough
-# for BM25 ranking and keeps the tokenizer dependency-free).
-_TOKEN = re.compile(r"[a-z0-9]+|[\u4e00-\u9fff]")
+from loci.tokenizers import default_tokenize, get_tokenizer
 
 
 def tokenize(text: str) -> list[str]:
-    return _TOKEN.findall(text.lower())
+    """Backward-compatible alias for the default CJK single-char tokenizer."""
+    return default_tokenize(text)
 
 
 class BM25:
-    """Okapi BM25 over a corpus of {key: text}, built once per process."""
+    """Okapi BM25 over a corpus of {key: text}, built once per process.
 
-    def __init__(self, corpus: dict[str, str], k1: float = 1.5, b: float = 0.75):
+    `tokenizer` accepts a callable text->list[str]; defaults to the
+    dependency-free CJK single-char tokenizer. Pass `get_tokenizer("jieba")`
+    for word-level Chinese segmentation (requires the optional jieba package).
+    """
+
+    def __init__(self, corpus: dict[str, str], k1: float = 1.5, b: float = 0.75,
+                 tokenizer=None):
         self.k1, self.b = k1, b
+        self._tok = tokenizer or default_tokenize
         self.tf: dict[str, dict[str, int]] = {}
         self.dl: dict[str, int] = {}
         self.df: dict[str, int] = {}
         for key, text in corpus.items():
-            toks = tokenize(text)
+            toks = self._tok(text)
             freq: dict[str, int] = {}
             for t in toks:
                 freq[t] = freq.get(t, 0) + 1
@@ -42,7 +47,7 @@ def score(self, query: str) -> list[tuple[str, float]]:
             if not dl:
                 continue
             s = 0.0
-            for term in tokenize(query):
+            for term in self._tok(query):
                 df = self.df.get(term)
                 if not df:
                     continue

src/loci/cli.py

@@ -30,6 +30,8 @@ def build(cfg):
 
 
 def make_retriever(cfg, embedder, store):
+    """Build the Retriever from config — the single wiring point every entry
+    (CLI / HTTP API / WebUI) shares, so options can't drift apart."""
     from loci.retriever import Retriever
     import os as _os
     fb_path = _os.path.join(cfg.store["path"], "feedback.jsonl")
@@ -39,7 +41,9 @@ def make_retriever(cfg, embedder, store):
                      feedback_path=fb_path,
                      rerank_provider=cfg.retrieval.get("rerank_provider", "llm"),
                      local_rerank_model=cfg.retrieval.get(
-                         "local_rerank_model", "BAAI/bge-reranker-base"))
+                         "local_rerank_model", "BAAI/bge-reranker-base"),
+                     bm25_tokenizer=cfg.bm25.get("tokenizer", "default"),
+                     max_per_doc=cfg.retrieval.get("max_per_doc", 2))
 
 
 def cmd_ingest(cfg, force: bool = False) -> None:
@@ -126,6 +130,14 @@ def cmd_ask(cfg, question: str, rerank: bool | None = None,
         print("\n" + verify_answer(cfg.llm, question, reply, hits))
 
 
+def cmd_webui(cfg, host: str | None = None, port: int | None = None) -> None:
+    """Serve the Gradio web UI (optional `ui` extra) over the same pipeline."""
+    from loci.webui import run_app   # deferred: webui imports cli back
+    host = host or cfg.webui.get("host") or "127.0.0.1"
+    port = port or int(cfg.webui.get("port") or 7860)
+    run_app(cfg, host=host, port=port)
+
+
 def cmd_feedback(cfg, verdict: str) -> None:
     """Rate the chunks used in the last ask; bad ratings down-weight chunks."""
     last = Path(cfg.store["path"]) / ".last_ask.json"
@@ -240,7 +252,7 @@ def cmd_sync(cfg, direction: str) -> None:
         print(f"  {name}: {direction} done")
 
 
-def cmd_bench(cfg, cases_file: str, k: int = 5) -> None:
+def cmd_bench(cfg, cases_file: str, k: int = 5, tokenizer: str | None = None) -> None:
     import json as _json
     p = Path(cases_file)
     if not p.exists():
@@ -249,6 +261,8 @@ def cmd_bench(cfg, cases_file: str, k: int = 5) -> None:
     cases = [_json.loads(l) for l in p.read_text(encoding="utf-8").splitlines() if l.strip()]
     embedder, store = build(cfg)
     retriever = make_retriever(cfg, embedder, store)
+    if tokenizer:   # per-run override for A/B comparisons
+        retriever.bm25_tokenizer = tokenizer
     for label, hybrid in [("vector-only", False), ("hybrid", True)]:
         retriever.hybrid = hybrid
         ok = 0
@@ -455,7 +469,8 @@ def main() -> None:
                           help="only files modified on/after DATE (YYYY-MM-DD or YYYY-MM)")
     p_search.add_argument("-e", "--exact", metavar="PHRASE",
                           help="only hits containing this exact phrase")
-    p_search.add_argument("-k", type=int, help="override top_k")
+    p_search.add_argument("-k", "--top_k", dest="k", type=int,
+                          help="override top_k (retrieved chunks per query)")
     p_search.add_argument("--rerank", nargs="?", const=True, default=None,
                           metavar="PROVIDER", help="rerank candidates; optional "
                           "value llm|local overrides [retrieval] rerank_provider")
@@ -466,7 +481,8 @@ def main() -> None:
                        help="scope retrieval to paths containing this substring")
     p_ask.add_argument("--since", metavar="DATE", type=parse_since,
                        help="only files modified on/after DATE (YYYY-MM-DD or YYYY-MM)")
-    p_ask.add_argument("-k", type=int, help="override top_k")
+    p_ask.add_argument("-k", "--top_k", dest="k", type=int,
+                       help="override top_k (retrieved chunks per query)")
     p_ask.add_argument("--rerank", nargs="?", const=True, default=None,
                        metavar="PROVIDER", help="rerank candidates; optional "
                        "value llm|local overrides [retrieval] rerank_provider")
@@ -510,7 +526,16 @@ def main() -> None:
 
     p_bench = sub.add_parser("bench", help="run a retrieval benchmark (hit@k) on a cases file")
     p_bench.add_argument("cases", help="JSONL file with {query, expect} per line")
-    p_bench.add_argument("-k", type=int, default=5, help="hits per query (default 5)")
+    p_bench.add_argument("-k", "--top_k", dest="k", type=int, default=5,
+                         help="hits per query (default 5)")
+    p_bench.add_argument("--tokenizer", default=None,
+                         help="BM25 tokenizer override for this run (default|jieba)")
+
+    p_webui = sub.add_parser("webui", help="run the Gradio web UI (optional [ui] extra)")
+    p_webui.add_argument("--host", default=None,
+                         help="bind address (default: [webui] host, or 127.0.0.1)")
+    p_webui.add_argument("--port", type=int, default=None,
+                         help="port (default: [webui] port, or 7860)")
 
     sub.add_parser("stats", help="show what is in the index")
     sub.add_parser("doctor", help="check config, endpoints, and store health")
@@ -558,7 +583,10 @@ def main() -> None:
     elif args.cmd == "sync":
         cmd_sync(cfg, args.direction)
     elif args.cmd == "bench":
-        cmd_bench(cfg, args.cases, k=args.k)
+        cmd_bench(cfg, args.cases, k=args.k, tokenizer=args.tokenizer)
+    elif args.cmd == "webui":
+        cfg.validate()
+        cmd_webui(cfg, host=args.host, port=args.port)
     elif args.cmd == "wiki":
         cmd_wiki_suggest(cfg) if getattr(args, "suggest", False) else cmd_wiki(cfg, args.topic, k=args.k)
     elif args.cmd == "watch":

src/loci/config.py

@@ -16,13 +16,16 @@
     "top_k": {"search": 5},
     "retrieval": {"hybrid": True, "rrf_k": 60, "rerank": False,
                   "rerank_provider": "llm",
-                  "local_rerank_model": "BAAI/bge-reranker-base"},
+                  "local_rerank_model": "BAAI/bge-reranker-base",
+                  "max_per_doc": 2},   # chunks per document in one result list
+    "bm25": {"tokenizer": "default"},
     "watch": {"interval": 30},
     "memories": {"path": "./memories"},
     "wiki": {"path": "./wiki"},
     "chat": {"auto_extract": True},
     "sync": {"remote": ""},
     "http": {"host": "127.0.0.1", "port": 8765, "token": ""},
+    "webui": {"host": "127.0.0.1", "port": 7860},
     "store": {"path": "./chroma_db"},
 }
 
@@ -34,12 +37,14 @@ class Config:
     chunk: dict = field(default_factory=dict)
     top_k: dict = field(default_factory=dict)
     retrieval: dict = field(default_factory=dict)
+    bm25: dict = field(default_factory=dict)
     watch: dict = field(default_factory=dict)
     memories: dict = field(default_factory=dict)
     wiki: dict = field(default_factory=dict)
     chat: dict = field(default_factory=dict)
     sync: dict = field(default_factory=dict)
     http: dict = field(default_factory=dict)
+    webui: dict = field(default_factory=dict)
     store: dict = field(default_factory=dict)
     sources: list = field(default_factory=list)
 
@@ -107,4 +112,9 @@ def load(path: str = "config.toml") -> Config:
         cfg.top_k["search"] = max(int(cfg.top_k["search"]), 1)
     except (TypeError, ValueError):
         cfg.top_k["search"] = DEFAULTS["top_k"]["search"]
+    known_tokenizers = ("default", "jieba")
+    if cfg.bm25.get("tokenizer") not in known_tokenizers:
+        print(f"[warn] bm25.tokenizer={cfg.bm25.get('tokenizer')!r} is not one of "
+              f"{known_tokenizers} — using 'default'", file=sys.stderr)
+        cfg.bm25["tokenizer"] = "default"
     return cfg

src/loci/loaders.py

@@ -38,6 +38,18 @@
 
 _ocr_engine = None  # lazy singleton (model load is ~1s)
 
+_warned_missing: set[str] = set()   # extras already warned about this run
+
+
+def _warn_missing_extra(extra: str, install: str) -> None:
+    """Warn once per run that files need an optional extra, with the exact
+    install command — silently skipping .pdf/.docx files confuses users."""
+    if extra in _warned_missing:
+        return
+    _warned_missing.add(extra)
+    print(f"[warn] {extra} sources found but the optional dependency is not "
+          f"installed — skipped. Install it with: {install}")
+
 IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff"}
 
 SUFFIXES = {".md", ".txt"}
@@ -296,10 +308,12 @@ def _read_text(p: Path) -> str | None:
             except Exception as e:
                 print(f"[warn] pdf unreadable, skipping: {p.name} ({e})")
                 return None
-        return None  # no pdf extra installed; `doctor` mentions it
+        _warn_missing_extra("pdf", "pip install 'loci-rag[pdf]'")
+        return None  # no pdf extra installed
     if suffix == ".docx":
         if not HAS_DOCX:
-            return None  # silently skip; `doctor` mentions the optional extra
+            _warn_missing_extra("docx", "pip install 'loci-rag[docx]'")
+            return None
         try:
             document = _docx.Document(str(p))
             parts = [para.text for para in document.paragraphs if para.text.strip()]
@@ -389,9 +403,10 @@ def _get_ocr():
 
 
 def _ocr_image(p: Path) -> str | None:
-    """OCR one image file; None (skip silently) when OCR isn't installed."""
+    """OCR one image file; None (skip) when OCR isn't installed."""
     if not HAS_OCR:
-        return None  # `doctor` mentions the optional [ocr] extra
+        _warn_missing_extra("ocr", "pip install 'loci-rag[ocr]'")
+        return None
     try:
         result, _ = _get_ocr()(str(p))
         if not result:

src/loci/retriever.py

@@ -9,27 +9,45 @@
 
 from loci.bm25 import BM25
 from loci.reranker import local_rerank
+from loci.tokenizers import get_tokenizer
 
+# Strict grounding rules: no fabrication, explicit "not in KB" answer, and a
+# mandatory source list so every claim is traceable to a document section.
 SYSTEM_PROMPT = (
     "You are a Q&A assistant over the user's personal knowledge base. "
-    "Answer only from the provided excerpts; if they don't contain the answer, "
-    "say so plainly — do not invent. End your reply with a list of citations "
-    "in the form [source: file path > section]."
+    "Follow these rules strictly:\n"
+    "1. Answer ONLY from the provided excerpts. Never use outside knowledge "
+    "to fill gaps — do not invent or guess.\n"
+    "2. If the excerpts do not contain the answer, say so plainly in the "
+    "user's language (e.g. \"知识库中没有相关信息\") and stop.\n"
+    "3. Structure the answer as short bullet points; each point must be "
+    "grounded in at least one excerpt.\n"
+    "4. End the reply with a 'Sources:' line citing every excerpt used, in "
+    "the form [source: file path > section]."
 )
 
 
 class Retriever:
+    """Vector (+BM25) retrieval over a Store, fused via RRF.
+
+    `max_per_doc` caps how many chunks one source document may contribute to
+    a single result list, so one long file can't crowd out other relevant
+    docs (0 = no cap)."""
     def __init__(self, embedder, store, top_k: int,
                  hybrid: bool = True, rrf_k: int = 60,
                  rerank: bool = False, llm_cfg: dict | None = None,
                  rerank_provider: str = "llm",
                  local_rerank_model: str = "BAAI/bge-reranker-base",
-                 feedback_path: str | None = None):
+                 feedback_path: str | None = None,
+                 bm25_tokenizer: str = "default",
+                 max_per_doc: int = 2):
         self.embedder, self.store, self.top_k = embedder, store, top_k
         self.hybrid, self.rrf_k = hybrid, rrf_k
         self.rerank, self.llm_cfg = rerank, llm_cfg
         self.rerank_provider = rerank_provider
         self.local_rerank_model = local_rerank_model
+        self.bm25_tokenizer = bm25_tokenizer
+        self.max_per_doc = max_per_doc
         self._penalties: dict[str, float] = {}
         if feedback_path and Path(feedback_path).exists():
             import json as _json
@@ -95,7 +113,28 @@ def search(self, query: str, tag: str | None = None,
                 fused = local_rerank(query, fused, self.local_rerank_model)
             elif self.llm_cfg:
                 fused = rerank_hits(self.llm_cfg, query, fused)
-        return fused[:limit]
+        # diversity cap AFTER ranking/filters: the best chunk per doc keeps
+        # its slot, surplus chunks give way to other documents
+        return self._diversify(fused, limit)[:limit]
+
+    def _diversify(self, hits: list[dict], limit: int) -> list[dict]:
+        """Keep at most `max_per_doc` chunks per source document (order kept).
+
+        Applied as the last step before truncation, so a single document that
+        dominates the ranking loses its surplus slots to other docs."""
+        if self.max_per_doc <= 0 or not hits:
+            return hits
+        seen: dict[str, int] = {}
+        out: list[dict] = []
+        for h in hits:
+            src = h.get("source", "")
+            if seen.get(src, 0) >= self.max_per_doc:
+                continue
+            seen[src] = seen.get(src, 0) + 1
+            out.append(h)
+            if len(out) >= limit:
+                break
+        return out
 
     @staticmethod
     def _rrf_merge(ranked_lists: list[list[dict]], rrf_k: int = 60) -> list[dict]:
@@ -114,7 +153,9 @@ def _fuse(self, vhits: list[dict], query: str, fetch: int) -> list[dict]:
         ids, docs = self.store.all_chunks()
         if not ids:
             return vhits
-        ranked = BM25(dict(zip(ids, docs))).score(query)[:fetch]
+        ranked = BM25(dict(zip(ids, docs)),
+                      tokenizer=get_tokenizer(self.bm25_tokenizer)
+                      ).score(query)[:fetch]
         scores: dict[str, float] = {}
         for r, h in enumerate(vhits):
             scores[h["id"]] = scores.get(h["id"], 0.0) + 1.0 / (self.rrf_k + r + 1)