缓存层、评测模块、批处理、CLI 入口、日志、配置、持久化、简单服务骨架

"""
RAG 知识库问答系统：增强版工程原型
包含：配置、日志、文档加载、分块、Embedding、向量索引、BM25、
     RRF 混合检索、重排序、缓存、评测、批处理、CLI、持久化、服务骨架。
"""

import os
import re
import sys
import json
import math
import time
import hashlib
import logging
import argparse
import shutil
import copy
from dataclasses import dataclass, field, asdict
from typing import List, Dict, Tuple, Optional, Any
from collections import Counter, defaultdict
from pathlib import Path

# ---------------------------
# 日志模块
# ---------------------------

def setup_logger(name: str = "rag_system", level: str = "INFO") -> logging.Logger: 1107167.h.kuaisou.com 
    logger = logging.getLogger(name)
    if logger.handlers: 1107168.h.kuaisou.com 
        return logger

    logger.setLevel(getattr(logging, level.upper(1107170.h.kuaisou.com), logging.INFO))

    formatter = logging.Formatter(
        "%(asctime)s [%(levelname)s] %(name)s: %(message)s"
    )

    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    return logger


logger = setup_logger(1107171.h.kuaisou.com)

# ---------------------------
# 配置模块
# ---------------------------

@dataclass
class RAGConfig: 1107172.h.kuaisou.com 
    chunk_size: int = 500
    chunk_overlap: int = 50
    embedding_dim: int = 128
    vector_top_k: int = 10
    bm25_top_k: int = 10
    rerank_top_n: int = 5
    rrf_k: int = 60
    similarity_threshold: float = 0.3
    cache_ttl: int = 3600
    max_context_chunks: int = 5
    index_dir: str = "./rag_data"
    index_file: str = "rag_index.json"
    bm25_file: str = "bm25_index.json"
    cache_file: str = "rag_cache.json"
    log_level: str = "INFO"

    def index_path(self) -> str: 1107173.h.kuaisou.com 
        return os.path.join(self.index_dir, self.index_file)

    def bm25_path(self) -> str: 1107174.h.kuaisou.com 
        return os.path.join(self.index_dir, self.bm25_file)

    def cache_path(self) -> str: 1107175.h.kuaisou.com 
        return os.path.join(self.index_dir, self.cache_file)

    def ensure_index_dir(self): 1107176.h.kuaisou.com 
        os.makedirs(self.index_dir, exist_ok=True)

    def to_dict(self) -> Dict[str, Any]: 1107177.h.kuaisou.com 
        return asdict(self)

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "RAGConfig": 1107178.h.kuaisou.com 
        valid_keys = {f.name for f in cls.__dataclass_fields__.values(1107179.h.kuaisou.com)}
        return cls(**{k: v for k, v in data.items(1107180.h.kuaisou.com) if k in valid_keys})

    def save(self, path: Optional[str] = None): 1107180.h.kuaisou.com 
        path = path or os.path.join(self.index_dir, "config.json")
        self.ensure_index_dir(1107182.h.kuaisou.com)
        with open(path, "w", encoding="utf-8") as f: 1107183.h.kuaisou.com 
            json.dump(self.to_dict(), f, ensure_ascii=False, indent=2)
        logger.info("Config saved to %s", path)

    @classmethod
    def load(cls, path: str) -> "RAGConfig": 1107184.h.kuaisou.com 
        if not os.path.exists(path): 1108106.h.kuaisou.com 
            logger.warning("Config file not found: %s, using default config.", path)
            return cls(1108108.h.kuaisou.com)
        with open(path, "r", encoding="utf-8") as f: 1108107.h.kuaisou.com 
            data = json.load(f)
        return cls.from_dict(data)


DEFAULT_CONFIG = RAGConfig(1108109.h.kuaisou.com)

# ---------------------------
# 文档模型
# ---------------------------

