530 lines
17 KiB
Python
530 lines
17 KiB
Python
|
"""
|
|||
|
|
RAG 知识库模块
|
|||
|
|
支持用户上传品牌文档，自动分块、索引，生成内容时自动检索相关内容
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import json
|
|||
|
|
import hashlib
|
|||
|
|
import logging
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import List, Dict, Optional, Any
|
|||
|
|
from datetime import datetime
|
|||
|
|
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DocumentChunk:
|
|||
|
|
"""文档分块"""
|
|||
|
|
|
|||
|
|
def __init__(self, content: str, metadata: Dict[str, Any]):
|
|||
|
|
self.content = content
|
|||
|
|
self.metadata = metadata
|
|||
|
|
self.chunk_id = hashlib.md5(content.encode()).hexdigest()[:12]
|
|||
|
|
|
|||
|
|
def to_dict(self) -> Dict:
|
|||
|
|
return {
|
|||
|
|
"chunk_id": self.chunk_id,
|
|||
|
|
"content": self.content,
|
|||
|
|
"metadata": self.metadata
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
@classmethod
|
|||
|
|
def from_dict(cls, data: Dict) -> 'DocumentChunk':
|
|||
|
|
chunk = cls(data["content"], data["metadata"])
|
|||
|
|
chunk.chunk_id = data["chunk_id"]
|
|||
|
|
return chunk
|
|||
|
|
|
|||
|
|
|
|||
|
|
class KnowledgeBase:
|
|||
|
|
"""知识库管理器"""
|
|||
|
|
|
|||
|
|
def __init__(self, storage_path: str = "knowledge_base"):
|
|||
|
|
"""
|
|||
|
|
Args:
|
|||
|
|
storage_path: 知识库存储路径
|
|||
|
|
"""
|
|||
|
|
self.storage_path = Path(storage_path)
|
|||
|
|
self.storage_path.mkdir(parents=True, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 文档元数据
|
|||
|
|
self.documents_file = self.storage_path / "documents.json"
|
|||
|
|
# 分块数据
|
|||
|
|
self.chunks_file = self.storage_path / "chunks.json"
|
|||
|
|
|
|||
|
|
self.documents: Dict[str, Dict] = self._load_json(self.documents_file, {})
|
|||
|
|
self.chunks: List[Dict] = self._load_json(self.chunks_file, [])
|
|||
|
|
|
|||
|
|
def _load_json(self, path: Path, default: Any) -> Any:
|
|||
|
|
"""加载 JSON 文件"""
|
|||
|
|
if path.exists():
|
|||
|
|
try:
|
|||
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|||
|
|
return json.load(f)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.warning(f"加载 {path} 失败: {e}")
|
|||
|
|
return default
|
|||
|
|
|
|||
|
|
def _save_json(self, path: Path, data: Any):
|
|||
|
|
"""保存 JSON 文件"""
|
|||
|
|
try:
|
|||
|
|
with open(path, 'w', encoding='utf-8') as f:
|
|||
|
|
json.dump(data, f, ensure_ascii=False, indent=2)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"保存 {path} 失败: {e}")
|
|||
|
|
|
|||
|
|
def add_document(self, filename: str, content: str, doc_type: str = "text",
|
|||
|
|
metadata: Optional[Dict] = None) -> Dict:
|
|||
|
|
"""
|
|||
|
|
添加文档到知识库
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
filename: 文件名
|
|||
|
|
content: 文档内容
|
|||
|
|
doc_type: 文档类型 (text, markdown, faq, case, product)
|
|||
|
|
metadata: 额外元数据
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
文档信息
|
|||
|
|
"""
|
|||
|
|
doc_id = hashlib.md5(f"{filename}{datetime.now().isoformat()}".encode()).hexdigest()[:12]
|
|||
|
|
|
|||
|
|
doc_info = {
|
|||
|
|
"doc_id": doc_id,
|
|||
|
|
"filename": filename,
|
|||
|
|
"doc_type": doc_type,
|
|||
|
|
"content_length": len(content),
|
|||
|
|
"chunk_count": 0,
|
|||
|
|
"created_at": datetime.now().isoformat(),
|
|||
|
|
"metadata": metadata or {}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 分块
|
|||
|
|
chunks = self._split_document(content, doc_id, filename, doc_type)
|
|||
|
|
doc_info["chunk_count"] = len(chunks)
|
|||
|
|
|
|||
|
|
# 保存
|
|||
|
|
self.documents[doc_id] = doc_info
|
|||
|
|
self.chunks.extend([c.to_dict() for c in chunks])
|
|||
|
|
|
|||
|
|
self._save_json(self.documents_file, self.documents)
|
|||
|
|
self._save_json(self.chunks_file, self.chunks)
|
|||
|
|
|
|||
|
|
logger.info(f"文档 '{filename}' 已添加，分为 {len(chunks)} 个分块")
|
|||
|
|
return doc_info
|
|||
|
|
|
|||
|
|
def _split_document(self, content: str, doc_id: str, filename: str,
|
|||
|
|
doc_type: str, chunk_size: int = 500, overlap: int = 50) -> List[DocumentChunk]:
|
|||
|
|
"""
|
|||
|
|
将文档分割为多个分块
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
content: 文档内容
|
|||
|
|
doc_id: 文档 ID
|
|||
|
|
filename: 文件名
|
|||
|
|
doc_type: 文档类型
|
|||
|
|
chunk_size: 分块大小（字符数）
|
|||
|
|
overlap: 重叠字符数
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
分块列表
|
|||
|
|
"""
|
|||
|
|
chunks = []
|
|||
|
|
|
|||
|
|
# 根据文档类型选择分块策略
|
|||
|
|
if doc_type == "faq":
|
|||
|
|
# FAQ 文档按 Q&A 对分块
|
|||
|
|
chunks = self._split_faq(content, doc_id, filename)
|
|||
|
|
elif doc_type == "product":
|
|||
|
|
# 产品文档按功能/特性分块
|
|||
|
|
chunks = self._split_by_sections(content, doc_id, filename, doc_type)
|
|||
|
|
else:
|
|||
|
|
# 通用文档按段落/长度分块
|
|||
|
|
chunks = self._split_by_length(content, doc_id, filename, doc_type,
|
|||
|
|
chunk_size, overlap)
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def _split_faq(self, content: str, doc_id: str, filename: str) -> List[DocumentChunk]:
|
|||
|
|
"""FAQ 文档分块：每个 Q&A 对为一个分块"""
|
|||
|
|
chunks = []
|
|||
|
|
lines = content.split('\n')
|
|||
|
|
|
|||
|
|
current_q = ""
|
|||
|
|
current_a = ""
|
|||
|
|
|
|||
|
|
for line in lines:
|
|||
|
|
line = line.strip()
|
|||
|
|
if not line:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 检测问题行
|
|||
|
|
if line.startswith('Q:') or line.startswith('问：') or line.startswith('Q：'):
|
|||
|
|
# 保存上一个 Q&A 对
|
|||
|
|
if current_q and current_a:
|
|||
|
|
chunk_content = f"问题：{current_q}\n回答：{current_a}"
|
|||
|
|
chunks.append(DocumentChunk(
|
|||
|
|
content=chunk_content,
|
|||
|
|
metadata={
|
|||
|
|
"doc_id": doc_id,
|
|||
|
|
"filename": filename,
|
|||
|
|
"type": "faq",
|
|||
|
|
"question": current_q
|
|||
|
|
}
|
|||
|
|
))
|
|||
|
|
current_q = line[2:].strip()
|
|||
|
|
current_a = ""
|
|||
|
|
elif line.startswith('A:') or line.startswith('答：') or line.startswith('A：'):
|
|||
|
|
current_a = line[2:].strip()
|
|||
|
|
elif current_a:
|
|||
|
|
current_a += "\n" + line
|
|||
|
|
elif current_q and not current_a:
|
|||
|
|
current_q += "\n" + line
|
|||
|
|
|
|||
|
|
# 保存最后一个 Q&A 对
|
|||
|
|
if current_q and current_a:
|
|||
|
|
chunk_content = f"问题：{current_q}\n回答：{current_a}"
|
|||
|
|
chunks.append(DocumentChunk(
|
|||
|
|
content=chunk_content,
|
|||
|
|
metadata={
|
|||
|
|
"doc_id": doc_id,
|
|||
|
|
"filename": filename,
|
|||
|
|
"type": "faq",
|
|||
|
|
"question": current_q
|
|||
|
|
}
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def _split_by_sections(self, content: str, doc_id: str, filename: str,
|
|||
|
|
doc_type: str) -> List[DocumentChunk]:
|
|||
|
|
"""按章节分块（适用于产品文档、Markdown 等）"""
|
|||
|
|
chunks = []
|
|||
|
|
sections = content.split('\n# ')
|
|||
|
|
|
|||
|
|
for i, section in enumerate(sections):
|
|||
|
|
if not section.strip():
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
# 提取标题
|
|||
|
|
lines = section.split('\n', 1)
|
|||
|
|
title = lines[0].strip('# ').strip()
|
|||
|
|
body = lines[1].strip() if len(lines) > 1 else ""
|
|||
|
|
|
|||
|
|
if body:
|
|||
|
|
chunks.append(DocumentChunk(
|
|||
|
|
content=f"## {title}\n{body}",
|
|||
|
|
metadata={
|
|||
|
|
"doc_id": doc_id,
|
|||
|
|
"filename": filename,
|
|||
|
|
"type": doc_type,
|
|||
|
|
"section_title": title
|
|||
|
|
}
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def _split_by_length(self, content: str, doc_id: str, filename: str,
|
|||
|
|
doc_type: str, chunk_size: int, overlap: int) -> List[DocumentChunk]:
|
|||
|
|
"""按长度分块"""
|
|||
|
|
chunks = []
|
|||
|
|
paragraphs = content.split('\n\n')
|
|||
|
|
|
|||
|
|
current_chunk = ""
|
|||
|
|
for para in paragraphs:
|
|||
|
|
para = para.strip()
|
|||
|
|
if not para:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
if len(current_chunk) + len(para) > chunk_size and current_chunk:
|
|||
|
|
chunks.append(DocumentChunk(
|
|||
|
|
content=current_chunk,
|
|||
|
|
metadata={
|
|||
|
|
"doc_id": doc_id,
|
|||
|
|
"filename": filename,
|
|||
|
|
"type": doc_type
|
|||
|
|
}
|
|||
|
|
))
|
|||
|
|
# 保留重叠部分
|
|||
|
|
if overlap > 0:
|
|||
|
|
current_chunk = current_chunk[-overlap:] + "\n" + para
|
|||
|
|
else:
|
|||
|
|
current_chunk = para
|
|||
|
|
else:
|
|||
|
|
current_chunk = current_chunk + "\n\n" + para if current_chunk else para
|
|||
|
|
|
|||
|
|
# 保存最后一个分块
|
|||
|
|
if current_chunk.strip():
|
|||
|
|
chunks.append(DocumentChunk(
|
|||
|
|
content=current_chunk,
|
|||
|
|
metadata={
|
|||
|
|
"doc_id": doc_id,
|
|||
|
|
"filename": filename,
|
|||
|
|
"type": doc_type
|
|||
|
|
}
|
|||
|
|
))
|
|||
|
|
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def search(self, query: str, top_k: int = 5,
|
|||
|
|
doc_type: Optional[str] = None) -> List[Dict]:
|
|||
|
|
"""
|
|||
|
|
搜索知识库
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
query: 查询文本
|
|||
|
|
top_k: 返回结果数量
|
|||
|
|
doc_type: 过滤文档类型
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
相关分块列表
|
|||
|
|
"""
|
|||
|
|
if not self.chunks:
|
|||
|
|
return []
|
|||
|
|
|
|||
|
|
# 计算相似度分数
|
|||
|
|
scored_chunks = []
|
|||
|
|
query_lower = query.lower()
|
|||
|
|
query_keywords = set(query_lower.split())
|
|||
|
|
|
|||
|
|
for chunk_data in self.chunks:
|
|||
|
|
content_lower = chunk_data["content"].lower()
|
|||
|
|
|
|||
|
|
# 计算关键词匹配分数
|
|||
|
|
keyword_matches = sum(1 for kw in query_keywords if kw in content_lower)
|
|||
|
|
keyword_score = keyword_matches / len(query_keywords) if query_keywords else 0
|
|||
|
|
|
|||
|
|
# 计算内容相关性分数（包含查询词的比例）
|
|||
|
|
content_score = 0
|
|||
|
|
for kw in query_keywords:
|
|||
|
|
if kw in content_lower:
|
|||
|
|
# 计算关键词在内容中的密度
|
|||
|
|
count = content_lower.count(kw)
|
|||
|
|
content_score += count * len(kw) / len(content_lower)
|
|||
|
|
|
|||
|
|
# 综合分数
|
|||
|
|
total_score = keyword_score * 0.6 + content_score * 0.4
|
|||
|
|
|
|||
|
|
if total_score > 0:
|
|||
|
|
# 过滤文档类型
|
|||
|
|
if doc_type and chunk_data["metadata"].get("type") != doc_type:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
scored_chunks.append({
|
|||
|
|
"chunk": chunk_data,
|
|||
|
|
"score": total_score
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# 按分数排序并返回 top_k
|
|||
|
|
scored_chunks.sort(key=lambda x: x["score"], reverse=True)
|
|||
|
|
|
|||
|
|
return [
|
|||
|
|
{
|
|||
|
|
"content": item["chunk"]["content"],
|
|||
|
|
"metadata": item["chunk"]["metadata"],
|
|||
|
|
"score": item["score"]
|
|||
|
|
}
|
|||
|
|
for item in scored_chunks[:top_k]
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
def get_context_for_generation(self, query: str, brand: str,
|
|||
|
|
platform: str, top_k: int = 3) -> str:
|
|||
|
|
"""
|
|||
|
|
获取用于内容生成的上下文
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
query: 查询/主题
|
|||
|
|
brand: 品牌名
|
|||
|
|
platform: 目标平台
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
格式化的上下文字符串
|
|||
|
|
"""
|
|||
|
|
# 搜索相关文档
|
|||
|
|
results = self.search(query, top_k=top_k)
|
|||
|
|
|
|||
|
|
if not results:
|
|||
|
|
return ""
|
|||
|
|
|
|||
|
|
# 组装上下文
|
|||
|
|
context_parts = ["以下是相关的品牌/产品信息，可用于内容生成：\n"]
|
|||
|
|
|
|||
|
|
for i, result in enumerate(results, 1):
|
|||
|
|
source = result["metadata"].get("filename", "未知来源")
|
|||
|
|
content = result["content"]
|
|||
|
|
|
|||
|
|
context_parts.append(f"--- 参考资料 {i}（来源：{source}）---")
|
|||
|
|
context_parts.append(content)
|
|||
|
|
context_parts.append("")
|
|||
|
|
|
|||
|
|
return "\n".join(context_parts)
|
|||
|
|
|
|||
|
|
def list_documents(self) -> List[Dict]:
|
|||
|
|
"""列出所有文档"""
|
|||
|
|
return list(self.documents.values())
|
|||
|
|
|
|||
|
|
def delete_document(self, doc_id: str) -> bool:
|
|||
|
|
"""删除文档及其分块"""
|
|||
|
|
if doc_id not in self.documents:
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
# 删除文档
|
|||
|
|
del self.documents[doc_id]
|
|||
|
|
|
|||
|
|
# 删除相关分块
|
|||
|
|
self.chunks = [c for c in self.chunks if c["metadata"].get("doc_id") != doc_id]
|
|||
|
|
|
|||
|
|
# 保存
|
|||
|
|
self._save_json(self.documents_file, self.documents)
|
|||
|
|
self._save_json(self.chunks_file, self.chunks)
|
|||
|
|
|
|||
|
|
return True
|
|||
|
|
|
|||
|
|
def get_stats(self) -> Dict:
|
|||
|
|
"""获取知识库统计信息"""
|
|||
|
|
doc_types = {}
|
|||
|
|
for doc in self.documents.values():
|
|||
|
|
doc_type = doc.get("doc_type", "unknown")
|
|||
|
|
doc_types[doc_type] = doc_types.get(doc_type, 0) + 1
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"total_documents": len(self.documents),
|
|||
|
|
"total_chunks": len(self.chunks),
|
|||
|
|
"document_types": doc_types
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
class SourceVerifier:
|
|||
|
|
"""来源验证器"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self.claim_patterns = [
|
|||
|
|
"根据",
|
|||
|
|
"据",
|
|||
|
|
"报告显示",
|
|||
|
|
"数据表明",
|
|||
|
|
"研究表明",
|
|||
|
|
"调查发现",
|
|||
|
|
"据统计",
|
|||
|
|
"根据报告",
|
|||
|
|
"根据数据",
|
|||
|
|
"根据研究",
|
|||
|
|
"according to",
|
|||
|
|
"based on",
|
|||
|
|
"as reported by",
|
|||
|
|
"research shows",
|
|||
|
|
"data shows"
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
def extract_claims(self, content: str) -> List[Dict]:
|
|||
|
|
"""
|
|||
|
|
从内容中提取来源声明
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
content: 文本内容
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
声明列表
|
|||
|
|
"""
|
|||
|
|
claims = []
|
|||
|
|
sentences = content.replace('。', '。\n').replace('.', '.\n').split('\n')
|
|||
|
|
|
|||
|
|
for sentence in sentences:
|
|||
|
|
sentence = sentence.strip()
|
|||
|
|
if not sentence:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
for pattern in self.claim_patterns:
|
|||
|
|
if pattern in sentence.lower():
|
|||
|
|
claims.append({
|
|||
|
|
"text": sentence,
|
|||
|
|
"pattern": pattern,
|
|||
|
|
"verified": False,
|
|||
|
|
"verification_result": None
|
|||
|
|
})
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
return claims
|
|||
|
|
|
|||
|
|
def generate_verification_prompt(self, claim: str) -> str:
|
|||
|
|
"""
|
|||
|
|
生成验证提示词
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
claim: 来源声明
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
验证提示词
|
|||
|
|
"""
|
|||
|
|
return f"""请验证以下声明的真实性：
|
|||
|
|
|
|||
|
|
声明：{claim}
|
|||
|
|
|
|||
|
|
请回答：
|
|||
|
|
1. 这个声明是否包含可验证的具体来源（如具体报告名称、机构名称、数据年份）？
|
|||
|
|
2. 如果包含，请判断这个来源是否可能存在且可信。
|
|||
|
|
3. 如果无法验证或来源可疑，请说明原因。
|
|||
|
|
|
|||
|
|
回答格式：
|
|||
|
|
- 来源具体性：[具体/模糊/无来源]
|
|||
|
|
- 可信度评估：[高/中/低/无法判断]
|
|||
|
|
- 建议：[保留/修改/删除]
|
|||
|
|
- 原因：[简要说明]"""
|
|||
|
|
|
|||
|
|
def assess_source_quality(self, content: str) -> Dict:
|
|||
|
|
"""
|
|||
|
|
评估内容的来源质量
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
content: 文本内容
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
质量评估结果
|
|||
|
|
"""
|
|||
|
|
claims = self.extract_claims(content)
|
|||
|
|
|
|||
|
|
if not claims:
|
|||
|
|
return {
|
|||
|
|
"has_sources": False,
|
|||
|
|
"claim_count": 0,
|
|||
|
|
"quality_score": 0,
|
|||
|
|
"suggestions": ["内容中没有引用任何来源，建议添加数据支撑"]
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 分析来源质量
|
|||
|
|
specific_count = 0
|
|||
|
|
vague_count = 0
|
|||
|
|
|
|||
|
|
for claim in claims:
|
|||
|
|
text = claim["text"]
|
|||
|
|
# 检查是否有具体来源指标
|
|||
|
|
has_specific = any([
|
|||
|
|
any(year in text for year in ["2020", "2021", "2022", "2023", "2024", "2025"]),
|
|||
|
|
any(org in text for org in ["Gartner", "IDC", "Forrester", "McKinsey",
|
|||
|
|
"哈佛", "MIT", "斯坦福", "中科院"]),
|
|||
|
|
"报告" in text and ("《" in text or "年" in text),
|
|||
|
|
"数据" in text and any(c.isdigit() for c in text)
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
if has_specific:
|
|||
|
|
specific_count += 1
|
|||
|
|
else:
|
|||
|
|
vague_count += 1
|
|||
|
|
|
|||
|
|
quality_score = min(100, (specific_count / len(claims)) * 100)
|
|||
|
|
|
|||
|
|
suggestions = []
|
|||
|
|
if vague_count > 0:
|
|||
|
|
suggestions.append(f"有 {vague_count} 个来源描述模糊，建议补充具体报告名称或数据年份")
|
|||
|
|
if specific_count == 0:
|
|||
|
|
suggestions.append("所有来源都不够具体，建议引用真实的行业报告或权威机构数据")
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
"has_sources": True,
|
|||
|
|
"claim_count": len(claims),
|
|||
|
|
"specific_count": specific_count,
|
|||
|
|
"vague_count": vague_count,
|
|||
|
|
"quality_score": quality_score,
|
|||
|
|
"claims": claims,
|
|||
|
|
"suggestions": suggestions
|
|||
|
|
}
|