Replace custom RAG with LlamaIndex

README.md

@@ -7,7 +7,7 @@
 - **后端**: FastAPI
 - **前端**: Streamlit
 - **LLM**: LiteLLM 调用通义千问（`DEFAULT_LLM` 需为已登记模型）
-- **检索**: 进程内分块 + 关键字重叠检索（无向量库）
+- **检索**: LlamaIndex 句子分块 + BM25 检索（无向量库）
 
 ## 工作方式
 

pyproject.toml

@@ -14,6 +14,8 @@ dependencies = [
     "uvicorn~=0.24.0",
     "streamlit~=1.52.1",
     "requests~=2.32.5",
+    "llama-index-core>=0.14.23",
+    "llama-index-retrievers-bm25>=0.7.1",
 ]
 
 [dependency-groups]

src/backend/infra/rag/loader.py

@@ -1,6 +1,8 @@
-import uuid
 from pathlib import Path
 
+from llama_index.core import Document, Settings, SimpleDirectoryReader
+from llama_index.core.node_parser import SentenceSplitter
+
 from src.backend.infra.rag.store import (
     DocumentChunk,
     KnowledgeStore,
@@ -12,23 +14,22 @@
 
 SUPPORTED_SUFFIXES = {".txt", ".md", ".markdown"}
 
+Settings.llm = None
+Settings.embed_model = None
+
+
+def _splitter(chunk_size: int, overlap: int) -> SentenceSplitter:
+    return SentenceSplitter(chunk_size=chunk_size, chunk_overlap=overlap)
+
 
 def chunk_text(text: str, chunk_size: int, overlap: int) -> list[str]:
     text = text.strip()
     if not text:
         return []
-    if len(text) <= chunk_size:
-        return [text]
-
-    chunks: list[str] = []
-    start = 0
-    while start < len(text):
-        end = start + chunk_size
-        chunks.append(text[start:end])
-        if end >= len(text):
-            break
-        start = max(end - overlap, start + 1)
-    return chunks
+    nodes = _splitter(chunk_size, overlap).get_nodes_from_documents(
+        [Document(text=text)]
+    )
+    return [node.get_content() for node in nodes]
 
 
 def validate_chunk_params(chunk_size: int, overlap: int) -> None:
@@ -58,21 +59,25 @@ def ingest_text(
     overlap: int,
 ) -> int:
     store.remove_document(doc_name)
-    parts = chunk_text(content, chunk_size, overlap)
-    doc_id = str(uuid.uuid4())
-    for idx, part in enumerate(parts):
+    nodes = _splitter(chunk_size, overlap).get_nodes_from_documents(
+        [Document(text=content.strip())]
+    )
+    for idx, node in enumerate(nodes):
+        node.metadata["doc_name"] = doc_name
+        node.metadata["chunk_index"] = idx
+        store.nodes.append(node)
         store.chunks.append(
             DocumentChunk(
-                doc_id=doc_id,
+                doc_id=node.node_id,
                 doc_name=doc_name,
-                content=part,
+                content=node.get_content(),
                 chunk_index=idx,
             )
         )
     store.document_names.append(doc_name)
     store.document_count += 1
-    logger.info(f"Ingested RAG doc: {doc_name} ({len(parts)} chunks)")
-    return len(parts)
+    logger.info(f"Ingested RAG doc: {doc_name} ({len(nodes)} chunks)")
+    return len(nodes)
 
 
 def load_docs(
@@ -81,7 +86,7 @@ def load_docs(
     chunk_size: int | None = None,
     overlap: int | None = None,
 ) -> KnowledgeStore:
-    """从目录加载全部文档到内存知识库。"""
+    """用 LlamaIndex 从目录加载文档并分块。"""
     size, ov = _resolve_chunk_params(chunk_size, overlap)
     path = Path(docs_path or settings.rag_settings.docs_path).expanduser().resolve()
     store = KnowledgeStore(
@@ -100,24 +105,34 @@ def load_docs(
         logger.warning(f"RAG docs path is not a directory: {path}")
         return store
 
-    files = sorted(
-        p
-        for p in path.rglob("*")
-        if p.is_file() and p.suffix.lower() in SUPPORTED_SUFFIXES
-    )
+    try:
+        documents = SimpleDirectoryReader(
+            input_dir=str(path),
+            required_exts=sorted(SUPPORTED_SUFFIXES),
+            recursive=True,
+            filename_as_id=True,
+        ).load_data()
+    except ValueError:
+        logger.warning(f"No supported documents found in: {path}")
+        return store
 
-    for file_path in files:
+    for document in documents:
+        file_path = Path(
+            document.metadata.get("file_path")
+            or document.metadata.get("file_name")
+            or ""
+        )
         try:
-            content = file_path.read_text(encoding="utf-8")
-        except UnicodeDecodeError:
-            content = file_path.read_text(encoding="utf-8", errors="ignore")
-            logger.warning(f"Decoded with errors ignored: {file_path}")
-        except Exception as e:
-            logger.exception(f"Failed to read {file_path}: {e}")
-            continue
-
-        doc_name = str(file_path.relative_to(path))
-        ingest_text(store, doc_name, content, chunk_size=size, overlap=ov)
+            doc_name = str(file_path.relative_to(path))
+        except ValueError:
+            doc_name = file_path.name or "document"
+        ingest_text(
+            store,
+            doc_name,
+            document.get_content(),
+            chunk_size=size,
+            overlap=ov,
+        )
 
     logger.info(
         f"RAG knowledge loaded: {store.document_count} docs, "

src/backend/infra/rag/retriever.py

@@ -1,53 +1,43 @@
-import re
+from llama_index.retrievers.bm25 import BM25Retriever
 
 from src.backend.infra.rag.store import DocumentChunk, get_knowledge_store
 from src.shared.config import settings
+from src.shared.logger import logger
 
-
-def _tokenize(text: str) -> set[str]:
-    # 中英文简单切词：连续中文/字母数字为一 token
-    tokens = re.findall(r"[\u4e00-\u9fff]+|[a-zA-Z0-9]+", text.lower())
-    result: set[str] = set()
-    for token in tokens:
-        if re.fullmatch(r"[\u4e00-\u9fff]+", token) and len(token) > 1:
-            # 中文再拆成字与 bigram，提高短查询命中率
-            result.update(token)
-            result.update(token[i : i + 2] for i in range(len(token) - 1))
-        else:
-            result.add(token)
-    return {t for t in result if t}
-
-
-def _score(query_tokens: set[str], chunk: DocumentChunk) -> float:
-    """
-    计算查询词与文档内容的相似度
-    """
-    if not query_tokens:
-        return 0.0
-    chunk_tokens = _tokenize(chunk.content)
-    if not chunk_tokens:
-        return 0.0
-    overlap = query_tokens & chunk_tokens
-    if not overlap:
-        # 退化为子串匹配
-        q = "".join(sorted(query_tokens, key=len, reverse=True)[:3])
-        if q and q in chunk.content.lower():
-            return 0.1
-        return 0.0
-    return len(overlap) / len(query_tokens)
+_BM25_TOKEN_PATTERN = r"(?u)[\u4e00-\u9fff]|[a-zA-Z0-9]+"
 
 
 def retrieve(query: str, top_k: int | None = None) -> list[DocumentChunk]:
     store = get_knowledge_store()
-    if not store.chunks:
+    if not store.nodes:
         return []
 
     k = top_k if top_k is not None else settings.rag_settings.top_k
-    query_tokens = _tokenize(query)
-    scored = [(_score(query_tokens, chunk), chunk) for chunk in store.chunks]
-    scored = [(s, c) for s, c in scored if s > 0]
-    scored.sort(key=lambda x: x[0], reverse=True)
-    return [c for _, c in scored[:k]]
+    try:
+        retriever = BM25Retriever.from_defaults(
+            nodes=store.nodes,
+            similarity_top_k=k,
+            skip_stemming=True,
+            language="zh",
+            token_pattern=_BM25_TOKEN_PATTERN,
+        )
+        results = retriever.retrieve(query)
+    except Exception:
+        logger.exception("LlamaIndex BM25 retrieve failed")
+        return []
+
+    chunks: list[DocumentChunk] = []
+    for result in results:
+        node = result.node
+        chunks.append(
+            DocumentChunk(
+                doc_id=node.node_id,
+                doc_name=str(node.metadata.get("doc_name") or ""),
+                content=node.get_content(),
+                chunk_index=int(node.metadata.get("chunk_index") or 0),
+            )
+        )
+    return chunks
 
 
 def format_context(chunks: list[DocumentChunk]) -> str:

src/backend/infra/rag/store.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from typing import Any
 
 
 @dataclass
@@ -11,25 +12,30 @@ class DocumentChunk:
 
 @dataclass
 class KnowledgeStore:
-    """进程内知识库"""
+    """进程内知识库（LlamaIndex nodes + 预览分块）"""
 
     docs_path: str = ""
     chunk_size: int = 500
     chunk_overlap: int = 50
     document_count: int = 0
     document_names: list[str] = field(default_factory=list)
     chunks: list[DocumentChunk] = field(default_factory=list)
+    nodes: list[Any] = field(default_factory=list)
 
     def clear(self) -> None:
         self.document_count = 0
         self.document_names = []
         self.chunks = []
+        self.nodes = []
 
     def remove_document(self, doc_name: str) -> None:
         if doc_name in self.document_names:
             self.document_names.remove(doc_name)
             self.document_count = max(0, self.document_count - 1)
         self.chunks = [chunk for chunk in self.chunks if chunk.doc_name != doc_name]
+        self.nodes = [
+            node for node in self.nodes if node.metadata.get("doc_name") != doc_name
+        ]
 
 
 _store: KnowledgeStore | None = None

src/frontend/components/sidebar/sidebar.py

@@ -28,28 +28,28 @@ def _render_knowledge_panel():
             f"{status['document_count']} 个文档 · {status['chunk_count']} 个分块"
         )
         chunk_size = st.number_input(
-            "分块大小（字符）",
+            "分块大小",
             min_value=50,
             max_value=10000,
             value=int(status.get("chunk_size", 500)),
             step=50,
-            help="每个分块最多包含多少字符",
+            help="每个分块的目标长度（按句子切分）",
         )
         chunk_overlap = st.number_input(
-            "重叠长度（字符）",
+            "重叠长度",
             min_value=0,
             max_value=5000,
             value=int(status.get("chunk_overlap", 50)),
             step=10,
-            help="相邻两个分块之间重复多少字符，需小于分块大小",
+            help="相邻两个分块之间重复的内容长度，需小于分块大小",
         )
     else:
         st.warning("无法连接后端知识库服务")
         chunk_size = st.number_input(
-            "分块大小（字符）", min_value=50, max_value=10000, value=500
+            "分块大小", min_value=50, max_value=10000, value=500
         )
         chunk_overlap = st.number_input(
-            "重叠长度（字符）", min_value=0, max_value=5000, value=50
+            "重叠长度", min_value=0, max_value=5000, value=50
         )
 
     uploaded_files = st.file_uploader(

tests/unittest/test_backend/test_app/test_knowledge_service.py

@@ -29,10 +29,10 @@ def reset_store():
 
 class TestChunkHelpers:
     def test_chunk_text_with_overlap(self):
-        text = "a" * 100
+        text = "年假政策说明。" * 30
         chunks = chunk_text(text, chunk_size=40, overlap=10)
         assert len(chunks) >= 2
-        assert all(len(chunk) <= 40 for chunk in chunks)
+        assert all(chunk.strip() for chunk in chunks)
 
     def test_validate_chunk_params_rejects_large_overlap(self):
         with pytest.raises(ValueError, match="chunk_overlap"):

tests/unittest/test_backend/test_infra/test_rag/test_rag.py

@@ -15,6 +15,7 @@ def test_load_docs_from_sample_dir():
     store = load_docs(str(docs_path))
     assert store.document_count >= 1
     assert len(store.chunks) >= 1
+    assert len(store.nodes) >= 1
 
 
 def test_retrieve_leave_policy():