import os
import re
import json
import math
import time
import hashlib
import logging
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional, Any
from collections import Counter, defaultdict
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("rag_system")class Document: 15015.h.kuaisou.com
"""表示一个待处理的文档对象。"""
def __init__(
self,
doc_id: str,
text: str,
source: str = "",
metadata: Optional[Dict[str, Any]] = None,
):
self.doc_id = doc_id
self.text = text
self.source = source
self.metadata = metadata or {}
def __repr__(self) -> str: 15025.h.kuaisou.com
return f"Document(doc_id={self.doc_id!r}, source={self.source!r}, len={len(self.text)})"
class DocumentLoader: 15027.h.kuaisou.com
"""
文档加载器,支持从文本字符串、JSON、TXT 文件加载文档。
实际项目中可以继续扩展 PDF、Markdown、Word 等格式。
"""
def load_from_text(self, text: str, doc_id: str = "doc_0", source: str = "") -> Document: 15029.h.kuaisou.com
return Document(doc_id=doc_id, text=text, source=source)
def load_from_json(self, file_path: str) -> List[Document]: 15030.h.kuaisou.com
if not os.path.exists(file_path): 15031.h.kuaisou.com
raise FileNotFoundError(f"JSON file not found: {file_path}")
with open(file_path, "r", encoding="utf-8") as f: 15032.h.kuaisou.com
data = json.load(f)
documents = []
for i, item in enumerate(data): 15033.h.kuaisou.com
doc = Document(
doc_id=item.get("doc_id", f"doc_{i}"),
text=item.get("text", ""),
source=item.get("source", file_path),
metadata=item.get("metadata", {}),
)
documents.append(doc)
return documents
def load_from_txt(self, file_path: str, doc_id: str = "doc_0") -> Document: 15034.h.kuaisou.com
if not os.path.exists(file_path): 15035.h.kuaisou.com
raise FileNotFoundError(f"TXT file not found: {file_path}")
with open(file_path, "r", encoding="utf-8") as f: 15036.h.kuaisou.com
return Document(
doc_id=doc_id,
text=text,
source=file_path,
metadata={"file_type": "txt"},
)@dataclass
class TextChunk: 15037.h.kuaisou.com
"""文本块对象,用于后续 Embedding 和检索。"""
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: 15038.h.kuaisou.com
return len(self.text)
class RecursiveCharacterTextSplitter: 15039.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,
):
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]: 16001.h.kuaisou.com
return self._split_recursive(text, self.separators)
def _split_recursive(self, text: str, separators: List[str]) -> List[str]: 16002.h.kuaisou.com
if not text: 16003.h.kuaisou.com
return []
separator = separators[-1]
new_separators = []
for i, sep in enumerate(separators): 16004.h.kuaisou.com
if sep == "" or sep in text: 16005.h.kuaisou.com
separator = sep
new_separators = separators[i + 1:]
break
if separator: 16006.h.kuaisou.com
splits = text.split(separator)
else: 16007.h.kuaisou.com
splits = list(text)
chunks = []
current_chunk: List[str] = []
current_length = 0
for split in splits: 16008.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: 16009.h.kuaisou.com
current_chunk.append(split)
current_length += split_len + sep_len
else: 16010.h.kuaisou.com
if current_chunk: 16011.h.kuaisou.com
chunk_text = separator.join(current_chunk)
if chunk_text: 16012.h.kuaisou.com
chunks.append(chunk_text)
overlap_text = self._build_overlap(current_chunk, separator)
if overlap_text: 16013.h.kuaisou.com
current_chunk = [overlap_text, split]
current_length = self.length_function(overlap_text) + split_len + sep_len
else: 16014.h.kuaisou.com
current_chunk = [split]
current_length = split_len + sep_len
else: 16015.h.kuaisou.com
if split_len > self.chunk_size and new_separators: 16016.h.kuaisou.com
sub_chunks = self._split_recursive(split, new_separators)
chunks.extend(sub_chunks)
else: 16017.h.kuaisou.com
chunks.append(split)
current_chunk = []
current_length = 0
if current_chunk: 16018.h.kuaisou.com
chunk_text = separator.join(current_chunk)
if chunk_text: 16019.h.kuaisou.com
chunks.append(chunk_text)
return chunks
def _build_overlap(self, parts: List[str], separator: str) -> str: 16020.h.kuaisou.com
if not parts or self.chunk_overlap <= 0: 16021.h.kuaisou.com
return ""
overlap_parts = []
overlap_len = 0
for part in reversed(parts): 16022.h.kuaisou.com
part_len = self.length_function(part)
if overlap_len + part_len <= self.chunk_overlap: 16023.h.kuaisou.com
overlap_parts.insert(0, part)
overlap_len += part_len
else: 16024.h.kuaisou.com
remaining = self.chunk_overlap - overlap_len
if remaining > 0: 16025.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 = []
pos = 0
for i, chunk_text in enumerate(raw_chunks): 16026.h.kuaisou.com
start = text.find(chunk_text, pos)
if start == -1: 16027.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 chunksclass SimpleEmbeddingModel: 16028.h.kuaisou.com
"""
简易 Embedding 模型。
生产环境应替换为 OpenAI、BGE、SentenceTransformer 等真实向量模型。
"""
def __init__(self, dim: int = 128): 16029.h.kuaisou.com
self.dim = dim
def _hash_word(self, word: str) -> int: 16030.h.kuaisou.com
return int(hashlib.md5(word.lower(16032.h.kuaisou.com).encode(16033.h.kuaisou.com)).hexdigest(), 16)
def embed(self, text: str) -> List[float]: 16031.h.kuaisou.com
words = re.findall(r'\w+|[^\w\s]', text.lower(16034.h.kuaisou.com))
vector = [0.0] * self.dim
for word in words: 16035.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: 16036.h.kuaisou.com
vector = [v / norm for v in vector]
return vector
def embed_batch(self, texts: List[str]) -> List[List[float]]: 16037.h.kuaisou.com
return [self.embed(t) for t in texts]
def cosine_similarity(a: List[float], b: List[float]) -> float: 16038.h.kuaisou.com
if len(a) != len(b): 16039.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: 16040.h.kuaisou.com
return 0.0
return dot / (norm_a * norm_b)
class VectorIndex: 16041.h.kuaisou.com
"""内存向量索引,支持添加、检索、删除。"""
def __init__(self, dim: int = 128): 16042.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: 16043.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): 16044.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 = []
for cid, vec in self.vectors.items(16046.h.kuaisou.com): 16045.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]]]: 16047.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: 16048.h.kuaisou.com
self.vectors.pop(chunk_id, None)
self.metadata.pop(chunk_id, None)
return True
def clear(self): 16049.h.kuaisou.com
self.vectors.clear(16050.h.kuaisou.com)
self.metadata.clear(16051.h.kuaisou.com)
def size(self) -> int: 16052.h.kuaisou.com
return len(self.vectors)class BM25Retriever: 16053.h.kuaisou.com
"""
BM25 关键词检索器。
用于补充向量检索,提升精确词匹配能力。
"""
def __init__(self, k1: float = 1.5, b: float = 0.75): 16054.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): 16055.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): 16056.h.kuaisou.com
self.doc_freq[term] += 1
self.num_docs += 1
total_length = sum(self.doc_lengths.values(16057.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]): 16059.h.kuaisou.com
for doc_id, text in zip(doc_ids, texts): 16058.h.kuaisou.com
self.add_document(doc_id, text)
def search(self, query: str, k: int = 5) -> List[Tuple[str, float]]: 16060.h.kuaisou.com
query_terms = self._tokenize(query)
scores: Dict[str, float] = defaultdict(float)
for term in query_terms: 16061.h.kuaisou.com
df = self.doc_freq.get(term, 0)
if df == 0: 16062.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(16064.h.kuaisou.com): 16063.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(17001.h.kuaisou.com), key=lambda x: -x[1])
return ranked[:k]
def _tokenize(self, text: str) -> List[str]: 17002.h.kuaisou.com
return [t.lower() for t in re.findall(r'\w+|[^\w\s]', text)]class RRFMerger: 17003.h.kuaisou.com
"""
RRF:Reciprocal Rank Fusion。
用于合并向量检索和 BM25 检索结果。
"""
def __init__(self, k: int = 60):
self.k = k
def merge(
self,
result_lists: List[List[Tuple[str, float]]],
top_n: int = 5,
) -> List[Tuple[str, float]]: 17004.h.kuaisou.com
rrf_scores: Dict[str, float] = {}
for result_list in result_lists: 17005.h.kuaisou.com
for rank, (chunk_id, _) in enumerate(result_list, start=1): 17006.h.kuaisou.com
rrf_scores[chunk_id] = rrf_scores.get(chunk_id, 0.0) + 1.0 / (self.k + rank)
merged = sorted(rrf_scores.items(17007.h.kuaisou.com), key=lambda x: -x[1])
return merged[:top_n]
class CrossEncoderReranker: 17008.h.kuaisou.com
"""
简易重排序器。
生产环境可替换为 BGE-reranker、Cohere Rerank 等。
"""
def __init__(self, embedding_model: SimpleEmbeddingModel): 17009.h.kuaisou.com
self.embedding_model = embedding_model
def rerank(
self,
query: str,
candidates: List[Tuple[str, str, float]],
top_n: int = 5,
) -> List[Tuple[str, str, float]]: 17010.h.kuaisou.com
query_vec = self.embedding_model.embed(query)
scored = []
for chunk_id, chunk_text, init_score in candidates: 17011.h.kuaisou.com
chunk_vec = self.embedding_model.embed(chunk_text)
sim = cosine_similarity(query_vec, chunk_vec)
final_score = 0.4 * init_score + 0.6 * sim
scored.append((chunk_id, chunk_text, final_score))
scored.sort(key=lambda x: -x[2])
return scored[:top_n]
class HybridRetriever: 17012.h.kuaisou.com
"""
混合检索器:
1. 向量检索
2. BM25 检索
3. RRF 融合
4. 重排序
"""
def __init__(
self,
embedding_model: SimpleEmbeddingModel,
vector_index: VectorIndex,
bm25_retriever: BM25Retriever,
rrf_k: int = 60,
rerank_top_n: int = 5,
similarity_threshold: float = 0.3,
):
self.embedding_model = embedding_model
self.vector_index = vector_index
self.bm25_retriever = bm25_retriever
self.rrf_merger = RRFMerger(k=rrf_k)
self.reranker = CrossEncoderReranker(embedding_model)
self.rerank_top_n = rerank_top_n
self.similarity_threshold = similarity_threshold
def retrieve(self, query: str, k: int = 5) -> List[Tuple[str, str, float]]:
query_vec = self.embedding_model.embed(query)
vec_results = self.vector_index.search_with_threshold(
query_vec,
k=k * 3,
threshold=self.similarity_threshold,
)
vec_ranked = [(cid, score) for cid, score, _ in vec_results]
bm25_results = self.bm25_retriever.search(query, k=k * 3)
merged = self.rrf_merger.merge(
[vec_ranked, bm25_results],
top_n=k * 2,
)
candidates = []
for cid, rrf_score in merged: 17013.h.kuaisou.com
text = self.vector_index.metadata.get(cid, {}).get("text", "")
if text: 17014.h.kuaisou.com
candidates.append((cid, text, rrf_score))
reranked = self.reranker.rerank(query, candidates, top_n=self.rerank_top_n)
return rerankedclass RAGPromptBuilder: 17015.h.kuaisou.com
"""
构建防幻觉 RAG Prompt。
明确要求模型只基于上下文回答,避免编造。
"""
SYSTEM_PROMPT = """你是一个专业的知识库问答助手。请严格遵循以下规则:
1. 只根据【上下文】中的内容回答问题。
2. 如果上下文不足以回答问题,请回复:“根据现有知识库无法解答该问题。”
3. 不要编造上下文之外的信息。
4. 回答末尾标注引用来源。
5. 回答尽量简洁、准确、专业。"""
USER_TEMPLATE = """【上下文】
{context}
【用户问题】
{question}
请基于上述上下文回答问题:"""
def __init__(self, max_context_chunks: int = 5): 17016.h.kuaisou.com
self.max_context_chunks = max_context_chunks
def build_prompt(self, question: str, retrieved_chunks: List[Dict[str, str]]) -> str: 17017.h.kuaisou.com
context_parts = []
for i, chunk in enumerate(retrieved_chunks[: self.max_context_chunks], 1): 17018.h.kuaisou.com
context_parts.append(
f"[片段{i}] (来源: {chunk.get('source', 'unknown')}, "
f"ID: {chunk['chunk_id']})\n{chunk['text']}\n"
)
context = "\n".join(context_parts)
user_prompt = self.USER_TEMPLATE.format(context=context, question=question)
return self.SYSTEM_PROMPT + "\n\n" + user_prompt
class RAGGenerator: 17019.h.kuaisou.com
"""
RAG 答案生成器。
实际项目中应接入 OpenAI、DeepSeek、Qwen 等 LLM。
"""
def __init__(self, prompt_builder: Optional[RAGPromptBuilder] = None): 17020.h.kuaisou.com
self.prompt_builder = prompt_builder or RAGPromptBuilder(17021.h.kuaisou.com)
def generate(
self,
question: str,
retrieved_chunks: List[Dict[str, str]],
) -> Dict[str, Any]: 17022.h.kuaisou.com
if not retrieved_chunks: 17023.h.kuaisou.com
return {
"answer": "根据现有知识库无法解答该问题。",
"sources": [],
"question": question,
}
prompt = self.prompt_builder.build_prompt(question, retrieved_chunks)
answer = self._call_llm(prompt)
sources = [
{"chunk_id": c["chunk_id"], "source": c.get("source", "")}
for c in retrieved_chunks
]
return {
"answer": answer,
"sources": sources,
"question": question,
}
def _call_llm(self, prompt: str) -> str: 17024.h.kuaisou.com
"""
占位方法,实际项目中替换为真实 LLM 调用。
示例:OpenAI、DeepSeek、Qwen、本地部署模型。
"""
return f"[LLM模拟回答] 基于 {len(prompt)} 字符上下文生成的答案。"class RAGCache: 17121.h.kuaisou.com
"""
简单缓存层,用于缓存相同问题的检索和生成结果。
生产环境可替换为 Redis、Memcached 等。
"""
def __init__(self, ttl: int = 3600): 17025.h.kuaisou.com
self.cache: Dict[str, Dict[str, Any]] = {}
self.timestamps: Dict[str, float] = {}
self.ttl = ttl
def _make_key(self, question: str) -> str: 17026.h.kuaisou.com
return hashlib.sha256(question.strip().encode("utf-8")).hexdigest()
def get(self, question: 17117.h.kuaisou.com 原创声明:本文系作者授权腾讯云开发者社区发表,未经许可,不得转载。
如有侵权,请联系 cloudcommunity@tencent.com 删除。