init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
14
services/__init__.py
Normal file
14
services/__init__.py
Normal file
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
业务服务模块
|
||||
|
||||
包含：
|
||||
- session: 会话管理（开发环境）
|
||||
- feedback: 问答质量闭环（反馈、FAQ、质量报告）
|
||||
- outline: 纲要生成与关联推荐
|
||||
"""
|
||||
|
||||
from services.session import SessionManager
|
||||
|
||||
__all__ = [
|
||||
'SessionManager',
|
||||
]
|
||||
1314
services/feedback.py
Normal file
1314
services/feedback.py
Normal file
File diff suppressed because it is too large
Load Diff
986
services/outline.py
Normal file
986
services/outline.py
Normal file
@@ -0,0 +1,986 @@
|
||||
"""
|
||||
纲要生成与关联推荐服务
|
||||
|
||||
功能：
|
||||
1. GKPT-MIND-020 自动化纲要生成
|
||||
- AI提取制度文件章节结构
|
||||
- 生成思维导图数据
|
||||
- 支持多种格式导出
|
||||
|
||||
2. GKPT-READ-005 关联推荐
|
||||
- 基于向量相似度推荐相关文档
|
||||
- 支持标签匹配
|
||||
- 综合排序
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import hashlib
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import Optional, List, Dict, Any, Tuple
|
||||
from dataclasses import dataclass, asdict, field
|
||||
|
||||
from data.db import get_connection, init_databases
|
||||
from core.llm_utils import call_llm, parse_json_from_response
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ==================== 数据类定义 ====================
|
||||
|
||||
@dataclass
|
||||
class OutlineNode:
|
||||
"""纲要节点"""
|
||||
id: str = ""
|
||||
title: str = ""
|
||||
summary: str = ""
|
||||
level: int = 1
|
||||
order: int = 1
|
||||
page: int = 0
|
||||
children: List['OutlineNode'] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
"id": self.id,
|
||||
"title": self.title,
|
||||
"summary": self.summary,
|
||||
"level": self.level,
|
||||
"order": self.order,
|
||||
"page": self.page,
|
||||
"children": [child.to_dict() for child in self.children]
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict) -> 'OutlineNode':
|
||||
"""从字典创建"""
|
||||
node = cls(
|
||||
id=data.get("id", ""),
|
||||
title=data.get("title", ""),
|
||||
summary=data.get("summary", ""),
|
||||
level=data.get("level", 1),
|
||||
order=data.get("order", 1),
|
||||
page=data.get("page", 0)
|
||||
)
|
||||
for child_data in data.get("children", []):
|
||||
node.children.append(cls.from_dict(child_data))
|
||||
return node
|
||||
|
||||
|
||||
@dataclass
|
||||
class DocumentOutline:
|
||||
"""文档纲要"""
|
||||
document_id: str = ""
|
||||
document_name: str = ""
|
||||
total_pages: int = 0
|
||||
content_hash: str = ""
|
||||
generated_at: str = ""
|
||||
outline: List[OutlineNode] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
"""转换为字典"""
|
||||
return {
|
||||
"document_id": self.document_id,
|
||||
"document_name": self.document_name,
|
||||
"total_pages": self.total_pages,
|
||||
"content_hash": self.content_hash,
|
||||
"generated_at": self.generated_at,
|
||||
"outline": [node.to_dict() for node in self.outline]
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: Dict) -> 'DocumentOutline':
|
||||
"""从字典创建"""
|
||||
return cls(
|
||||
document_id=data.get("document_id", ""),
|
||||
document_name=data.get("document_name", ""),
|
||||
total_pages=data.get("total_pages", 0),
|
||||
content_hash=data.get("content_hash", ""),
|
||||
generated_at=data.get("generated_at", ""),
|
||||
outline=[OutlineNode.from_dict(n) for n in data.get("outline", [])]
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Recommendation:
|
||||
"""推荐结果"""
|
||||
document_id: str = ""
|
||||
document_name: str = ""
|
||||
summary: str = ""
|
||||
similarity: float = 0.0
|
||||
tag_score: float = 0.0
|
||||
final_score: float = 0.0
|
||||
tags: List[str] = field(default_factory=list)
|
||||
reason: str = "" # 推荐理由
|
||||
|
||||
def to_dict(self) -> Dict:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
# ==================== 数据库管理 ====================
|
||||
|
||||
class OutlineDB:
|
||||
"""纲要缓存数据库"""
|
||||
|
||||
def __init__(self):
|
||||
init_databases()
|
||||
self._init_db()
|
||||
|
||||
def _init_db(self):
|
||||
"""初始化数据库表"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 纲要缓存表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS outline_cache (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id TEXT NOT NULL UNIQUE,
|
||||
document_name TEXT,
|
||||
total_pages INTEGER DEFAULT 0,
|
||||
content_hash TEXT NOT NULL,
|
||||
outline_json TEXT NOT NULL,
|
||||
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
|
||||
# 文档向量缓存表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS document_vectors (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id TEXT NOT NULL UNIQUE,
|
||||
document_name TEXT,
|
||||
vector_hash TEXT,
|
||||
vector_json TEXT NOT NULL,
|
||||
tags_json TEXT,
|
||||
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
|
||||
# 推荐缓存表
|
||||
cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS recommendation_cache (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id TEXT NOT NULL,
|
||||
recommendations_json TEXT NOT NULL,
|
||||
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
|
||||
# 创建索引
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_outline_doc ON outline_cache(document_id)")
|
||||
cursor.execute("CREATE INDEX IF NOT EXISTS idx_vector_doc ON document_vectors(document_id)")
|
||||
|
||||
# ==================== 纲要缓存 ====================
|
||||
|
||||
def save_outline(self, outline: DocumentOutline) -> int:
|
||||
"""保存纲要"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
INSERT OR REPLACE INTO outline_cache
|
||||
(document_id, document_name, total_pages, content_hash, outline_json, generated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
outline.document_id,
|
||||
outline.document_name,
|
||||
outline.total_pages,
|
||||
outline.content_hash,
|
||||
json.dumps(outline.to_dict(), ensure_ascii=False),
|
||||
outline.generated_at
|
||||
))
|
||||
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_outline(self, document_id: str) -> Optional[DocumentOutline]:
|
||||
"""获取纲要"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
SELECT document_id, document_name, total_pages, content_hash,
|
||||
outline_json, generated_at
|
||||
FROM outline_cache WHERE document_id = ?
|
||||
""", (document_id,))
|
||||
|
||||
row = cursor.fetchone()
|
||||
|
||||
if not row:
|
||||
return None
|
||||
|
||||
outline_data = json.loads(row[4])
|
||||
return DocumentOutline.from_dict(outline_data)
|
||||
|
||||
def delete_outline(self, document_id: str) -> bool:
|
||||
"""删除纲要缓存"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("DELETE FROM outline_cache WHERE document_id = ?", (document_id,))
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def list_outlines(self, limit: int = 50) -> List[Dict]:
|
||||
"""获取纲要列表"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
SELECT document_id, document_name, total_pages, generated_at
|
||||
FROM outline_cache ORDER BY generated_at DESC LIMIT ?
|
||||
""", (limit,))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
|
||||
return [
|
||||
{
|
||||
"document_id": row[0],
|
||||
"document_name": row[1],
|
||||
"total_pages": row[2],
|
||||
"generated_at": row[3]
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
|
||||
# ==================== 文档向量 ====================
|
||||
|
||||
def save_document_vector(self, document_id: str, document_name: str,
|
||||
vector: List[float], tags: List[str] = None) -> int:
|
||||
"""保存文档向量"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
vector_hash = hashlib.md5(str(vector[:10]).encode()).hexdigest()[:8]
|
||||
|
||||
cursor.execute("""
|
||||
INSERT OR REPLACE INTO document_vectors
|
||||
(document_id, document_name, vector_hash, vector_json, tags_json, updated_at)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
document_id,
|
||||
document_name,
|
||||
vector_hash,
|
||||
json.dumps(vector),
|
||||
json.dumps(tags or [], ensure_ascii=False),
|
||||
datetime.now().isoformat()
|
||||
))
|
||||
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_document_vector(self, document_id: str) -> Optional[Dict]:
|
||||
"""获取文档向量"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
SELECT document_id, document_name, vector_json, tags_json
|
||||
FROM document_vectors WHERE document_id = ?
|
||||
""", (document_id,))
|
||||
|
||||
row = cursor.fetchone()
|
||||
|
||||
if not row:
|
||||
return None
|
||||
|
||||
return {
|
||||
"document_id": row[0],
|
||||
"document_name": row[1],
|
||||
"vector": json.loads(row[2]),
|
||||
"tags": json.loads(row[3]) if row[3] else []
|
||||
}
|
||||
|
||||
def get_all_document_vectors(self) -> List[Dict]:
|
||||
"""获取所有文档向量"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
SELECT document_id, document_name, vector_json, tags_json
|
||||
FROM document_vectors
|
||||
""")
|
||||
|
||||
rows = cursor.fetchall()
|
||||
|
||||
return [
|
||||
{
|
||||
"document_id": row[0],
|
||||
"document_name": row[1],
|
||||
"vector": json.loads(row[2]),
|
||||
"tags": json.loads(row[3]) if row[3] else []
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
|
||||
# ==================== 推荐缓存 ====================
|
||||
|
||||
def save_recommendations(self, document_id: str, recommendations: List[Recommendation]) -> int:
|
||||
"""保存推荐结果"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 先删除旧缓存
|
||||
cursor.execute("DELETE FROM recommendation_cache WHERE document_id = ?", (document_id,))
|
||||
|
||||
cursor.execute("""
|
||||
INSERT INTO recommendation_cache (document_id, recommendations_json, generated_at)
|
||||
VALUES (?, ?, ?)
|
||||
""", (
|
||||
document_id,
|
||||
json.dumps([r.to_dict() for r in recommendations], ensure_ascii=False),
|
||||
datetime.now().isoformat()
|
||||
))
|
||||
|
||||
return cursor.lastrowid
|
||||
|
||||
def get_recommendations(self, document_id: str) -> Optional[List[Recommendation]]:
|
||||
"""获取推荐结果"""
|
||||
with get_connection("knowledge") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("""
|
||||
SELECT recommendations_json FROM recommendation_cache WHERE document_id = ?
|
||||
""", (document_id,))
|
||||
|
||||
row = cursor.fetchone()
|
||||
|
||||
if not row:
|
||||
return None
|
||||
|
||||
return [Recommendation(**r) for r in json.loads(row[0])]
|
||||
|
||||
|
||||
# ==================== 纲要生成服务 ====================
|
||||
|
||||
class OutlineGenerator:
|
||||
"""纲要生成服务"""
|
||||
|
||||
def __init__(self, db: OutlineDB, documents_path: str = "./documents"):
|
||||
self.db = db
|
||||
self.documents_path = documents_path
|
||||
self.llm_client = None
|
||||
self.embedding_model = None
|
||||
self._init_llm()
|
||||
|
||||
def _init_llm(self):
|
||||
"""初始化LLM客户端"""
|
||||
try:
|
||||
from config import API_KEY, BASE_URL, MODEL
|
||||
from openai import OpenAI
|
||||
|
||||
self.llm_client = OpenAI(api_key=API_KEY, base_url=BASE_URL)
|
||||
self.model = MODEL
|
||||
logger.info("LLM客户端初始化成功")
|
||||
except ImportError:
|
||||
logger.warning("未找到LLM配置，纲要生成功能受限")
|
||||
self.llm_client = None
|
||||
|
||||
def generate_outline(self, document_id: str, force: bool = False) -> DocumentOutline:
|
||||
"""
|
||||
生成文档纲要
|
||||
|
||||
Args:
|
||||
document_id: 文档ID（相对路径，如 public/差旅管理办法.txt）
|
||||
force: 是否强制重新生成
|
||||
|
||||
Returns:
|
||||
DocumentOutline
|
||||
"""
|
||||
# 1. 检查缓存
|
||||
if not force:
|
||||
cached = self.db.get_outline(document_id)
|
||||
if cached:
|
||||
# 检查内容是否变化
|
||||
current_hash = self._get_document_hash(document_id)
|
||||
if current_hash == cached.content_hash:
|
||||
logger.info(f"使用缓存的纲要: {document_id}")
|
||||
return cached
|
||||
|
||||
# 2. 获取文档内容
|
||||
document_content = self._read_document(document_id)
|
||||
if not document_content:
|
||||
raise ValueError(f"文档不存在或无法读取: {document_id}")
|
||||
|
||||
document_name = os.path.basename(document_id)
|
||||
|
||||
# 3. 使用LLM提取结构
|
||||
outline_data = self._extract_structure(document_content, document_name)
|
||||
|
||||
# 4. 构建纲要对象
|
||||
outline = DocumentOutline(
|
||||
document_id=document_id,
|
||||
document_name=document_name,
|
||||
total_pages=0, # 可以后续计算
|
||||
content_hash=self._get_document_hash(document_id),
|
||||
generated_at=datetime.now().isoformat(),
|
||||
outline=[OutlineNode.from_dict(n) for n in outline_data.get("children", [])]
|
||||
)
|
||||
|
||||
# 5. 保存缓存
|
||||
self.db.save_outline(outline)
|
||||
logger.info(f"纲要生成完成: {document_id}")
|
||||
|
||||
return outline
|
||||
|
||||
def _read_document(self, document_id: str) -> Optional[str]:
|
||||
"""读取文档内容"""
|
||||
# 处理不同的文档格式
|
||||
file_path = os.path.join(self.documents_path, document_id)
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
logger.error(f"文档不存在: {file_path}")
|
||||
return None
|
||||
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
|
||||
try:
|
||||
if ext == '.txt':
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
return f.read()
|
||||
|
||||
elif ext == '.pdf':
|
||||
# 使用 pdfplumber 提取文本
|
||||
try:
|
||||
import pdfplumber
|
||||
text_parts = []
|
||||
with pdfplumber.open(file_path) as pdf:
|
||||
for page in pdf.pages:
|
||||
text = page.extract_text()
|
||||
if text:
|
||||
text_parts.append(text)
|
||||
return '\n'.join(text_parts)
|
||||
except ImportError:
|
||||
logger.warning("pdfplumber未安装，无法读取PDF")
|
||||
return None
|
||||
|
||||
elif ext in ['.docx', '.doc']:
|
||||
# 使用 python-docx 提取文本
|
||||
try:
|
||||
from docx import Document
|
||||
doc = Document(file_path)
|
||||
return '\n'.join([p.text for p in doc.paragraphs if p.text.strip()])
|
||||
except ImportError:
|
||||
logger.warning("python-docx未安装，无法读取Word文档")
|
||||
return None
|
||||
|
||||
elif ext == '.xlsx':
|
||||
# 使用 openpyxl 提取文本
|
||||
try:
|
||||
from openpyxl import load_workbook
|
||||
wb = load_workbook(file_path)
|
||||
text_parts = []
|
||||
for sheet in wb.worksheets:
|
||||
for row in sheet.iter_rows(values_only=True):
|
||||
text_parts.extend([str(cell) for cell in row if cell])
|
||||
return '\n'.join(text_parts)
|
||||
except ImportError:
|
||||
logger.warning("openpyxl未安装，无法读取Excel")
|
||||
return None
|
||||
|
||||
else:
|
||||
# 尝试作为文本读取
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
return f.read()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"读取文档失败: {e}")
|
||||
return None
|
||||
|
||||
def _get_document_hash(self, document_id: str) -> str:
|
||||
"""计算文档哈希"""
|
||||
file_path = os.path.join(self.documents_path, document_id)
|
||||
if not os.path.exists(file_path):
|
||||
return ""
|
||||
|
||||
with open(file_path, 'rb') as f:
|
||||
return hashlib.md5(f.read()).hexdigest()
|
||||
|
||||
def _extract_structure(self, content: str, document_name: str) -> Dict:
|
||||
"""使用LLM提取文档结构"""
|
||||
if not self.llm_client:
|
||||
# 返回基本结构
|
||||
return {
|
||||
"title": document_name,
|
||||
"summary": "无法生成摘要（LLM未配置）",
|
||||
"children": []
|
||||
}
|
||||
|
||||
# 限制内容长度
|
||||
max_length = 8000
|
||||
if len(content) > max_length:
|
||||
content = content[:max_length] + "\n...(内容已截断)"
|
||||
|
||||
prompt = f"""请分析以下制度文档，提取章节结构和核心要点。
|
||||
|
||||
文档名称：{document_name}
|
||||
|
||||
文档内容：
|
||||
{content}
|
||||
|
||||
请按以下格式返回JSON：
|
||||
{{
|
||||
"title": "文档标题",
|
||||
"summary": "文档概述（50字以内）",
|
||||
"children": [
|
||||
{{
|
||||
"id": "1",
|
||||
"title": "第一章 一级标题",
|
||||
"summary": "本章核心要点（30字以内）",
|
||||
"level": 1,
|
||||
"order": 1,
|
||||
"children": [
|
||||
{{
|
||||
"id": "1.1",
|
||||
"title": "1.1 二级标题",
|
||||
"summary": "核心要点",
|
||||
"level": 2,
|
||||
"order": 1,
|
||||
"children": []
|
||||
}}
|
||||
]
|
||||
}}
|
||||
]
|
||||
}}
|
||||
|
||||
要求：
|
||||
1. 识别文档的一级、二级、三级标题（如有）
|
||||
2. 每个章节提取核心要点，不超过30字
|
||||
3. 保持层级关系，最多3层
|
||||
4. 只返回JSON，不要有其他内容"""
|
||||
|
||||
try:
|
||||
result_text = call_llm(
|
||||
self.llm_client,
|
||||
prompt=prompt,
|
||||
model=self.model,
|
||||
temperature=0.3,
|
||||
max_tokens=3000
|
||||
)
|
||||
|
||||
if not result_text:
|
||||
logger.error("LLM提取结构返回空")
|
||||
return {
|
||||
"title": document_name,
|
||||
"summary": "生成失败: LLM返回空",
|
||||
"children": []
|
||||
}
|
||||
|
||||
# 清理可能的markdown标记
|
||||
if result_text.startswith("```json"):
|
||||
result_text = result_text[7:]
|
||||
if result_text.startswith("```"):
|
||||
result_text = result_text[3:]
|
||||
if result_text.endswith("```"):
|
||||
result_text = result_text[:-3]
|
||||
result_text = result_text.strip()
|
||||
|
||||
return json.loads(result_text)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"LLM提取结构失败: {e}")
|
||||
return {
|
||||
"title": document_name,
|
||||
"summary": f"生成失败: {str(e)}",
|
||||
"children": []
|
||||
}
|
||||
|
||||
def export_outline(self, outline: DocumentOutline, format: str = "json") -> str:
|
||||
"""
|
||||
导出纲要
|
||||
|
||||
Args:
|
||||
outline: 纲要对象
|
||||
format: 导出格式 (json/markdown/markmap)
|
||||
|
||||
Returns:
|
||||
导出内容
|
||||
"""
|
||||
if format == "json":
|
||||
return json.dumps(outline.to_dict(), ensure_ascii=False, indent=2)
|
||||
|
||||
elif format == "markdown":
|
||||
return self._export_markdown(outline)
|
||||
|
||||
elif format == "markmap":
|
||||
# markmap 格式（可渲染为思维导图的 Markdown）
|
||||
return self._export_markmap(outline)
|
||||
|
||||
else:
|
||||
raise ValueError(f"不支持的导出格式: {format}")
|
||||
|
||||
def _export_markdown(self, outline: DocumentOutline, node: OutlineNode = None, level: int = 0) -> str:
|
||||
"""导出为 Markdown 格式"""
|
||||
lines = []
|
||||
|
||||
if node is None:
|
||||
# 根级别
|
||||
lines.append(f"# {outline.document_name}\n")
|
||||
lines.append(f"> 生成时间: {outline.generated_at}\n")
|
||||
for child in outline.outline:
|
||||
lines.append(self._export_markdown(outline, child, 1))
|
||||
else:
|
||||
# 节点级别
|
||||
prefix = "#" * (level + 1)
|
||||
lines.append(f"{prefix} {node.title}\n")
|
||||
if node.summary:
|
||||
lines.append(f"{node.summary}\n")
|
||||
for child in node.children:
|
||||
lines.append(self._export_markdown(outline, child, level + 1))
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def _export_markmap(self, outline: DocumentOutline) -> str:
|
||||
"""导出为 markmap 格式（思维导图 Markdown）"""
|
||||
lines = [f"# {outline.document_name}"]
|
||||
|
||||
def render_node(node: OutlineNode, level: int):
|
||||
indent = " " * level
|
||||
lines.append(f"{indent}- {node.title}")
|
||||
if node.summary:
|
||||
lines.append(f"{indent} - *{node.summary}*")
|
||||
for child in node.children:
|
||||
render_node(child, level + 1)
|
||||
|
||||
for node in outline.outline:
|
||||
render_node(node, 1)
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def batch_generate(self, document_ids: List[str], force: bool = False) -> Dict[str, DocumentOutline]:
|
||||
"""
|
||||
批量生成纲要
|
||||
|
||||
Args:
|
||||
document_ids: 文档ID列表
|
||||
force: 是否强制重新生成
|
||||
|
||||
Returns:
|
||||
{document_id: DocumentOutline}
|
||||
"""
|
||||
results = {}
|
||||
for doc_id in document_ids:
|
||||
try:
|
||||
results[doc_id] = self.generate_outline(doc_id, force)
|
||||
except Exception as e:
|
||||
logger.error(f"生成纲要失败 {doc_id}: {e}")
|
||||
results[doc_id] = None
|
||||
return results
|
||||
|
||||
|
||||
# ==================== 关联推荐服务 ====================
|
||||
|
||||
class RecommendationService:
|
||||
"""关联推荐服务"""
|
||||
|
||||
def __init__(self, db: OutlineDB, documents_path: str = "./documents",
|
||||
chroma_collection=None, embedding_model=None):
|
||||
self.db = db
|
||||
self.documents_path = documents_path
|
||||
self.chroma_collection = chroma_collection
|
||||
self.embedding_model = embedding_model
|
||||
|
||||
def get_recommendations(self, document_id: str, top_k: int = 5,
|
||||
use_cache: bool = True) -> List[Recommendation]:
|
||||
"""
|
||||
获取关联推荐
|
||||
|
||||
Args:
|
||||
document_id: 当前文档ID
|
||||
top_k: 返回数量
|
||||
use_cache: 是否使用缓存
|
||||
|
||||
Returns:
|
||||
推荐列表
|
||||
"""
|
||||
# 1. 检查缓存
|
||||
if use_cache:
|
||||
cached = self.db.get_recommendations(document_id)
|
||||
if cached:
|
||||
logger.info(f"使用缓存的推荐: {document_id}")
|
||||
return cached[:top_k]
|
||||
|
||||
# 2. 获取当前文档向量
|
||||
current_vector = self._get_or_compute_vector(document_id)
|
||||
if current_vector is None:
|
||||
logger.warning(f"无法获取文档向量: {document_id}")
|
||||
return []
|
||||
|
||||
current_doc = self.db.get_document_vector(document_id)
|
||||
current_tags = current_doc.get("tags", []) if current_doc else []
|
||||
|
||||
# 3. 检索相似文档
|
||||
similar_docs = self._search_similar(current_vector, top_k * 3, exclude_id=document_id)
|
||||
|
||||
# 4. 计算综合得分
|
||||
recommendations = []
|
||||
for doc in similar_docs:
|
||||
# 标签匹配得分
|
||||
doc_tags = doc.get("tags", [])
|
||||
tag_overlap = len(set(current_tags) & set(doc_tags))
|
||||
tag_score = min(tag_overlap * 0.15, 0.3) # 最高0.3
|
||||
|
||||
# 综合得分
|
||||
similarity = doc.get("similarity", 0)
|
||||
final_score = similarity * 0.7 + tag_score
|
||||
|
||||
# 推荐理由
|
||||
reasons = []
|
||||
if similarity > 0.8:
|
||||
reasons.append("内容高度相似")
|
||||
elif similarity > 0.6:
|
||||
reasons.append("内容相关")
|
||||
if tag_overlap > 0:
|
||||
reasons.append(f"包含{tag_overlap}个相同标签")
|
||||
|
||||
recommendation = Recommendation(
|
||||
document_id=doc.get("document_id", ""),
|
||||
document_name=doc.get("document_name", ""),
|
||||
summary=doc.get("summary", "")[:100] if doc.get("summary") else "",
|
||||
similarity=round(similarity, 3),
|
||||
tag_score=round(tag_score, 3),
|
||||
final_score=round(final_score, 3),
|
||||
tags=doc_tags[:5],
|
||||
reason="、".join(reasons) if reasons else "相关推荐"
|
||||
)
|
||||
recommendations.append(recommendation)
|
||||
|
||||
# 5. 排序并截取
|
||||
recommendations.sort(key=lambda x: x.final_score, reverse=True)
|
||||
results = recommendations[:top_k]
|
||||
|
||||
# 6. 缓存结果
|
||||
if results:
|
||||
self.db.save_recommendations(document_id, results)
|
||||
|
||||
return results
|
||||
|
||||
def _get_or_compute_vector(self, document_id: str) -> Optional[List[float]]:
|
||||
"""获取或计算文档向量"""
|
||||
# 检查缓存
|
||||
cached = self.db.get_document_vector(document_id)
|
||||
if cached and cached.get("vector"):
|
||||
return cached["vector"]
|
||||
|
||||
# 计算向量
|
||||
if not self.embedding_model:
|
||||
logger.warning("嵌入模型未初始化，无法计算向量")
|
||||
return None
|
||||
|
||||
# 读取文档内容
|
||||
file_path = os.path.join(self.documents_path, document_id)
|
||||
if not os.path.exists(file_path):
|
||||
return None
|
||||
|
||||
try:
|
||||
# 使用 _read_document 方法读取文档（支持多种格式）
|
||||
content = self._read_document(file_path)
|
||||
if not content:
|
||||
logger.warning(f"无法读取文档内容: {document_id}")
|
||||
return None
|
||||
|
||||
# 计算向量
|
||||
vector = self.embedding_model.encode(content[:5000]) # 限制长度
|
||||
|
||||
# 缓存
|
||||
doc_name = os.path.basename(document_id)
|
||||
self.db.save_document_vector(document_id, doc_name, vector.tolist())
|
||||
|
||||
return vector.tolist()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"计算文档向量失败: {e}")
|
||||
return None
|
||||
|
||||
def _search_similar(self, query_vector: List[float], top_k: int,
|
||||
exclude_id: str = None) -> List[Dict]:
|
||||
"""搜索相似文档"""
|
||||
if not self.chroma_collection:
|
||||
# 使用数据库缓存
|
||||
all_docs = self.db.get_all_document_vectors()
|
||||
results = []
|
||||
|
||||
import numpy as np
|
||||
query_vec = np.array(query_vector)
|
||||
|
||||
for doc in all_docs:
|
||||
if exclude_id and doc["document_id"] == exclude_id:
|
||||
continue
|
||||
|
||||
doc_vec = np.array(doc["vector"])
|
||||
# 余弦相似度
|
||||
similarity = np.dot(query_vec, doc_vec) / (
|
||||
np.linalg.norm(query_vec) * np.linalg.norm(doc_vec)
|
||||
)
|
||||
|
||||
results.append({
|
||||
"document_id": doc["document_id"],
|
||||
"document_name": doc["document_name"],
|
||||
"similarity": float(similarity),
|
||||
"tags": doc.get("tags", [])
|
||||
})
|
||||
|
||||
results.sort(key=lambda x: x["similarity"], reverse=True)
|
||||
return results[:top_k]
|
||||
|
||||
# 使用 ChromaDB
|
||||
try:
|
||||
results = self.chroma_collection.query(
|
||||
query_embeddings=[query_vector],
|
||||
n_results=top_k + 1 # 多取一个，排除自己
|
||||
)
|
||||
|
||||
docs = []
|
||||
for i, doc_id in enumerate(results.get("ids", [[]])[0]):
|
||||
if exclude_id and doc_id == exclude_id:
|
||||
continue
|
||||
|
||||
metadata = results.get("metadatas", [[]])[0][i] if results.get("metadatas") else {}
|
||||
distance = results.get("distances", [[]])[0][i] if results.get("distances") else 0
|
||||
|
||||
# 转换距离为相似度
|
||||
similarity = 1 - distance if distance < 1 else 0
|
||||
|
||||
docs.append({
|
||||
"document_id": doc_id,
|
||||
"document_name": metadata.get("source", doc_id),
|
||||
"similarity": similarity,
|
||||
"tags": metadata.get("tags", "").split(",") if metadata.get("tags") else []
|
||||
})
|
||||
|
||||
return docs[:top_k]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"相似文档检索失败: {e}")
|
||||
return []
|
||||
|
||||
def compute_all_vectors(self) -> int:
|
||||
"""计算所有文档的向量"""
|
||||
if not self.embedding_model:
|
||||
logger.warning("嵌入模型未初始化")
|
||||
return 0
|
||||
|
||||
count = 0
|
||||
for root, dirs, files in os.walk(self.documents_path):
|
||||
for file in files:
|
||||
if file.endswith(('.txt', '.pdf', '.docx', '.md')):
|
||||
document_id = os.path.relpath(
|
||||
os.path.join(root, file),
|
||||
self.documents_path
|
||||
).replace("\\", "/")
|
||||
|
||||
try:
|
||||
self._get_or_compute_vector(document_id)
|
||||
count += 1
|
||||
except Exception as e:
|
||||
logger.error(f"计算向量失败 {document_id}: {e}")
|
||||
|
||||
logger.info(f"计算了 {count} 个文档的向量")
|
||||
return count
|
||||
|
||||
|
||||
# ==================== 便捷函数 ====================
|
||||
|
||||
def create_services(documents_path: str = "./documents",
|
||||
chroma_collection=None,
|
||||
embedding_model=None) -> Tuple[OutlineDB, OutlineGenerator, RecommendationService]:
|
||||
"""
|
||||
创建服务实例
|
||||
|
||||
Args:
|
||||
documents_path: 文档目录
|
||||
chroma_collection: ChromaDB集合
|
||||
embedding_model: 嵌入模型
|
||||
|
||||
Returns:
|
||||
(数据库实例, 纲要生成服务, 推荐服务)
|
||||
"""
|
||||
db = OutlineDB()
|
||||
outline_generator = OutlineGenerator(db, documents_path)
|
||||
recommendation_service = RecommendationService(
|
||||
db, documents_path, chroma_collection, embedding_model
|
||||
)
|
||||
|
||||
return db, outline_generator, recommendation_service
|
||||
|
||||
|
||||
# ==================== 使用示例 ====================
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
# 设置编码
|
||||
if sys.platform == 'win32':
|
||||
sys.stdout.reconfigure(encoding='utf-8')
|
||||
|
||||
print("=" * 60)
|
||||
print("纲要生成与关联推荐服务测试")
|
||||
print("=" * 60)
|
||||
|
||||
# 创建服务
|
||||
db, outline_svc, rec_svc = create_services(
|
||||
documents_path="./documents"
|
||||
)
|
||||
|
||||
# 测试纲要生成
|
||||
print("\n[1] 测试纲要生成...")
|
||||
|
||||
# 查找一个测试文档
|
||||
documents_path = "./documents"
|
||||
test_doc = None
|
||||
for root, dirs, files in os.walk(documents_path):
|
||||
for file in files:
|
||||
if file.endswith('.txt'):
|
||||
test_doc = os.path.relpath(
|
||||
os.path.join(root, file),
|
||||
documents_path
|
||||
).replace("\\", "/")
|
||||
break
|
||||
if test_doc:
|
||||
break
|
||||
|
||||
if test_doc:
|
||||
print(f" 测试文档: {test_doc}")
|
||||
try:
|
||||
outline = outline_svc.generate_outline(test_doc)
|
||||
print(f" 标题: {outline.document_name}")
|
||||
print(f" 章节数: {len(outline.outline)}")
|
||||
for node in outline.outline[:3]:
|
||||
print(f" - {node.title}: {node.summary[:30]}...")
|
||||
|
||||
# 测试导出
|
||||
print("\n[2] 测试导出...")
|
||||
md_content = outline_svc.export_outline(outline, "markdown")
|
||||
print(f" Markdown 导出长度: {len(md_content)} 字符")
|
||||
|
||||
markmap_content = outline_svc.export_outline(outline, "markmap")
|
||||
print(f" Markmap 导出长度: {len(markmap_content)} 字符")
|
||||
|
||||
except Exception as e:
|
||||
print(f" 纲要生成失败: {e}")
|
||||
else:
|
||||
print(" 未找到测试文档")
|
||||
|
||||
# 测试推荐服务
|
||||
print("\n[3] 测试关联推荐...")
|
||||
if test_doc:
|
||||
try:
|
||||
recommendations = rec_svc.get_recommendations(test_doc, top_k=3)
|
||||
print(f" 推荐数量: {len(recommendations)}")
|
||||
for rec in recommendations:
|
||||
print(f" - {rec.document_name} (相似度: {rec.similarity}, 理由: {rec.reason})")
|
||||
except Exception as e:
|
||||
print(f" 推荐失败: {e}")
|
||||
|
||||
# 列出纲要缓存
|
||||
print("\n[4] 测试纲要缓存...")
|
||||
outlines = db.list_outlines()
|
||||
print(f" 缓存数量: {len(outlines)}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("测试完成")
|
||||
print("=" * 60)
|
||||
437
services/session.py
Normal file
437
services/session.py
Normal file
@@ -0,0 +1,437 @@
|
||||
"""
|
||||
会话管理器 - 支持多用户对话历史
|
||||
|
||||
功能：
|
||||
1. 多用户会话隔离
|
||||
2. 对话历史持久化（SQLite）
|
||||
3. 历史压缩（避免上下文过长）
|
||||
4. 会话过期清理
|
||||
|
||||
使用方式：
|
||||
from services.session import SessionManager
|
||||
|
||||
sm = SessionManager()
|
||||
session_id = sm.create_session("user_123")
|
||||
|
||||
# 添加对话
|
||||
sm.add_message(session_id, "user", "出差补助标准是什么？")
|
||||
sm.add_message(session_id, "assistant", "根据规定...")
|
||||
|
||||
# 获取历史
|
||||
history = sm.get_history(session_id)
|
||||
|
||||
# 生成带历史的提示词
|
||||
context = sm.build_context(session_id, "那请假呢？")
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional, List, Dict
|
||||
|
||||
from data.db import get_connection
|
||||
|
||||
|
||||
class SessionManager:
|
||||
"""会话管理器"""
|
||||
|
||||
def __init__(self, session_expire_hours: int = 24):
|
||||
"""
|
||||
初始化会话管理器
|
||||
|
||||
Args:
|
||||
session_expire_hours: 会话过期时间（小时）
|
||||
"""
|
||||
self.session_expire_hours = session_expire_hours
|
||||
self._init_db()
|
||||
|
||||
def _init_db(self):
|
||||
"""初始化数据库表"""
|
||||
from data.db import init_databases
|
||||
init_databases()
|
||||
|
||||
def create_session(self, user_id: str, metadata: dict = None) -> str:
|
||||
"""
|
||||
创建新会话
|
||||
|
||||
Args:
|
||||
user_id: 用户ID
|
||||
metadata: 可选的元数据（如用户名、IP等）
|
||||
|
||||
Returns:
|
||||
session_id: 会话ID
|
||||
"""
|
||||
session_id = str(uuid.uuid4())
|
||||
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT INTO sessions (session_id, user_id, metadata)
|
||||
VALUES (?, ?, ?)
|
||||
''', (session_id, user_id, json.dumps(metadata or {})))
|
||||
|
||||
return session_id
|
||||
|
||||
def get_or_create_session(self, user_id: str, session_id: str = None) -> str:
|
||||
"""
|
||||
获取或创建会话
|
||||
|
||||
Args:
|
||||
user_id: 用户ID
|
||||
session_id: 可选的会话ID，如果提供则验证归属
|
||||
|
||||
Returns:
|
||||
session_id: 有效会话ID
|
||||
"""
|
||||
if session_id:
|
||||
# 验证会话是否存在且属于该用户
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT session_id FROM sessions
|
||||
WHERE session_id = ? AND user_id = ?
|
||||
''', (session_id, user_id))
|
||||
|
||||
result = cursor.fetchone()
|
||||
|
||||
if result:
|
||||
# 更新最后活跃时间
|
||||
self._update_last_active(session_id)
|
||||
return session_id
|
||||
|
||||
# 创建新会话
|
||||
return self.create_session(user_id)
|
||||
|
||||
def _update_last_active(self, session_id: str):
|
||||
"""更新会话最后活跃时间"""
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
UPDATE sessions
|
||||
SET last_active = CURRENT_TIMESTAMP
|
||||
WHERE session_id = ?
|
||||
''', (session_id,))
|
||||
|
||||
def add_message(self, session_id: str, role: str, content: str, metadata: dict = None):
|
||||
"""
|
||||
添加消息到会话历史
|
||||
|
||||
Args:
|
||||
session_id: 会话ID
|
||||
role: 角色 (user/assistant)
|
||||
content: 消息内容
|
||||
metadata: 可选的扩展数据（如sources, images, is_rag）
|
||||
"""
|
||||
import json
|
||||
meta_str = json.dumps(metadata) if metadata else None
|
||||
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO messages (session_id, role, content, metadata)
|
||||
VALUES (?, ?, ?, ?)
|
||||
''', (session_id, role, content, meta_str))
|
||||
|
||||
# 更新会话活跃时间
|
||||
cursor.execute('''
|
||||
UPDATE sessions
|
||||
SET last_active = CURRENT_TIMESTAMP
|
||||
WHERE session_id = ?
|
||||
''', (session_id,))
|
||||
|
||||
def get_history(self, session_id: str, limit: int = 20) -> List[Dict]:
|
||||
"""
|
||||
获取会话历史
|
||||
|
||||
Args:
|
||||
session_id: 会话ID
|
||||
limit: 最大消息数
|
||||
|
||||
Returns:
|
||||
[{"id": 1, "role": "user/assistant", "content": "...", "metadata": {...}}, ...]
|
||||
"""
|
||||
import json
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT id, role, content, metadata, created_at
|
||||
FROM messages
|
||||
WHERE session_id = ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?
|
||||
''', (session_id, limit))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
|
||||
# 按时间正序排列（旧的在前）
|
||||
history = []
|
||||
for row in reversed(rows):
|
||||
meta_str = row[3]
|
||||
meta = {}
|
||||
if meta_str:
|
||||
try:
|
||||
meta = json.loads(meta_str)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
pass
|
||||
|
||||
history.append({
|
||||
"id": row[0],
|
||||
"role": row[1],
|
||||
"content": row[2],
|
||||
"metadata": meta,
|
||||
"created_at": row[4]
|
||||
})
|
||||
|
||||
return history
|
||||
|
||||
def get_history_text(self, session_id: str, limit: int = 10) -> str:
|
||||
"""
|
||||
获取历史文本格式（用于Prompt）
|
||||
|
||||
Args:
|
||||
session_id: 会话ID
|
||||
limit: 最大轮次（一问一答为一轮）
|
||||
|
||||
Returns:
|
||||
格式化的历史文本
|
||||
"""
|
||||
history = self.get_history(session_id, limit=limit * 2)
|
||||
|
||||
if not history:
|
||||
return ""
|
||||
|
||||
lines = []
|
||||
for msg in history:
|
||||
role_name = "用户" if msg["role"] == "user" else "助手"
|
||||
lines.append(f"{role_name}：{msg['content']}")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
def build_context(self, session_id: str, current_query: str,
|
||||
max_history_tokens: int = 1500) -> str:
|
||||
"""
|
||||
构建带历史的上下文（智能压缩）
|
||||
|
||||
Args:
|
||||
session_id: 会话ID
|
||||
current_query: 当前问题
|
||||
max_history_tokens: 历史最大token数（估算）
|
||||
|
||||
Returns:
|
||||
包含历史的上下文文本
|
||||
"""
|
||||
history = self.get_history(session_id, limit=20)
|
||||
|
||||
if not history:
|
||||
return current_query
|
||||
|
||||
# 构建历史摘要
|
||||
history_text = self._compress_history(history, max_history_tokens)
|
||||
|
||||
context = f"""【对话历史】
|
||||
{history_text}
|
||||
|
||||
【当前问题】
|
||||
{current_query}"""
|
||||
|
||||
return context
|
||||
|
||||
def _compress_history(self, history: List[Dict], max_tokens: int) -> str:
|
||||
"""
|
||||
压缩历史（避免上下文过长）
|
||||
|
||||
策略：
|
||||
1. 保留最近的完整对话
|
||||
2. 较早的对话进行摘要压缩
|
||||
"""
|
||||
# 简单估算：1个中文字约等于1.5个token
|
||||
def estimate_tokens(text: str) -> int:
|
||||
return len(text) * 1.5
|
||||
|
||||
# 从最新开始，保留尽可能多的完整对话
|
||||
selected = []
|
||||
total_tokens = 0
|
||||
|
||||
for msg in reversed(history):
|
||||
msg_tokens = estimate_tokens(msg["content"])
|
||||
|
||||
if total_tokens + msg_tokens > max_tokens:
|
||||
# 超出限制，停止
|
||||
break
|
||||
|
||||
selected.insert(0, msg)
|
||||
total_tokens += msg_tokens
|
||||
|
||||
if not selected:
|
||||
return ""
|
||||
|
||||
# 如果有更早的历史被省略，添加提示
|
||||
if len(selected) < len(history):
|
||||
omitted_count = len(history) - len(selected)
|
||||
prefix = f"[省略了 {omitted_count} 条较早的历史消息]\n\n"
|
||||
else:
|
||||
prefix = ""
|
||||
|
||||
# 格式化
|
||||
lines = []
|
||||
for msg in selected:
|
||||
role_name = "用户" if msg["role"] == "user" else "助手"
|
||||
lines.append(f"{role_name}：{msg['content']}")
|
||||
|
||||
return prefix + "\n".join(lines)
|
||||
|
||||
def clear_history(self, session_id: str):
|
||||
"""清空会话历史"""
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
DELETE FROM messages WHERE session_id = ?
|
||||
''', (session_id,))
|
||||
|
||||
def delete_session(self, session_id: str):
|
||||
"""删除会话"""
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
DELETE FROM messages WHERE session_id = ?
|
||||
''', (session_id,))
|
||||
|
||||
cursor.execute('''
|
||||
DELETE FROM sessions WHERE session_id = ?
|
||||
''', (session_id,))
|
||||
|
||||
def cleanup_expired_sessions(self):
|
||||
"""清理过期会话"""
|
||||
expire_time = datetime.now() - timedelta(hours=self.session_expire_hours)
|
||||
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 删除过期会话的消息
|
||||
cursor.execute('''
|
||||
DELETE FROM messages
|
||||
WHERE session_id IN (
|
||||
SELECT session_id FROM sessions
|
||||
WHERE last_active < ?
|
||||
)
|
||||
''', (expire_time,))
|
||||
|
||||
# 删除过期会话
|
||||
cursor.execute('''
|
||||
DELETE FROM sessions
|
||||
WHERE last_active < ?
|
||||
''', (expire_time,))
|
||||
|
||||
deleted = cursor.rowcount
|
||||
|
||||
return deleted
|
||||
|
||||
def get_user_sessions(self, user_id: str, limit: int = 10) -> List[Dict]:
|
||||
"""
|
||||
获取用户的所有会话
|
||||
|
||||
Returns:
|
||||
[{"session_id": "...", "created_at": "...", "last_active": "..."}, ...]
|
||||
"""
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
SELECT session_id, created_at, last_active, metadata
|
||||
FROM sessions
|
||||
WHERE user_id = ?
|
||||
ORDER BY last_active DESC
|
||||
LIMIT ?
|
||||
''', (user_id, limit))
|
||||
|
||||
rows = cursor.fetchall()
|
||||
|
||||
sessions = []
|
||||
for row in rows:
|
||||
sessions.append({
|
||||
"session_id": row[0],
|
||||
"created_at": row[1],
|
||||
"last_active": row[2],
|
||||
"metadata": json.loads(row[3]) if row[3] else {}
|
||||
})
|
||||
|
||||
return sessions
|
||||
|
||||
def get_stats(self) -> Dict:
|
||||
"""获取统计信息"""
|
||||
with get_connection("session") as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM sessions')
|
||||
total_sessions = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(*) FROM messages')
|
||||
total_messages = cursor.fetchone()[0]
|
||||
|
||||
cursor.execute('SELECT COUNT(DISTINCT user_id) FROM sessions')
|
||||
total_users = cursor.fetchone()[0]
|
||||
|
||||
return {
|
||||
"total_sessions": total_sessions,
|
||||
"total_messages": total_messages,
|
||||
"total_users": total_users
|
||||
}
|
||||
|
||||
|
||||
# 测试代码
|
||||
if __name__ == "__main__":
|
||||
# 创建会话管理器
|
||||
sm = SessionManager()
|
||||
|
||||
print("=" * 50)
|
||||
print("会话管理器测试")
|
||||
print("=" * 50)
|
||||
|
||||
# 测试1: 创建会话
|
||||
print("\n【测试1】创建会话")
|
||||
user_id = "test_user_001"
|
||||
session_id = sm.create_session(user_id, {"name": "测试用户"})
|
||||
print(f"用户ID: {user_id}")
|
||||
print(f"会话ID: {session_id}")
|
||||
|
||||
# 测试2: 添加对话
|
||||
print("\n【测试2】添加对话")
|
||||
sm.add_message(session_id, "user", "出差补助标准是什么？")
|
||||
sm.add_message(session_id, "assistant", "根据公司规定，出差补助包括伙食费、交通费和住宿费。伙食补助每天100元...")
|
||||
sm.add_message(session_id, "user", "那请假流程呢？")
|
||||
sm.add_message(session_id, "assistant", "请假流程如下：1. 提交请假申请...")
|
||||
|
||||
print("已添加4条消息")
|
||||
|
||||
# 测试3: 获取历史
|
||||
print("\n【测试3】获取历史")
|
||||
history = sm.get_history(session_id)
|
||||
for msg in history:
|
||||
role = "用户" if msg["role"] == "user" else "助手"
|
||||
print(f" {role}: {msg['content'][:30]}...")
|
||||
|
||||
# 测试4: 构建上下文
|
||||
print("\n【测试4】构建上下文")
|
||||
context = sm.build_context(session_id, "婚假有多少天？")
|
||||
print(context)
|
||||
|
||||
# 测试5: 统计信息
|
||||
print("\n【测试5】统计信息")
|
||||
stats = sm.get_stats()
|
||||
print(f" 总会话数: {stats['total_sessions']}")
|
||||
print(f" 总消息数: {stats['total_messages']}")
|
||||
print(f" 总用户数: {stats['total_users']}")
|
||||
|
||||
# 测试6: 获取用户会话列表
|
||||
print("\n【测试6】用户会话列表")
|
||||
sessions = sm.get_user_sessions(user_id)
|
||||
for s in sessions:
|
||||
print(f" 会话ID: {s['session_id'][:8]}...")
|
||||
print(f" 创建时间: {s['created_at']}")
|
||||
print(f" 最后活跃: {s['last_active']}")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("测试完成")
|
||||
Reference in New Issue
Block a user