coder
8c7a99d83f
新增: - 系统设置页面 (/settings) - 支持动态配置LLM、索引、文档处理参数 - 配置API - 保存配置、测试LLM连接 - 前端JS交互文件 - 搜索、文档管理功能 修复: - 首页搜索框无法正常工作的问题（缺少main.js） - 服务支持动态读取配置（无需重启生效） 改进: - LLM/索引/文档配置支持热更新 - 添加测试LLM连接功能
729 lines
22 KiB
Python
729 lines
22 KiB
Python
"""
|
|
文档索引服务
|
|
使用LLM分析文档并构建索引
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
import json
|
|
import math
|
|
from datetime import datetime
|
|
from collections import Counter
|
|
from openai import OpenAI
|
|
from flask import current_app
|
|
|
|
from config import LLM_CONFIG, DOC_CONFIG, INDEX_CONFIG
|
|
from models import db, Document, DocumentChunk, InvertedIndex, IndexStats
|
|
|
|
|
|
def get_llm_config():
|
|
"""获取有效的LLM配置（支持动态更新）"""
|
|
user_config_file = os.path.join(os.path.dirname(__file__), 'user_config.json')
|
|
if os.path.exists(user_config_file):
|
|
with open(user_config_file, 'r', encoding='utf-8') as f:
|
|
user_config = json.load(f)
|
|
return {**LLM_CONFIG, **user_config.get('llm', {})}
|
|
return LLM_CONFIG
|
|
|
|
|
|
def get_doc_config():
|
|
"""获取有效的文档配置"""
|
|
user_config_file = os.path.join(os.path.dirname(__file__), 'user_config.json')
|
|
if os.path.exists(user_config_file):
|
|
with open(user_config_file, 'r', encoding='utf-8') as f:
|
|
user_config = json.load(f)
|
|
return {**DOC_CONFIG, **user_config.get('doc', {})}
|
|
return DOC_CONFIG
|
|
|
|
|
|
def get_index_config():
|
|
"""获取有效的索引配置"""
|
|
user_config_file = os.path.join(os.path.dirname(__file__), 'user_config.json')
|
|
if os.path.exists(user_config_file):
|
|
with open(user_config_file, 'r', encoding='utf-8') as f:
|
|
user_config = json.load(f)
|
|
return {**INDEX_CONFIG, **user_config.get('index', {})}
|
|
return INDEX_CONFIG
|
|
|
|
|
|
class LLMService:
|
|
"""LLM服务封装"""
|
|
|
|
def __init__(self):
|
|
pass # 不再在初始化时设置配置
|
|
|
|
def _get_client(self):
|
|
"""获取LLM客户端"""
|
|
config = get_llm_config()
|
|
return OpenAI(
|
|
api_key=config['api_key'],
|
|
base_url=config['api_base'],
|
|
)
|
|
|
|
def _get_config(self):
|
|
"""获取当前配置"""
|
|
return get_llm_config()
|
|
|
|
def analyze_document(self, content, title=None):
|
|
"""
|
|
分析文档，提取关键信息
|
|
|
|
Returns:
|
|
dict: {summary, keywords, topics, entities, category}
|
|
"""
|
|
prompt = f"""请分析以下文档内容，提取关键信息。
|
|
|
|
{'文档标题：' + title if title else ''}
|
|
|
|
文档内容（前3000字）：
|
|
{content[:3000]}
|
|
|
|
请以JSON格式返回以下信息：
|
|
{{
|
|
"summary": "文档摘要（100-200字）",
|
|
"keywords": ["关键词1", "关键词2", ...最多10个],
|
|
"topics": ["主题1", "主题2", ...最多5个],
|
|
"category": "文档分类",
|
|
"entities": {{
|
|
"persons": ["人名"],
|
|
"organizations": ["组织名"],
|
|
"locations": ["地点"],
|
|
"dates": ["日期"],
|
|
"others": ["其他实体"]
|
|
}}
|
|
}}
|
|
|
|
只返回JSON，不要其他内容。"""
|
|
|
|
try:
|
|
config = self._get_config()
|
|
client = self._get_client()
|
|
response = client.chat.completions.create(
|
|
model=config['model'],
|
|
messages=[{"role": "user", "content": prompt}],
|
|
max_tokens=1000,
|
|
temperature=0.3,
|
|
)
|
|
|
|
result = response.choices[0].message.content.strip()
|
|
# 清理可能的markdown标记
|
|
result = re.sub(r'^```json\s*', '', result)
|
|
result = re.sub(r'\s*```$', '', result)
|
|
|
|
return json.loads(result)
|
|
|
|
except Exception as e:
|
|
print(f"LLM分析失败: {e}")
|
|
return {
|
|
"summary": "",
|
|
"keywords": [],
|
|
"topics": [],
|
|
"category": "",
|
|
"entities": {}
|
|
}
|
|
|
|
def analyze_chunk(self, content):
|
|
"""
|
|
分析文档块，提取关键词
|
|
|
|
Returns:
|
|
dict: {summary, keywords, topics}
|
|
"""
|
|
prompt = f"""分析以下文本片段，提取关键信息。
|
|
|
|
文本：
|
|
{content[:1500]}
|
|
|
|
请以JSON格式返回：
|
|
{{
|
|
"summary": "片段摘要（50字以内）",
|
|
"keywords": ["关键词", ...最多8个],
|
|
"topics": ["主题", ...最多3个]
|
|
}}
|
|
|
|
只返回JSON。"""
|
|
|
|
try:
|
|
config = self._get_config()
|
|
client = self._get_client()
|
|
response = client.chat.completions.create(
|
|
model=config['model'],
|
|
messages=[{"role": "user", "content": prompt}],
|
|
max_tokens=500,
|
|
temperature=0.3,
|
|
)
|
|
|
|
result = response.choices[0].message.content.strip()
|
|
result = re.sub(r'^```json\s*', '', result)
|
|
result = re.sub(r'\s*```$', '', result)
|
|
|
|
return json.loads(result)
|
|
|
|
except Exception as e:
|
|
return {"summary": "", "keywords": [], "topics": []}
|
|
|
|
def process_query(self, query):
|
|
"""
|
|
处理查询，提取意图和关键词
|
|
|
|
Returns:
|
|
dict: {intent, keywords, expanded_terms, entities}
|
|
"""
|
|
prompt = f"""分析以下查询，提取搜索意图和关键词。
|
|
|
|
查询：{query}
|
|
|
|
请以JSON格式返回：
|
|
{{
|
|
"intent": "查询意图（如：查找信息、比较、解释、列表等）",
|
|
"keywords": ["主要关键词", ...最多10个],
|
|
"expanded_terms": ["同义词或相关词", ...最多5个],
|
|
"entities": {{
|
|
"persons": [],
|
|
"organizations": [],
|
|
"locations": [],
|
|
"dates": [],
|
|
"others": []
|
|
}}
|
|
}}
|
|
|
|
只返回JSON。"""
|
|
|
|
try:
|
|
config = self._get_config()
|
|
client = self._get_client()
|
|
response = client.chat.completions.create(
|
|
model=config['model'],
|
|
messages=[{"role": "user", "content": prompt}],
|
|
max_tokens=500,
|
|
temperature=0.3,
|
|
)
|
|
|
|
result = response.choices[0].message.content.strip()
|
|
result = re.sub(r'^```json\s*', '', result)
|
|
result = re.sub(r'\s*```$', '', result)
|
|
|
|
return json.loads(result)
|
|
|
|
except Exception as e:
|
|
return {
|
|
"intent": "search",
|
|
"keywords": query.split(),
|
|
"expanded_terms": [],
|
|
"entities": {}
|
|
}
|
|
|
|
|
|
class DocumentIndexer:
|
|
"""文档索引器"""
|
|
|
|
def __init__(self):
|
|
self.llm = LLMService()
|
|
|
|
def _get_chunk_config(self):
|
|
"""获取分块配置"""
|
|
config = get_doc_config()
|
|
return config['chunk_size'], config['chunk_overlap']
|
|
|
|
def index_document(self, doc_id):
|
|
"""
|
|
索引单个文档
|
|
|
|
Args:
|
|
doc_id: 文档ID
|
|
|
|
Returns:
|
|
bool: 是否成功
|
|
"""
|
|
doc = Document.query.get(doc_id)
|
|
if not doc:
|
|
return False
|
|
|
|
try:
|
|
doc.status = 'processing'
|
|
db.session.commit()
|
|
|
|
# 读取文档内容
|
|
content = self._read_document(doc.filepath)
|
|
if not content:
|
|
raise Exception("无法读取文档内容")
|
|
|
|
# 存储原文
|
|
doc.content = content
|
|
doc.word_count = len(content)
|
|
|
|
# 使用LLM分析整个文档
|
|
print(f" 正在分析文档: {doc.filename}")
|
|
analysis = self.llm.analyze_document(content, doc.title)
|
|
|
|
doc.summary = analysis.get('summary', '')
|
|
doc.set_keywords(analysis.get('keywords', []))
|
|
doc.set_topics(analysis.get('topics', []))
|
|
doc.category = analysis.get('category', '')
|
|
doc.set_entities(analysis.get('entities', {}))
|
|
|
|
# 分块处理
|
|
chunks = self._split_content(content)
|
|
doc.chunk_count = len(chunks)
|
|
|
|
# 清理旧分块
|
|
DocumentChunk.query.filter_by(document_id=doc.id).delete()
|
|
|
|
# 索引每个分块
|
|
for i, chunk_content in enumerate(chunks):
|
|
chunk = DocumentChunk(
|
|
document_id=doc.id,
|
|
chunk_index=i,
|
|
content=chunk_content,
|
|
start_char=0,
|
|
end_char=len(chunk_content)
|
|
)
|
|
|
|
# LLM分析分块
|
|
chunk_analysis = self.llm.analyze_chunk(chunk_content)
|
|
chunk.summary = chunk_analysis.get('summary', '')
|
|
chunk.set_keywords(chunk_analysis.get('keywords', []))
|
|
|
|
# 计算词频
|
|
term_freq = self._compute_term_freq(chunk_content)
|
|
chunk.set_term_freq(term_freq)
|
|
|
|
db.session.add(chunk)
|
|
|
|
db.session.commit()
|
|
|
|
# 更新倒排索引
|
|
self._update_inverted_index(doc.id)
|
|
|
|
# 标记完成
|
|
doc.status = 'indexed'
|
|
doc.indexed_at = datetime.utcnow()
|
|
db.session.commit()
|
|
|
|
# 更新统计
|
|
IndexStats.get_stats().update_stats()
|
|
|
|
print(f" ✓ 文档索引完成: {doc.filename}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
doc.status = 'failed'
|
|
doc.error_message = str(e)
|
|
db.session.commit()
|
|
print(f" ✗ 索引失败: {e}")
|
|
return False
|
|
|
|
def _read_document(self, filepath):
|
|
"""读取文档内容"""
|
|
ext = os.path.splitext(filepath)[1].lower()
|
|
|
|
# 尝试读取文本文件（包括没有扩展名的）
|
|
try:
|
|
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
|
|
content = f.read()
|
|
if content.strip(): # 如果能读取到内容
|
|
return content
|
|
except:
|
|
pass
|
|
|
|
# 按扩展名处理特定格式
|
|
if ext in ['.txt', '.md', '.json', '.html']:
|
|
with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
|
|
return f.read()
|
|
|
|
elif ext == '.pdf':
|
|
try:
|
|
from pypdf import PdfReader
|
|
reader = PdfReader(filepath)
|
|
text = ''
|
|
for page in reader.pages:
|
|
text += page.extract_text() + '\n'
|
|
return text
|
|
except:
|
|
pass
|
|
|
|
elif ext == '.docx':
|
|
try:
|
|
from docx import Document as DocxDocument
|
|
doc = DocxDocument(filepath)
|
|
return '\n'.join([p.text for p in doc.paragraphs])
|
|
except:
|
|
pass
|
|
|
|
# 最后尝试以二进制方式读取并解码
|
|
try:
|
|
with open(filepath, 'rb') as f:
|
|
content = f.read()
|
|
# 尝试多种编码
|
|
for encoding in ['utf-8', 'gbk', 'gb2312', 'latin-1']:
|
|
try:
|
|
return content.decode(encoding)
|
|
except:
|
|
continue
|
|
except:
|
|
pass
|
|
|
|
return None
|
|
|
|
def _split_content(self, content):
|
|
"""
|
|
分割内容为块
|
|
|
|
Args:
|
|
content: 文档内容
|
|
|
|
Returns:
|
|
list: 内容块列表
|
|
"""
|
|
chunk_size, _ = self._get_chunk_config()
|
|
chunks = []
|
|
|
|
# 按段落分割
|
|
paragraphs = content.split('\n\n')
|
|
|
|
current_chunk = ""
|
|
for para in paragraphs:
|
|
if len(current_chunk) + len(para) < chunk_size:
|
|
current_chunk += para + '\n\n'
|
|
else:
|
|
if current_chunk.strip():
|
|
chunks.append(current_chunk.strip())
|
|
current_chunk = para + '\n\n'
|
|
|
|
if current_chunk.strip():
|
|
chunks.append(current_chunk.strip())
|
|
|
|
return chunks if chunks else [content[:chunk_size]]
|
|
|
|
def _compute_term_freq(self, content):
|
|
"""计算词频"""
|
|
# 简单分词（中英文混合）
|
|
# 中文按字符，英文按空格
|
|
terms = []
|
|
|
|
# 提取中文词汇（简单按字，实际可用jieba）
|
|
chinese = re.findall(r'[\u4e00-\u9fff]+', content)
|
|
for text in chinese:
|
|
# 简单的双字词分割
|
|
if len(text) >= 2:
|
|
for i in range(len(text) - 1):
|
|
terms.append(text[i:i+2])
|
|
terms.extend(list(text))
|
|
|
|
# 提取英文单词
|
|
english = re.findall(r'[a-zA-Z]+', content.lower())
|
|
terms.extend(english)
|
|
|
|
# 统计词频
|
|
return dict(Counter(terms))
|
|
|
|
def _update_inverted_index(self, doc_id):
|
|
"""更新倒排索引"""
|
|
chunks = DocumentChunk.query.filter_by(document_id=doc_id).all()
|
|
|
|
# 收集所有词及其位置
|
|
term_postings = {}
|
|
|
|
for chunk in chunks:
|
|
# 从词频获取词
|
|
tf = chunk.get_term_freq()
|
|
|
|
# 从关键词获取词
|
|
keywords = chunk.get_keywords()
|
|
|
|
for kw in keywords:
|
|
if kw not in term_postings:
|
|
term_postings[kw] = []
|
|
term_postings[kw].append({
|
|
'doc_id': doc_id,
|
|
'chunk_id': chunk.id,
|
|
'tf': tf.get(kw, 1),
|
|
'weight': INDEX_CONFIG['keyword_weight']
|
|
})
|
|
|
|
for term, freq in tf.items():
|
|
if term not in term_postings:
|
|
term_postings[term] = []
|
|
term_postings[term].append({
|
|
'doc_id': doc_id,
|
|
'chunk_id': chunk.id,
|
|
'tf': freq,
|
|
'weight': INDEX_CONFIG['content_weight']
|
|
})
|
|
|
|
# 更新数据库
|
|
for term, postings in term_postings.items():
|
|
index = InvertedIndex.get_or_create(term)
|
|
existing = index.get_postings()
|
|
|
|
# 合并postings（去除旧的同一文档的记录）
|
|
existing = [p for p in existing if p['doc_id'] != doc_id]
|
|
existing.extend(postings)
|
|
|
|
index.set_postings(existing)
|
|
|
|
db.session.commit()
|
|
|
|
|
|
class SearchEngine:
|
|
"""搜索引擎"""
|
|
|
|
def __init__(self):
|
|
self.llm = LLMService()
|
|
|
|
def _get_bm25_params(self):
|
|
"""获取BM25参数"""
|
|
config = get_index_config()
|
|
return config['bm25_k1'], config['bm25_b']
|
|
|
|
def search(self, query, top_k=10):
|
|
"""
|
|
搜索文档
|
|
|
|
Args:
|
|
query: 查询字符串
|
|
top_k: 返回结果数
|
|
|
|
Returns:
|
|
list: 搜索结果 [{doc, score, highlights}]
|
|
"""
|
|
start_time = datetime.now()
|
|
|
|
# 1. LLM处理查询
|
|
print(f"处理查询: {query}")
|
|
query_analysis = self.llm.process_query(query)
|
|
|
|
keywords = query_analysis.get('keywords', [])
|
|
expanded = query_analysis.get('expanded_terms', [])
|
|
|
|
# 合并关键词
|
|
all_terms = keywords + expanded
|
|
|
|
print(f" 关键词: {keywords}")
|
|
print(f" 扩展词: {expanded}")
|
|
|
|
# 2. 检索
|
|
results = self._retrieve(all_terms)
|
|
|
|
# 3. 计算BM25分数
|
|
scored_results = self._score_results(results, all_terms)
|
|
|
|
# 4. 排序
|
|
scored_results.sort(key=lambda x: x['score'], reverse=True)
|
|
|
|
# 5. 返回top_k
|
|
final_results = scored_results[:top_k]
|
|
|
|
retrieval_time = (datetime.now() - start_time).total_seconds()
|
|
|
|
# 6. 记录日志
|
|
self._log_query(query, query_analysis, final_results, retrieval_time)
|
|
|
|
return final_results
|
|
|
|
def _retrieve(self, terms):
|
|
"""
|
|
检索包含关键词的文档
|
|
|
|
Args:
|
|
terms: 关键词列表
|
|
|
|
Returns:
|
|
dict: {doc_id: {chunks: [], terms: []}}
|
|
"""
|
|
results = {}
|
|
|
|
for term in terms:
|
|
# 查询倒排索引
|
|
index = InvertedIndex.query.filter(
|
|
InvertedIndex.term.ilike(f'%{term}%')
|
|
).all()
|
|
|
|
for idx in index:
|
|
postings = idx.get_postings()
|
|
|
|
for p in postings:
|
|
doc_id = p['doc_id']
|
|
|
|
if doc_id not in results:
|
|
results[doc_id] = {
|
|
'chunks': set(),
|
|
'terms': {},
|
|
'postings': []
|
|
}
|
|
|
|
results[doc_id]['chunks'].add(p['chunk_id'])
|
|
results[doc_id]['terms'][term] = results[doc_id]['terms'].get(term, 0) + p.get('tf', 1)
|
|
results[doc_id]['postings'].append({
|
|
'term': term,
|
|
'chunk_id': p['chunk_id'],
|
|
'tf': p.get('tf', 1),
|
|
'weight': p.get('weight', 1.0)
|
|
})
|
|
|
|
return results
|
|
|
|
def _score_results(self, results, query_terms):
|
|
"""
|
|
使用BM25计算分数
|
|
|
|
Args:
|
|
results: 检索结果
|
|
query_terms: 查询词
|
|
|
|
Returns:
|
|
list: [{doc, score, chunks}]
|
|
"""
|
|
scored = []
|
|
|
|
# 计算平均文档长度
|
|
total_docs = Document.query.filter_by(status='indexed').count()
|
|
if total_docs == 0:
|
|
return []
|
|
|
|
avg_doc_len = db.session.query(
|
|
db.func.avg(Document.word_count)
|
|
).filter(Document.status == 'indexed').scalar() or 1000
|
|
|
|
for doc_id, data in results.items():
|
|
doc = Document.query.get(doc_id)
|
|
if not doc or doc.status != 'indexed':
|
|
continue
|
|
|
|
# BM25计算
|
|
k1, b = self._get_bm25_params()
|
|
score = 0
|
|
doc_len = doc.word_count or 1000
|
|
|
|
for term in query_terms:
|
|
# 查询倒排索引获取文档频率
|
|
index = InvertedIndex.query.filter_by(term=term).first()
|
|
df = index.doc_freq if index else 1
|
|
|
|
# IDF
|
|
idf = math.log((total_docs - df + 0.5) / (df + 0.5) + 1)
|
|
|
|
# TF
|
|
tf = data['terms'].get(term, 0)
|
|
|
|
# BM25公式
|
|
tf_component = (tf * (k1 + 1)) / (
|
|
tf + k1 * (1 - b + b * doc_len / avg_doc_len)
|
|
)
|
|
|
|
score += idf * tf_component
|
|
|
|
# 获取匹配的chunk内容
|
|
chunk_ids = list(data['chunks'])[:3] # 最多取3个chunk
|
|
chunks = DocumentChunk.query.filter(DocumentChunk.id.in_(chunk_ids)).all()
|
|
|
|
scored.append({
|
|
'doc': doc.to_dict(),
|
|
'score': score,
|
|
'matched_chunks': [c.to_dict() for c in chunks],
|
|
'matched_terms': list(data['terms'].keys())
|
|
})
|
|
|
|
return scored
|
|
|
|
def _log_query(self, query, analysis, results, retrieval_time):
|
|
"""记录查询日志"""
|
|
log = QueryLog(
|
|
original_query=query,
|
|
processed_query=' '.join(analysis.get('keywords', [])),
|
|
expanded_terms=json.dumps(analysis.get('expanded_terms', [])),
|
|
intent=analysis.get('intent'),
|
|
entities=json.dumps(analysis.get('entities', {})),
|
|
result_count=len(results),
|
|
top_doc_ids=json.dumps([r['doc']['id'] for r in results[:5]]),
|
|
retrieval_time=retrieval_time,
|
|
total_time=retrieval_time
|
|
)
|
|
db.session.add(log)
|
|
db.session.commit()
|
|
|
|
|
|
class RAGGenerator:
|
|
"""RAG生成器"""
|
|
|
|
def __init__(self):
|
|
self.llm = LLMService()
|
|
self.search_engine = SearchEngine()
|
|
|
|
def answer(self, query, top_k=5):
|
|
"""
|
|
RAG回答
|
|
|
|
Args:
|
|
query: 用户查询
|
|
top_k: 检索文档数
|
|
|
|
Returns:
|
|
dict: {answer, sources, confidence}
|
|
"""
|
|
# 1. 检索相关文档
|
|
results = self.search_engine.search(query, top_k)
|
|
|
|
if not results:
|
|
return {
|
|
'answer': '抱歉，没有找到相关信息。',
|
|
'sources': [],
|
|
'confidence': 0
|
|
}
|
|
|
|
# 2. 构建上下文
|
|
context_parts = []
|
|
sources = []
|
|
|
|
for i, r in enumerate(results[:3]): # 最多使用3个文档
|
|
doc = r['doc']
|
|
chunks = r['matched_chunks']
|
|
|
|
context_parts.append(f"【文档{i+1}】{doc.get('title', doc['filename'])}")
|
|
|
|
for chunk in chunks[:2]: # 每个文档最多2个chunk
|
|
context_parts.append(chunk.get('content', '')[:500])
|
|
|
|
sources.append({
|
|
'id': doc['id'],
|
|
'title': doc.get('title', doc['filename']),
|
|
'score': r['score']
|
|
})
|
|
|
|
context = '\n\n'.join(context_parts)
|
|
|
|
# 3. LLM生成回答
|
|
prompt = f"""基于以下参考信息回答问题。如果参考信息中没有相关内容，请说明。
|
|
|
|
问题：{query}
|
|
|
|
参考信息：
|
|
{context}
|
|
|
|
请给出准确、简洁的回答，并标注信息来源。"""
|
|
|
|
try:
|
|
llm_config = get_llm_config()
|
|
client = OpenAI(
|
|
api_key=llm_config['api_key'],
|
|
base_url=llm_config['api_base'],
|
|
)
|
|
|
|
response = client.chat.completions.create(
|
|
model=llm_config['model'],
|
|
messages=[{"role": "user", "content": prompt}],
|
|
max_tokens=1000,
|
|
temperature=0.5,
|
|
)
|
|
|
|
answer = response.choices[0].message.content
|
|
|
|
except Exception as e:
|
|
answer = f"生成回答时出错: {e}"
|
|
|
|
# 4. 返回结果
|
|
return {
|
|
'answer': answer,
|
|
'sources': sources,
|
|
'confidence': min(1.0, results[0]['score'] / 10) if results else 0
|
|
} |