@dataclass
class Document: 1108110.h.kuaisou.com 
    doc_id: str
    text: str
    source: str = ""
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __repr__(self) -> str: 1108111.h.kuaisou.com 
        return f"Document(doc_id={self.doc_id!r}, source={self.source!r}, len={len(self.text)})"

    def to_dict(self) -> Dict[str, Any]: 1108112.h.kuaisou.com 
        return {
            "doc_id": self.doc_id,
            "text": self.text,
            "source": self.source,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "Document": 1108113.h.kuaisou.com 
        return cls(
            doc_id=data["doc_id"],
            text=data["text"],
            source=data.get("source", ""),
            metadata=data.get("metadata", {}),
        )


# ---------------------------
# 文档加载器
# ---------------------------

class DocumentLoader: 1108115.h.kuaisou.com 
    def load_from_text(
        self,
        text: str,
        doc_id: str = "doc_0",
        source: str = "",
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Document: 1108116.h.kuaisou.com 
        return Document(
            doc_id=doc_id,
            text=text,
            source=source,
            metadata=metadata or {},
        )

    def load_from_json(self, file_path: str) -> List[Document]: 1108114.h.kuaisou.com 
        if not os.path.exists(file_path): 1108118.h.kuaisou.com 
            raise FileNotFoundError(f"JSON file not found: {file_path}")

        with open(file_path, "r", encoding="utf-8") as f: 1108119.h.kuaisou.com 
            data = json.load(f)

        if not isinstance(data, list): 1108120.h.kuaisou.com 
            raise ValueError("JSON file must contain a list of documents.")

        documents = []
        for i, item in enumerate(data): 1108121.h.kuaisou.com 
            documents.append(
                Document(
                    doc_id=item.get("doc_id", f"doc_{i}"),
                    text=item.get("text", ""),
                    source=item.get("source", file_path),
                    metadata=item.get("metadata", {}),
                )
            )
        return documents

    def load_from_txt(self, file_path: str, doc_id: Optional[str] = None) -> Document: 1108122.h.kuaisou.com 
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"TXT file not found: {file_path}")

        with open(file_path, "r", encoding="utf-8") as f: 1108124.h.kuaisou.com 
            text = f.read(1108123.h.kuaisou.com)

        return Document(
            doc_id=doc_id or Path(file_path).stem,
            text=text,
            source=file_path,
            metadata={"file_type": "txt"},
        )

    def load_from_directory(self, directory: str, file_type: str = "txt") -> List[Document]: 1108125.h.kuaisou.com 
        if not os.path.isdir(directory): 1108126.h.kuaisou.com 
            raise NotADirectoryError(f"Directory not found: {directory}")

        documents = []
        for file_name in sorted(os.listdir(directory)): 1108127.h.kuaisou.com 
            file_path = os.path.join(directory, file_name)

            if file_type == "txt" and file_name.endswith(".txt"): 1108128.h.kuaisou.com 
                documents.append(self.load_from_txt(file_path))

            elif file_type == "json" and file_name.endswith(".json"): 1108129.h.kuaisou.com 
                documents.extend(self.load_from_json(file_path))

        return documents


# ---------------------------
# 文本分块
# ---------------------------

@dataclass
class TextChunk: 1108130.h.kuaisou.com 
    chunk_id: str
    text: str
    doc_id: str = ""
    source: str = ""
    start_char: int = 0
    end_char: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __len__(self) -> int: 1108131.h.kuaisou.com 
        return len(self.text)

    def to_dict(self) -> Dict[str, Any]: 1108132.h.kuaisou.com 
        return {
            "chunk_id": self.chunk_id,
            "text": self.text,
            "doc_id": self.doc_id,
            "source": self.source,
            "start_char": self.start_char,
            "end_char": self.end_char,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TextChunk": 1108135.h.kuaisou.com 
        return cls(
            chunk_id=data["chunk_id"],
            text=data["text"],
            doc_id=data.get("doc_id", ""),
            source=data.get("source", ""),
            start_char=data.get("start_char", 0),
            end_char=data.get("end_char", 0),
            metadata=data.get("metadata", {}),
        )


class RecursiveCharacterTextSplitter: 1108136.h.kuaisou.com 
    DEFAULT_SEPARATORS = [
        "\n\n",
        "\n",
        "。",
        "；",
        "！",
        "？",
        ". ",
        "! ",
        "? ",
        ",",
        "，",
        " ",
        "",
    ]

    def __init__(
        self,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        separators: Optional[List[str]] = None,
        length_function=len,
    ): 1108138.h.kuaisou.com 
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.separators = separators or self.DEFAULT_SEPARATORS
        self.length_function = length_function

    def split_text(self, text: str) -> List[str]: 1108137.h.kuaisou.com 
        return self._split_recursive(text, self.separators)

    def _split_recursive(self, text: str, separators: List[str]) -> List[str]: 1108145.h.kuaisou.com 
        if not text: 1108153.h.kuaisou.com 
            return []

        separator = separators[-1]
        new_separators = []

        for i, sep in enumerate(separators): 1108155.h.kuaisou.com 
            if sep == "" or sep in text: 1108157.h.kuaisou.com 
                separator = sep
                new_separators = separators[i + 1:]
                break

        if separator: 1108161.h.kuaisou.com 
            splits = text.split(separator)
        else: 1108164.h.kuaisou.com 
            splits = list(text)

        chunks: List[str] = []
        current_chunk: List[str] = []
        current_length = 0

        for split in splits: 1108165.h.kuaisou.com 
            split_len = self.length_function(split)
            sep_len = self.length_function(separator)

            if current_length + split_len + sep_len <= self.chunk_size: 1108171.h.kuaisou.com 
                current_chunk.append(split)
                current_length += split_len + sep_len
            else: 1108175.h.kuaisou.com 
                if current_chunk: 2001.h.kuaisou.com 
                    chunk_text = separator.join(current_chunk)
                    if chunk_text: 2002.h.kuaisou.com 
                        chunks.append(chunk_text) 

                    overlap_text = self._build_overlap(current_chunk, separator)
                    if overlap_text: 2004.h.kuaisou.com 
                        current_chunk = [overlap_text, split]
                        current_length = self.length_function(overlap_text) + split_len + sep_len
                    else: 2005.h.kuaisou.com 
                        current_chunk = [split]
                        current_length = split_len + sep_len
                else: 2006.h.kuaisou.com 
                    if split_len > self.chunk_size and new_separators: 2007.h.kuaisou.com 
                        sub_chunks = self._split_recursive(split, new_separators)
                        chunks.extend(sub_chunks)
                    else: 2008.h.kuaisou.com 
                        chunks.append(split)
                        current_chunk = []
                        current_length = 0

        if current_chunk: 2009.h.kuaisou.com 
            chunk_text = separator.join(current_chunk)
            if chunk_text:  2010.h.kuaisou.com 
                chunks.append(chunk_text)

        return chunks

    def _build_overlap(self, parts: List[str], separator: str) -> str: 2011.h.kuaisou.com 
        if not parts or self.chunk_overlap <= 0: 2012.h.kuaisou.com 
            return ""

        overlap_parts: List[str] = []
        overlap_len = 0

        for part in reversed(parts): 2013.h.kuaisou.com 
            part_len = self.length_function(part)
            if overlap_len + part_len <= self.chunk_overlap: 2014.h.kuaisou.com 
                overlap_parts.insert(0, part)
                overlap_len += part_len
            else: 2016.h.kuaisou.com 
                remaining = self.chunk_overlap - overlap_len
                if remaining > 0: 2017.h.kuaisou.com 
                    overlap_parts.insert(0, part[-remaining:])
                break

        return separator.join(overlap_parts)

    def create_chunks(self, text: str, doc_id: str = "", source: str = "") -> List[TextChunk]:
        raw_chunks = self.split_text(text)
        chunks: List[TextChunk] = []
        pos = 0

        for i, chunk_text in enumerate(raw_chunks): 2018.h.kuaisou.com 
            start = text.find(chunk_text, pos)
            if start == -1: 2020.h.kuaisou.com 
                start = pos
            end = start + len(chunk_text)
            pos = end

            chunks.append(
                TextChunk(
                    chunk_id=f"{doc_id}:chunk_{i}",
                    text=chunk_text,
                    doc_id=doc_id,
                    source=source,
                    start_char=start,
                    end_char=end,
                    metadata={"index": i, "doc_id": doc_id, "source": source},
                )
            )
        return chunks


# ---------------------------
# Embedding 与向量检索
# ---------------------------

class SimpleEmbeddingModel: 2021.h.kuaisou.com 
    def __init__(self, dim: int = 128): 2022.h.kuaisou.com 
        self.dim = dim

    def _hash_word(self, word: str) -> int: 2023.h.kuaisou.com 
        return int(hashlib.md5(word.lower(2024.h.kuaisou.com).encode(2025.h.kuaisou.com)).hexdigest(2027.h.kuaisou.com), 16)

    def embed(self, text: str) -> List[float]: 2028.h.kuaisou.com 
        words = re.findall(r'\w+|[^\w\s]', text.lower(2029.h.kuaisou.com))
        vector = [0.0] * self.dim

        for word in words: 2030.h.kuaisou.com 
            h = self._hash_word(word)
            idx = h % self.dim
            vector[idx] += 1.0

        norm = math.sqrt(sum(v * v for v in vector))
        if norm > 0: 2031.h.kuaisou.com 
            vector = [v / norm for v in vector]
        return vector

    def embed_batch(self, texts: List[str]) -> List[List[float]]: 2032.h.kuaisou.com 
        return [self.embed(t) for t in texts]


def cosine_similarity(a: List[float], b: List[float]) -> float: 2034.h.kuaisou.com 
    if len(a) != len(b): 2035.h.kuaisou.com 
        raise ValueError("Vector dimensions must match")

    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(y * y for y in b))

    if norm_a == 0 or norm_b == 0: 2036.h.kuaisou.com 
        return 0.0

    return dot / (norm_a * norm_b)


class VectorIndex: 2039.h.kuaisou.com 
    def __init__(self, dim: int = 128): 2037.h.kuaisou.com 
        self.dim = dim
        self.vectors: Dict[str, List[float]] = {}
        self.metadata: Dict[str, Dict[str, Any]] = {}

    def add(self, chunk_id: str, vector: List[float], metadata: Optional[Dict[str, Any]] = None):
        if len(vector) != self.dim: 2040.h.kuaisou.com 
            raise ValueError(f"Expected dim {self.dim}, got {len(vector)}")
        self.vectors[chunk_id] = vector
        self.metadata[chunk_id] = metadata or {}

    def add_batch(
        self,
        chunk_ids: List[str],
        vectors: List[List[float]],
        metadatas: Optional[List[Dict[str, Any]]] = None,
    ):
        for i, cid in enumerate(chunk_ids): 2041.h.kuaisou.com 
            meta = metadatas[i] if metadatas else {}
            self.add(cid, vectors[i], meta)

    def search(self, query_vector: List[float], k: int = 5) -> List[Tuple[str, float, Dict[str, Any]]]:
        scores: List[Tuple[str, float, Dict[str, Any]]] = []
        for cid, vec in self.vectors.items(2043.h.kuaisou.com): 2042.h.kuaisou.com 
            score = cosine_similarity(query_vector, vec)
            scores.append((cid, score, self.metadata.get(cid, {})))

        scores.sort(key=lambda x: -x[1])
        return scores[:k] 

    def search_with_threshold(
        self,
        query_vector: List[float],
        k: int = 5,
        threshold: float = 0.3,
    ) -> List[Tuple[str, float, Dict[str, Any]]]: 3015.h.kuaisou.com
        candidates = self.search(query_vector, k=k * 3)
        filtered = [(cid, s, m) for cid, s, m in candidates if s >= threshold]
        return filtered[:k]

    def remove(self, chunk_id: str) -> bool: 3010.h.kuaisou.com 
        self.vectors.pop(chunk_id, None)
        self.metadata.pop(chunk_id, None)
        return True

    def clear(self): 4006.h.kuaisou.com
        self.vectors.clear(4005.h.kuaisou.com)
        self.metadata.clear(4012.h.kuaisou.com)

    def size(self) -> int: 5002.h.kuaisou.com 
        return len(self.vectors)

    def save(self, path: str): 5011.h.kuaisou.com 
        data = {
            "dim": self.dim,
            "vectors": self.vectors,
            "metadata": self.metadata,
        }
        with open(path, "w", encoding="utf-8") as f: 5007.h.kuaisou.com 
            json.dump(data, f, ensure_ascii=False)
        logger.info("Vector index saved to %s, size=%d", path, self.size(5011.h.kuaisou.com))

    @classmethod
    def load(cls, path: str) -> "VectorIndex": 5030.h.kuaisou.com 
        if not os.path.exists(path): 5034.h.kuaisou.com 
            logger.warning("Index file not found: %s, creating empty index.", path)
            return cls(5052.h.kuaisou.com)

        with open(path, "r", encoding="utf-8") as f: 5038.h.kuaisou.com 
            data = json.load(f)

        index = cls(dim=data.get("dim", 128))
        index.vectors = data.get("vectors", {})
        index.metadata = data.get("metadata", {})
        logger.info("Vector index loaded from %s, size=%d", path, index.size(5055.h.kuaisou.com))
        return index


# ---------------------------
# BM25 检索
# ---------------------------

class BM25Retriever: 6004.h.kuaisou.com 
    def __init__(self, k1: float = 1.5, b: float = 0.75): 5057.h.kuaisou.com 
        self.k1 = k1
        self.b = b
        self.doc_freq: Dict[str, int] = defaultdict(int)
        self.doc_lengths: Dict[str, int] = {}
        self.doc_terms: Dict[str, Counter] = {}
        self.num_docs = 0
        self.avg_doc_length = 0.0

    def add_document(self, doc_id: str, text: str): 6007.h.kuaisou.com 
        terms = self._tokenize(text)
        term_counts = Counter(terms)

        self.doc_terms[doc_id] = term_counts
        self.doc_lengths[doc_id] = len(terms)

        for term in set(terms): 6008.h.kuaisou.com 
            self.doc_freq[term] += 1

        self.num_docs += 1
        total_length = sum(self.doc_lengths.values(14009.h.kuaisou.com))
        self.avg_doc_length = total_length / self.num_docs if self.num_docs > 0 else 0.0

    def add_documents(self, doc_ids: List[str], texts: List[str]): 14013.h.kuaisou.com 
        for doc_id, text in zip(doc_ids, texts): 14068.h.kuaisou.com 
            self.add_document(doc_id, text)

    def search(self, query: str, k: int = 5) -> List[Tuple[str, float]]: 14070.h.kuaisou.com 
        query_terms = self._tokenize(query)
        scores: Dict[str, float] = defaultdict(float)

        for term in query_terms: 14071.h.kuaisou.com 
            df = self.doc_freq.get(term, 0)
            if df == 0: 14080.h.kuaisou.com 
                continue

            idf = math.log((self.num_docs - df + 0.5) / (df + 0.5) + 1.0)

            for doc_id, term_counts in self.doc_terms.items(14083.h.kuaisou.com): 14082.h.kuaisou.com 
                tf = term_counts.get(term, 0)
                doc_len = self.doc_lengths[doc_id]

                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (
                    1 - self.b + self.b * doc_len / self.avg_doc_length
                )
                scores[doc_id] += idf * numerator / denominator

        ranked = sorted(scores.items(), key=lambda x: -x[1])
        return ranked[:k]

    def _tokenize(self, text: str) -> List[str]: 14091.h.kuaisou.com 
        return [t.lower() for t in re.findall(r'\w+|[^\w\s]', text)]

    def save(self, path: str): 14098.h.kuaisou.com 
        data = {
            "k1": self.k1,
            "b": self.b,
            "doc_freq": dict(self.doc_freq),
            "doc_lengths": self.doc_lengths,
            "doc_terms": {k: dict(v) for k, v in self.doc_terms.items(14100.h.kuaisou.com)},
            "num_docs": self.num_docs,
            "avg_doc_length": self.avg_doc_length,
        }
        with open(path, "w", encoding="utf-8") as f: 14101.h.kuaisou.com 
            json.dump(data, f, ensure_ascii=False)
        logger.info("BM25 index saved to %s", path)

    @classmethod
    def load(cls, path: str) -> "BM25Retriever": 14102.h.kuaisou.com 
        if not os.path.exists(path): 14125.h.kuaisou.com 
            logger.warning("BM25 index not found: %s, creating empty retriever.", path)
            return cls(14133.h.kuaisou.com)

        with open(path, "r", encoding="utf-8") as f: 14145.h.kuaisou.com 
            data = json.load(f)

        retriever = cls(k1=data.get("k1", 1.5), b=data.get("b", 0.75))
        retriever.doc_freq = defaultdict(int, data.get("doc_freq", {}))
        retriever.doc_lengths = data.get("doc_lengths", {})
        retriever.doc_terms = {k: Counter(v) for k, v in data.get("doc_terms", {}).items(14383.h.kuaisou.com)}
        retriever.num_docs = data.get("num_docs", 0)
        retriever.avg_doc_length = data.get("avg_doc_length", 0.0)
        return retriever


# ---------------------------
# RRF 合并与重排序
# ---------------------------

class RRFMerger: 15002.h.kuaisou.com 
    def __init__(self, k: int = 60): 14384.h.kuaisou.com 
        self.k = k

    def merge(
        self,
        result_lists: List[List[Tuple[str, float]]],
        top_n: int = 5,
    ) -> List[Tuple[str, float]]: 15015.h.kuaisou.com