diff --git a/backend/python/app/api/knowledge.py b/backend/python/app/api/knowledge.py index bcec7ac..13c2ee2 100644 --- a/backend/python/app/api/knowledge.py +++ b/backend/python/app/api/knowledge.py @@ -1,1447 +1,1515 @@ -from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile, File, Form -from sqlalchemy.orm import Session -from sqlalchemy import and_, desc, asc -from pydantic import BaseModel, Field -from typing import Optional, List -import uuid -import os -import json -import re -import logging -import httpx -from datetime import datetime - -from ..database import get_db -from ..models import ( - User, KnowledgeBase, KnowledgeDocument, KnowledgeDirectory, Asset, - KnowledgeGraphNode, KnowledgeGraphEdge, KnowledgeGraphBuildLog, - AIModel, ModelProvider, - gen_uuid -) -from ..auth import get_current_user -from ..middleware.log_service import LogService - -router = APIRouter() -logger = logging.getLogger("aiflow") - -# 知识库文件存储根目录 -KNOWLEDGE_STORAGE_DIR = os.environ.get("KNOWLEDGE_STORAGE_DIR", "/data/share/project/data/knowledge") - - -class CreateKnowledgeBaseRequest(BaseModel): - name: str = Field(min_length=1, max_length=200) - description: Optional[str] = None - icon: Optional[str] = None - - -class UpdateKnowledgeBaseRequest(BaseModel): - name: Optional[str] = None - description: Optional[str] = None - icon: Optional[str] = None - status: Optional[str] = None - - -class CreateDocumentRequest(BaseModel): - filename: str - file_size: Optional[int] = None - file_type: Optional[str] = None - - -class CreateDirectoryRequest(BaseModel): - name: str = Field(min_length=1, max_length=200) - parent_id: Optional[str] = None - knowledge_base_id: Optional[str] = None - - -class KnowledgeQARequest(BaseModel): - base_id: str - question: str = Field(min_length=1, max_length=5000) - - -@router.get("/bases") -def list_knowledge_bases( - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - bases = ( - db.query(KnowledgeBase) - .filter(KnowledgeBase.user_id == current_user.id) - .order_by(desc(KnowledgeBase.updated_at)) - .all() - ) - - # 如果用户没有知识库，自动创建两个默认知识库（每个用户只创建一次） - if not bases: - # 检查是否已存在名为"默认知识库"或"知识图谱库"的知识库（兜底检查） - existing_default = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.user_id == current_user.id, KnowledgeBase.name == "默认知识库") - ).first() - existing_graph = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.user_id == current_user.id, KnowledgeBase.name == "知识图谱库") - ).first() - - if not existing_default: - default_kb = KnowledgeBase( - id=gen_uuid(), - user_id=current_user.id, - name="默认知识库", - description="联网采集、直接上传的内容默认存储", - directory_path="/", - document_count=0, - created_at=datetime.now(), - updated_at=datetime.now(), - ) - db.add(default_kb) - LogService.log_operation( - user_id=current_user.id, username=current_user.username, action="auto_create", - target="knowledge_base", target_id=default_kb.id, - detail="自动创建默认知识库", - ) - - if not existing_graph: - graph_kb = KnowledgeBase( - id=gen_uuid(), - user_id=current_user.id, - name="知识图谱库", - description="对话转知识图谱、构建知识图谱的默认存储", - directory_path="/", - document_count=0, - created_at=datetime.now(), - updated_at=datetime.now(), - ) - db.add(graph_kb) - LogService.log_operation( - user_id=current_user.id, username=current_user.username, action="auto_create", - target="knowledge_base", target_id=graph_kb.id, - detail="自动创建知识图谱库", - ) - - db.commit() - - # 重新查询 - bases = ( - db.query(KnowledgeBase) - .filter(KnowledgeBase.user_id == current_user.id) - .order_by(desc(KnowledgeBase.updated_at)) - .all() - ) - - # 批量查询所有知识库的真实图谱节点数（从 knowledge_graph_nodes 表统计） - kb_ids = [kb.id for kb in bases] - graph_node_counts = {} - if kb_ids: - from sqlalchemy import func as sa_func - node_stats = db.query( - KnowledgeGraphNode.knowledge_base_id, - sa_func.count(KnowledgeGraphNode.id) - ).filter( - KnowledgeGraphNode.knowledge_base_id.in_(kb_ids) - ).group_by(KnowledgeGraphNode.knowledge_base_id).all() - graph_node_counts = {row[0]: row[1] for row in node_stats} - - result = [] - for kb in bases: - # 使用真实图谱数据，如果图谱表中有数据则用它，否则用旧字段（可能残留旧值） - real_node_count = graph_node_counts.get(kb.id) - if real_node_count is not None and real_node_count > 0: - display_graph_count = real_node_count - else: - display_graph_count = 0 # 清零旧残留值 - - result.append({ - "id": kb.id, - "name": kb.name, - "description": kb.description, - "document_count": kb.document_count or 0, - "knowledge_graph_count": display_graph_count, - "created_at": str(kb.created_at) if kb.created_at else None, - "updated_at": str(kb.updated_at) if kb.updated_at else None, - }) - - # 计算存储空间信息 - # 存储已用：所有知识库文档的真实文件大小总和 - if bases: - from sqlalchemy import func as sa_func - total_size_result = db.query(sa_func.sum(KnowledgeDocument.file_size)).filter( - KnowledgeDocument.knowledge_base_id.in_([kb.id for kb in bases]) - ).scalar() - storage_used = total_size_result or 0 - else: - storage_used = 0 - - # 根据用户会员等级计算存储限额 - membership_level = current_user.membership_level or 0 - storage_limits = {0: 1 * 1024 * 1024 * 1024, 2: 10 * 1024 * 1024 * 1024, 3: 100 * 1024 * 1024 * 1024} - storage_limit = storage_limits.get(membership_level, 1 * 1024 * 1024 * 1024) - - return { - "code": 0, - "data": { - "bases": result, - "storage_used": storage_used, - "storage_limit": storage_limit, - "membership_level": membership_level, - }, - "message": "ok", - } - - -@router.post("/bases") -def create_knowledge_base( - req: CreateKnowledgeBaseRequest, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - kb = KnowledgeBase( - id=gen_uuid(), - user_id=current_user.id, - name=req.name, - description=req.description, - directory_path="/", - document_count=0, - created_at=datetime.now(), - updated_at=datetime.now(), - ) - db.add(kb) - db.commit() - db.refresh(kb) - - LogService.log_operation( - user_id=current_user.id, username=current_user.username, action="create", - target="knowledge_base", target_id=kb.id, - detail=f"创建知识库 {kb.name}", - ) - - return { - "code": 0, - "data": { - "id": kb.id, "name": kb.name, "description": kb.description, - "created_at": str(kb.created_at), - }, - "message": "创建成功", - } - - -@router.post("/bases/{base_id}/build-graph") -def build_knowledge_graph( - base_id: str, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - """构建知识图谱：调用LLM分析文档内容提取实体和关系，持久化到数据库""" - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - # 获取知识库中的已完成文档 - documents = db.query(KnowledgeDocument).filter( - and_( - KnowledgeDocument.knowledge_base_id == base_id, - KnowledgeDocument.status == "completed", - ) - ).all() - - if not documents: - raise HTTPException(status_code=400, detail="知识库中没有已完成的文档，请先上传文档并等待处理完成") - - # 获取LLM配置（获取默认启用的模型及其provider信息） - llm_config = _get_llm_config(db) - if not llm_config: - raise HTTPException(status_code=500, detail="没有可用的AI模型，请先在后台管理中配置模型") - - api_url, api_key, model_name, timeout = llm_config - - # 先清理旧图谱数据 - db.query(KnowledgeGraphEdge).filter( - KnowledgeGraphEdge.knowledge_base_id == base_id - ).delete() - db.query(KnowledgeGraphNode).filter( - KnowledgeGraphNode.knowledge_base_id == base_id - ).delete() - db.query(KnowledgeGraphBuildLog).filter( - KnowledgeGraphBuildLog.knowledge_base_id == base_id - ).delete() - db.flush() - - total_nodes = 0 - total_edges = 0 - - for doc in documents: - # 读取文档文件内容 - file_content = _read_document_content(doc) - - if not file_content or len(file_content.strip()) < 10: - logger.warning(f"文档 {doc.filename} 内容为空或过短，跳过图谱构建") - # 记录构建日志 - build_log = KnowledgeGraphBuildLog( - id=gen_uuid(), - knowledge_base_id=base_id, - document_id=doc.id, - node_count=0, - edge_count=0, - status="completed", - created_at=datetime.now(), - ) - db.add(build_log) - continue - - # 调用LLM提取实体和关系 - entities, relations = _extract_entities_with_llm( - api_url, api_key, model_name, timeout, doc.filename, file_content - ) - - # 将提取结果写入数据库 - doc_node_count, doc_edge_count = _persist_graph_data( - db, base_id, doc.id, entities, relations - ) - total_nodes += doc_node_count - total_edges += doc_edge_count - - # 记录构建日志 - build_log = KnowledgeGraphBuildLog( - id=gen_uuid(), - knowledge_base_id=base_id, - document_id=doc.id, - node_count=doc_node_count, - edge_count=doc_edge_count, - status="completed", - created_at=datetime.now(), - ) - db.add(build_log) - - db.commit() - - # 更新知识库的图谱计数 - kb.knowledge_graph_count = total_nodes - kb.updated_at = datetime.now() - db.commit() - - LogService.log_operation( - user_id=current_user.id, username=current_user.username, action="build_graph", - target="knowledge_base", target_id=base_id, - detail=f"为知识库 {kb.name} 构建知识图谱，共 {total_nodes} 节点、{total_edges} 边", - ) - - # 构建返回数据 - graph_data = _build_graph_response(db, base_id, total_nodes, total_edges, len(documents)) - - return { - "code": 0, - "data": graph_data, - "message": f"知识图谱构建完成，共 {total_nodes} 个节点、{total_edges} 条边", - } - - -def _get_llm_config(db: Session): - """获取LLM配置：优先默认模型，否则使用第一个启用的模型""" - logger.info("[图谱构建] 开始查找LLM配置...") - ai_model = db.query(AIModel).filter( - AIModel.is_default == True, - AIModel.is_enabled == True, - ).first() - - if not ai_model: - ai_model = db.query(AIModel).filter( - AIModel.is_enabled == True, - AIModel.endpoint.isnot(None), - ).first() - logger.info(f"[图谱构建] 无默认模型，尝试使用第一个启用模型: {ai_model.name if ai_model else 'None'}") - - if not ai_model: - logger.error("[图谱构建] 没有可用的AI模型！") - return None - - api_key = ai_model.api_key - api_url = ai_model.endpoint - - # 如果模型本身没有独立的api_key/endpoint，从provider获取 - if (not api_key or not api_url) and ai_model.provider_id: - provider = db.query(ModelProvider).filter(ModelProvider.id == ai_model.provider_id).first() - if provider: - if not api_key: - api_key = provider.api_key - if not api_url and provider.api_base_url: - api_url = f"{provider.api_base_url.rstrip('/')}/v1/chat/completions" - - if not api_url or not api_key: - logger.error(f"[图谱构建] LLM配置不完整: api_url={api_url}, api_key={'***' if api_key else 'None'}") - return None - - logger.info(f"[图谱构建] 找到LLM配置: model={ai_model.name}, url={api_url[:50]}..., timeout={ai_model.timeout or 60}") - return (api_url, api_key, ai_model.name, ai_model.timeout or 60) - - -def _read_document_content(doc: KnowledgeDocument) -> str: - """从磁盘读取文档的实际内容""" - storage_dir = os.environ.get("KNOWLEDGE_STORAGE_DIR", "/data/share/project/data/knowledge") - file_path = os.path.join(storage_dir, doc.knowledge_base_id, f"{doc.id}.{doc.file_type}") - - logger.info(f"[图谱构建] 读取文档: {doc.filename}, id={doc.id}, type={doc.file_type}, path={file_path}, status={doc.status}") - - if not os.path.exists(file_path): - # 尝试其他可能的路径 - alt_path = os.path.join(storage_dir, doc.knowledge_base_id, f"{doc.id}.{doc.file_type.lower()}") - if os.path.exists(alt_path): - file_path = alt_path - logger.info(f"[图谱构建] 使用备用路径: {file_path}") - else: - # 尝试直接读file_path字段 - if doc.file_path and os.path.exists(doc.file_path): - file_path = doc.file_path - logger.info(f"[图谱构建] 使用数据库file_path: {file_path}") - else: - logger.warning(f"[图谱构建] 文档文件不存在: {file_path} (备用: {alt_path}, 数据库file_path: {doc.file_path})") - return "" - - try: - with open(file_path, "r", encoding="utf-8") as f: - content = f.read() - return content - except Exception as e: - logger.error(f"读取文档文件失败: {file_path}, error={e}") - return "" - - -def _extract_entities_with_llm( - api_url: str, api_key: str, model_name: str, timeout: int, - doc_filename: str, file_content: str -): - """调用LLM分析文档内容，提取实体和关系""" - # 限制内容长度（防止token溢出） - max_content_len = 8000 - if len(file_content) > max_content_len: - truncated_content = file_content[:max_content_len] + "\n\n...(内容已截断)" - else: - truncated_content = file_content - - prompt = f"""分析以下文档内容，识别所有命名实体（人物、组织、地点、产品、技术术语等）和它们之间的关系。 - -文档标题：{doc_filename} - -文档内容： -{truncated_content} - -请以JSON格式返回，格式如下（只返回JSON，不要包含其他内容）： -{{"entities":[{{"name":"实体名","type":"entity/attribute/concept/event","description":"简介","attributes":{{"key":"value"}}}}],"relations":[{{"source":"源实体名","target":"目标实体名","relation":"关系类型","description":"关系描述"}}]}} - -要求： -1. 实体名称使用文档原文中的表述 -2. 实体类型只能从 entity（命名实体）、attribute（属性特征）、concept（核心概念）、event（事件）中选择 -3. 每个实体尽量给出简短描述 -4. 关系需要明确源实体和目标实体的名称 -5. 不要返回不存在的实体或关系""" - - try: - with httpx.Client(timeout=timeout) as client: - resp = client.post( - api_url, - headers={ - "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json", - }, - json={ - "model": model_name, - "messages": [ - {"role": "system", "content": "你是一个专业的知识图谱构建助手，擅长从文档中提取实体和关系。你只返回严格符合要求的JSON格式数据。"}, - {"role": "user", "content": prompt}, - ], - "max_tokens": 4096, - "temperature": 0.3, - }, - ) - - if resp.status_code != 200: - logger.error(f"LLM API调用失败 (HTTP {resp.status_code}): {resp.text[:500]}") - return [], [] - - data = resp.json() - content = data.get("choices", [{}])[0].get("message", {}).get("content", "") - - # 记录token使用量（可选） - usage = data.get("usage", {}) - logger.info( - f"LLM图谱构建 token使用: prompt={usage.get('prompt_tokens', 0)}, " - f"completion={usage.get('completion_tokens', 0)}, " - f"total={usage.get('total_tokens', 0)}" - ) - - if not content: - logger.error("LLM返回空内容") - return [], [] - - # 解析LLM返回的JSON - entities, relations = _parse_llm_graph_response(content) - return entities, relations - - except httpx.TimeoutException: - logger.error(f"LLM API调用超时: {api_url}") - return [], [] - except httpx.ConnectError: - logger.error(f"LLM API连接失败: {api_url}") - return [], [] - except Exception as e: - logger.error(f"LLM API调用异常: {e}") - return [], [] - - -def _parse_llm_graph_response(content: str): - """解析LLM返回的图谱JSON，做容错处理""" - entities = [] - relations = [] - - # 尝试多种方式提取JSON - # 方式1: 直接解析 - try: - result = json.loads(content) - entities = result.get("entities", []) - relations = result.get("relations", []) - return entities, relations - except json.JSONDecodeError: - pass - - # 方式2: 提取```json ... ```代码块 - json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', content) - if json_match: - try: - result = json.loads(json_match.group(1)) - entities = result.get("entities", []) - relations = result.get("relations", []) - return entities, relations - except json.JSONDecodeError: - pass - - # 方式3: 提取最外层的{...} - brace_match = re.search(r'\{[\s\S]*\}', content) - if brace_match: - try: - result = json.loads(brace_match.group(0)) - entities = result.get("entities", []) - relations = result.get("relations", []) - return entities, relations - except json.JSONDecodeError: - pass - - logger.warning(f"无法解析LLM返回的JSON: {content[:500]}") - return [], [] - - -def _persist_graph_data(db: Session, base_id: str, doc_id: str, entities: list, relations: list): - """将LLM提取的实体和关系持久化到数据库，返回(节点数, 边数)""" - # 先创建节点，记录name到node_id的映射 - name_to_node_id = {} - node_count = 0 - - for entity in entities: - if not isinstance(entity, dict): - continue - name = entity.get("name", "").strip() - if not name: - continue - entity_type = entity.get("type", "entity") - if entity_type not in ("entity", "attribute", "concept", "event"): - entity_type = "entity" - description = entity.get("description", "") - attributes = entity.get("attributes", {}) - if not isinstance(attributes, dict): - attributes = {} - - node_id = gen_uuid() - node = KnowledgeGraphNode( - id=node_id, - knowledge_base_id=base_id, - source_doc_id=doc_id, - label=name[:500], - type=entity_type, - properties=attributes, - description=description[:65535] if description else None, - created_at=datetime.now(), - ) - db.add(node) - name_to_node_id[name] = node_id - node_count += 1 - - # 必须flush才能获取节点ID用于外键关联 - db.flush() - - # 创建边 - edge_count = 0 - for rel in relations: - if not isinstance(rel, dict): - continue - source_name = rel.get("source", "").strip() - target_name = rel.get("target", "").strip() - if not source_name or not target_name: - continue - source_node_id = name_to_node_id.get(source_name) - target_node_id = name_to_node_id.get(target_name) - if not source_node_id or not target_node_id: - # 如果引用了LLM返回中不存在的实体名，跳过 - logger.debug(f"关系引用了未注册实体: {source_name} -> {target_name}") - continue - - edge = KnowledgeGraphEdge( - id=gen_uuid(), - knowledge_base_id=base_id, - source_node_id=source_node_id, - target_node_id=target_node_id, - relation=rel.get("relation", "关联")[:200], - weight=float(rel.get("weight", 1.0)) if rel.get("weight") else 1.0, - description=rel.get("description", "")[:65535] if rel.get("description") else None, - created_at=datetime.now(), - ) - db.add(edge) - edge_count += 1 - - return node_count, edge_count - - -def _build_graph_response(db: Session, base_id: str, total_nodes: int, total_edges: int, doc_count: int): - """从数据库读取图谱数据并构建响应""" - nodes = db.query(KnowledgeGraphNode).filter( - KnowledgeGraphNode.knowledge_base_id == base_id - ).order_by(KnowledgeGraphNode.created_at.asc()).all() - - edges = db.query(KnowledgeGraphEdge).filter( - KnowledgeGraphEdge.knowledge_base_id == base_id - ).order_by(KnowledgeGraphEdge.created_at.asc()).all() - - graph_nodes = [] - for node in nodes: - graph_nodes.append({ - "id": node.id, - "label": node.label, - "type": node.type, - "size": 10 + (len(node.description) % 10 if node.description else 5), - "group": node.type, - "description": node.description, - "properties": node.properties or {}, - "source_doc_id": node.source_doc_id, - }) - - graph_edges = [] - for edge in edges: - graph_edges.append({ - "id": edge.id, - "source": edge.source_node_id, - "target": edge.target_node_id, - "label": edge.relation, - "weight": edge.weight, - "description": edge.description, - }) - - # 统计各类型节点数量 - entity_types = {} - for node in nodes: - t = node.type - entity_types[t] = entity_types.get(t, 0) + 1 - - return { - "nodes": graph_nodes, - "edges": graph_edges, - "stats": { - "total_nodes": total_nodes, - "total_edges": total_edges, - "document_count": doc_count, - "entity_types": entity_types, - }, - } - - -@router.get("/bases/{base_id}/graph") -def get_knowledge_graph( - base_id: str, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - """获取知识图谱数据：从持久化表中读取已构建的图谱""" - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - # 从数据库表读取图谱节点 - nodes = db.query(KnowledgeGraphNode).filter( - KnowledgeGraphNode.knowledge_base_id == base_id - ).order_by(KnowledgeGraphNode.created_at.asc()).all() - - if not nodes: - # 清理旧残留数据：如果 knowledge_graph_count > 0 但图谱表中无节点，清零旧值 - if (kb.knowledge_graph_count or 0) > 0: - kb.knowledge_graph_count = 0 - db.commit() - - # 检查是否有文档但尚未构建图谱 - doc_count = db.query(KnowledgeDocument).filter( - and_( - KnowledgeDocument.knowledge_base_id == base_id, - KnowledgeDocument.status == "completed", - ) - ).count() - - return { - "code": 0, - "data": { - "nodes": [], - "edges": [], - "stats": {"total_nodes": 0, "total_edges": 0, "document_count": doc_count}, - }, - "message": "暂无可用于构建图谱的文档，请先构建知识图谱", - } - - # 从数据库表读取图谱边 - edges = db.query(KnowledgeGraphEdge).filter( - KnowledgeGraphEdge.knowledge_base_id == base_id - ).order_by(KnowledgeGraphEdge.created_at.asc()).all() - - graph_nodes = [] - for node in nodes: - graph_nodes.append({ - "id": node.id, - "label": node.label, - "type": node.type, - "size": 10 + (len(node.description) % 10 if node.description else 5), - "group": node.type, - "description": node.description, - "properties": node.properties or {}, - "source_doc_id": node.source_doc_id, - }) - - graph_edges = [] - for edge in edges: - graph_edges.append({ - "id": edge.id, - "source": edge.source_node_id, - "target": edge.target_node_id, - "label": edge.relation, - "weight": edge.weight, - "description": edge.description, - }) - - # 统计各类型节点数量 - entity_types = {} - for node in nodes: - t = node.type - entity_types[t] = entity_types.get(t, 0) + 1 - - graph_data = { - "nodes": graph_nodes, - "edges": graph_edges, - "stats": { - "total_nodes": len(graph_nodes), - "total_edges": len(graph_edges), - "document_count": len(set(n.source_doc_id for n in nodes if n.source_doc_id)), - "entity_types": entity_types, - }, - } - - return {"code": 0, "data": graph_data, "message": "ok"} - -@router.get("/bases/{base_id}") -def get_knowledge_base( - base_id: str, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - return { - "code": 0, - "data": { - "id": kb.id, "name": kb.name, "description": kb.description, - "document_count": kb.document_count or 0, - "created_at": str(kb.created_at), "updated_at": str(kb.updated_at), - }, - "message": "ok", - } - - -@router.put("/bases/{base_id}") -def update_knowledge_base( - base_id: str, - req: UpdateKnowledgeBaseRequest, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - if req.name is not None: - kb.name = req.name - if req.description is not None: - kb.description = req.description - if req.icon is not None: - pass # icon not in model - if req.status is not None: - pass # status not in model - kb.updated_at = datetime.now() - - db.commit() - db.refresh(kb) - - return {"code": 0, "data": { - "id": kb.id, "name": kb.name, "description": kb.description, - }, "message": "更新成功"} - - -@router.delete("/bases/{base_id}") -def delete_knowledge_base( - base_id: str, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - # 删除知识库下的所有文档 - db.query(KnowledgeDocument).filter(KnowledgeDocument.knowledge_base_id == base_id).delete() - db.delete(kb) - db.commit() - - return {"code": 0, "message": "知识库已删除"} - - -@router.get("/bases/{base_id}/documents") -def get_documents( - base_id: str, - status: Optional[str] = Query(None), - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - # 先验证知识库归属 - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - query = db.query(KnowledgeDocument).filter(KnowledgeDocument.knowledge_base_id == base_id) - if status: - query = query.filter(KnowledgeDocument.status == status) - - documents = query.order_by(desc(KnowledgeDocument.created_at)).all() - - result = [] - for doc in documents: - result.append({ - "id": doc.id, - "base_id": doc.knowledge_base_id, - "filename": doc.filename, - "file_size": doc.file_size, - "file_type": doc.file_type, - "status": doc.status, - "chunk_count": doc.chunk_count or 0, - "created_at": str(doc.created_at) if doc.created_at else None, - "updated_at": str(doc.updated_at) if doc.updated_at else None, - }) - - return {"code": 0, "data": result, "message": "ok"} - - -@router.post("/bases/{base_id}/documents") -async def upload_document( - base_id: str, - file: Optional[UploadFile] = File(None), - filename: str = Form(...), - file_type: Optional[str] = Form(None), - file_size: Optional[int] = Form(None), - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - # 确定文件扩展名和存储路径 - file_ext = (file_type or "txt").lower() - if file_ext.startswith("."): - file_ext = file_ext[1:] - - # 创建知识库存储目录 - storage_dir = os.path.join(KNOWLEDGE_STORAGE_DIR, base_id) - os.makedirs(storage_dir, exist_ok=True) - - # 生成唯一文件名并存储 - doc_id = gen_uuid() - stored_filename = f"{doc_id}.{file_ext}" - stored_path = os.path.join(storage_dir, stored_filename) - - # 文件内容解析 - actual_size = file_size or 0 - content = "" - uploaded_content = b"" - - # 如果有真实文件上传，先保存并读取内容 - if file and file.filename: - try: - uploaded_content = await file.read() - actual_size = len(uploaded_content) - # 先写入原始文件 - with open(stored_path, "wb") as f: - f.write(uploaded_content) - # 对于文本类格式，尝试直接解析 - text_exts = {"txt", "md", "csv", "json", "yaml", "yml", "xml", "py", "js", "ts", "html", "css", "log"} - if file_ext in text_exts: - try: - content = uploaded_content.decode("utf-8") - except UnicodeDecodeError: - try: - content = uploaded_content.decode("gbk") - except UnicodeDecodeError: - content = f"[二进制文件: {filename}, 大小: {actual_size} bytes]" - else: - # 非文本格式：记录文件但内容用摘要代替 - content = _generate_document_content(filename, file_ext, actual_size) - logger.info(f"非文本文件 {filename}({file_ext}) 已保存到 {stored_path}, 使用模拟内容用于检索") - except Exception as e: - logger.error(f"文件读取失败: {e}") - content = _generate_document_content(filename, file_ext, actual_size) - else: - # 无文件上传（兼容旧 JSON 模式）：生成模拟内容 - actual_size = file_size or len(filename.encode('utf-8')) * 100 - content = _generate_document_content(filename, file_ext, actual_size) - # 写入模拟内容 - try: - with open(stored_path, "w", encoding="utf-8") as f: - f.write(content) - except Exception as e: - logger.error(f"知识库文档写入失败: {stored_path}, error={e}") - raise HTTPException(status_code=500, detail=f"文件存储失败: {str(e)}") - - # 确保文件已写入（防止重复写入） - if uploaded_content and file_ext not in {"txt", "md", "csv", "json", "yaml", "yml", "xml", "py", "js", "ts", "html", "css", "log"}: - pass # 已在上面 wb 写入 - elif not os.path.exists(stored_path): - try: - with open(stored_path, "w", encoding="utf-8") as f: - f.write(content) - except Exception as e: - logger.error(f"知识库文档写入失败: {stored_path}, error={e}") - raise HTTPException(status_code=500, detail=f"文件存储失败: {str(e)}") - - logger.info(f"知识库文档已处理: {stored_path}, ext={file_ext}, size={actual_size}, content_len={len(content)}") - - # 文本分块 - chunks = _split_content_into_chunks(content, chunk_size=500) - - doc = KnowledgeDocument( - id=doc_id, - knowledge_base_id=base_id, - filename=filename, - file_size=actual_size, - file_type=file_ext, - status="completed", - chunk_count=len(chunks), - created_at=datetime.now(), - updated_at=datetime.now(), - ) - db.add(doc) - - # 更新知识库统计 - kb.document_count = (kb.document_count or 0) + 1 - kb.updated_at = datetime.now() - - # 同步创建Asset记录 - _sync_asset_for_knowledge_doc(db, current_user.id, doc_id, filename, file_ext, actual_size) - - db.commit() - db.refresh(doc) - - LogService.log_operation( - user_id=current_user.id, username=current_user.username, action="upload", - target="knowledge_document", target_id=doc.id, - detail=f"上传文档 {doc.filename} 到知识库 {kb.name}, 存储路径: {stored_path}, 分块数: {len(chunks)}", - ) - - return { - "code": 0, - "data": { - "id": doc.id, "base_id": doc.knowledge_base_id, "filename": doc.filename, - "file_size": doc.file_size, "file_type": doc.file_type, - "status": "completed", "chunk_count": doc.chunk_count, - "created_at": str(doc.created_at), - "stored_path": stored_path, - }, - "message": "上传成功", - } - - -def _generate_document_content(filename: str, file_ext: str, file_size: int) -> str: - """根据文件类型生成真实的模拟文档内容（文本解析后的结果）""" - name_base = filename.rsplit(".", 1)[0] - paragraphs = [] - - if file_ext in ("txt", "md"): - paragraphs = [ - f"# {name_base}", - f"", - f"## 概述", - f"本文档记录了关于 {name_base} 的详细内容。该文档由用户上传至知识库，" - f"经过文本解析引擎处理后，提取了以下关键信息。", - f"", - f"## 核心内容", - f"1. 主题分析：{name_base} 涉及的核心领域包括数据管理、知识工程和智能分析。", - f"2. 关键指标：通过对文档内容的深度解析，识别出 {max(3, file_size // 10000)} 个关键概念节点。", - f"3. 关联关系：文档中各主题之间存在 {max(2, file_size // 20000)} 条语义关联。", - f"", - f"## 详细分析", - f"在 {name_base} 的深入研究中，我们发现以下几点值得关注：", - f"- 数据结构层面的优化空间约为 15-20%", - f"- 知识图谱构建的准确率达到 87.3%", - f"- 检索效率相比传统方法提升了 2.3 倍", - f"", - f"## 结论与建议", - f"基于以上分析，建议在以下几个方面进行改进：", - f"1. 持续优化知识提取的准确度", - f"2. 扩展实体识别覆盖范围", - f"3. 增强跨文档关联分析能力", - f"", - f"---", - f"*本文档由 AIFlow 知识库系统自动解析生成，上载时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*", - ] - elif file_ext == "pdf": - paragraphs = [ - f"PDF文档: {name_base}", - f"", - f"【文档信息】", - f"标题：{name_base}", - f"页数估计：{max(1, file_size // 50000)} 页", - f"创建者：AIFlow 用户", - f"", - f"【提取的文本内容】", - f"第一章 引言", - f"本文档旨在阐述 {name_base} 相关的核心概念和实践方法。", - f"随着人工智能技术的不断发展，知识管理系统在企业数字化转型中扮演着越来越重要的角色。", - f"", - f"第二章 方法论", - f"本研究采用混合方法，结合定量分析与定性评估，对 {name_base} 进行全面探讨。", - f"数据来源包括：内部数据库查询结果、用户行为日志、以及与领域专家的深度访谈。", - f"", - f"第三章 结果分析", - f"通过对收集数据的系统分析，得出以下关键发现：", - f"1. 核心指标X的均值为 78.5，标准差为 12.3", - f"2. 关键因素Y与Z之间存在显著相关性（r=0.72, p<0.001）", - f"3. 85% 的受访者认为 {name_base} 是未来发展的关键方向", - f"", - f"第四章 结论", - f"综合以上分析，{name_base} 具有重要的研究价值和实践意义。", - ] - elif file_ext in ("docx", "doc"): - paragraphs = [ - f"Word文档: {name_base}", - f"", - f"{name_base} - 正式文档", - f"", - f"【摘要】", - f"本报告围绕 {name_base} 展开全面分析，涵盖市场趋势、技术演进和应用场景三个维度。", - f"通过系统性的研究和数据采集，为相关决策提供参考依据。", - f"", - f"【正文】", - f"一、背景介绍", - f"在当前的数字化转型浪潮中，{name_base} 已经成为组织提升竞争力的关键要素。", - f"根据Gartner的报告，到2026年将有超过60%的企业采用AI驱动的知识管理方案。", - f"", - f"二、技术架构", - f"系统采用微服务架构设计，包括数据采集层、处理引擎层、存储层和应用层四个核心模块。", - f"其中处理引擎层集成了NLP自然语言处理、实体识别和关系抽取等多种AI能力。", - f"", - f"三、实施计划", - f"第一阶段（1-2月）：完成基础架构搭建和数据迁移", - f"第二阶段（3-4月）：部署AI模型并开始增量学习", - f"第三阶段（5-6月）：全面上线并进行效果评估", - ] - elif file_ext in ("xlsx", "xls"): - paragraphs = [ - f"Excel表格: {name_base}", - f"", - f"【表格数据概览】", - f"| 指标类别 | Q1 | Q2 | Q3 | Q4 | 年度总计 |", - f"|---------|-----|-----|-----|-----|---------|", - f"| 用户增长 | 12% | 18% | 22% | 28% | 80% |", - f"| 营收增幅 | 8% | 14% | 19% | 25% | 66% |", - f"| 客户满意度 | 85 | 88 | 91 | 93 | 89.3 |", - f"| 市场份额 | 15% | 17% | 20% | 23% | -- |", - f"", - f"【关键指标分析】", - f"1. 用户增长率呈逐季上升趋势，Q4达到峰值28%", - f"2. 营收增长与用户增长高度正相关", - f"3. 客户满意度持续改善，Q4达到93分", - f"4. 市场份额稳步扩大，年度总增长约8个百分点", - f"", - f"【数据来源】内部经营分析系统，数据截止日期：{datetime.now().strftime('%Y-%m-%d')}", - ] - elif file_ext in ("pptx", "ppt"): - paragraphs = [ - f"PPT演示文稿: {name_base}", - f"", - f"【幻灯片结构】", - f"第1页：封面 - {name_base}", - f"第2页：目录与议程", - f"第3页：背景与现状分析", - f"第4页：核心挑战与机遇", - f"第5页：解决方案概述", - f"第6-8页：详细方案展开", - f"第9页：实施路线图", - f"第10页：总结与展望", - f"", - f"【关键要点提取】", - f"- 战略目标：实现 {name_base} 的全面数字化转型", - f"- 核心路径：AI赋能 + 数据驱动 + 敏捷迭代", - f"- 预期成果：效率提升30%、成本降低20%、用户满意度提升15%", - f"- 投入预算：约500万元（分三期投入）", - ] - else: - paragraphs = [ - f"文档: {name_base}", - f"", - f"文件类型：{file_ext.upper()}", - f"文件大小：{file_size} bytes", - f"上传时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", - f"", - f"【内容摘要】", - f"本文档类型为 {file_ext.upper()} 格式，包含与 {name_base} 相关的内容。", - f"系统已完成基本的文本解析和索引建立，支持关键词检索和语义匹配。", - ] - - return "\n".join(paragraphs) - - -def _split_content_into_chunks(content: str, chunk_size: int = 500) -> list: - """将文档内容分割为文本块（模拟embedding的chunking过程）""" - chunks = [] - paragraphs = content.split("\n\n") - current_chunk = "" - - for para in paragraphs: - if len(current_chunk) + len(para) + 2 <= chunk_size: - current_chunk = (current_chunk + "\n\n" + para).strip() - else: - if current_chunk: - chunks.append(current_chunk) - if len(para) > chunk_size: - # 长段落进一步切分 - for i in range(0, len(para), chunk_size): - chunks.append(para[i:i+chunk_size]) - current_chunk = "" - else: - current_chunk = para - - if current_chunk: - chunks.append(current_chunk) - - return chunks if chunks else [content[:chunk_size]] - - -def _sync_asset_for_knowledge_doc(db: Session, user_id: str, doc_id: str, filename: str, file_type: str, file_size: int): - """将知识库文档同步为资产记录（问题9：资产关联知识库文档用量）""" - try: - # 文件类型归一化为前端子分类标签（docx/doc->word 等），确保出现在「文件 > Word」子栏 - _ft_label = { - "pdf": "pdf", "docx": "word", "doc": "word", "xlsx": "excel", "xls": "excel", - "pptx": "ppt", "ppt": "ppt", "csv": "txt", "md": "md", "txt": "txt", - }.get((file_type or "").lower(), "other") - asset = Asset( - id=gen_uuid(), - user_id=user_id, - name=filename, - category="file", - description=f"知识库文档 (ID: {doc_id})", - file_type=_ft_label, - size=file_size, - status="active", - knowledge_doc_id=doc_id, - created_at=datetime.now(), - updated_at=datetime.now(), - ) - db.add(asset) - logger.info(f"已同步知识库文档到资产: {filename} -> asset_id={asset.id}") - except Exception as e: - logger.warning(f"同步知识库文档到资产失败: {e}") - - -@router.delete("/bases/{base_id}/documents/{doc_id}") -def delete_document( - base_id: str, - doc_id: str, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - doc = db.query(KnowledgeDocument).filter( - and_(KnowledgeDocument.id == doc_id, KnowledgeDocument.knowledge_base_id == base_id) - ).first() - if not doc: - raise HTTPException(status_code=404, detail="文档不存在") - - # 更新知识库统计 - kb.document_count = max(0, (kb.document_count or 0) - 1) - pass # total_size removed - kb.updated_at = datetime.now() - - # 同步删除知识库文档对应的资产记录（双向同步） - try: - linked_asset = db.query(Asset).filter(Asset.knowledge_doc_id == doc_id).first() - if linked_asset: - linked_asset.status = "deleted" - linked_asset.updated_at = datetime.now() - except Exception as _e: - logger.warning(f"同步删除关联资产失败: {_e}") - - db.delete(doc) - db.commit() - - return {"code": 0, "message": "文档已删除"} - - -@router.get("/directories") -def list_directories( - base_id: Optional[str] = Query(None), - parent_id: Optional[str] = Query(None), - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - query = db.query(KnowledgeDirectory).filter( - KnowledgeDirectory.user_id == current_user.id - ) - - if base_id: - query = query.filter(KnowledgeDirectory.knowledge_base_id == base_id) - if parent_id is not None: - query = query.filter(KnowledgeDirectory.parent_id == parent_id) - else: - # 默认查根目录 - query = query.filter(KnowledgeDirectory.parent_id.is_(None)) - - directories = query.order_by(asc(KnowledgeDirectory.name)).all() - - result = [] - for d in directories: - result.append({ - "id": d.id, - "name": d.name, - "parent_id": d.parent_id, - "knowledge_base_id": d.knowledge_base_id, - "created_at": str(d.created_at) if d.created_at else None, - }) - - return {"code": 0, "data": result, "message": "ok"} - - -@router.post("/directories") -def create_directory( - req: CreateDirectoryRequest, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - directory = KnowledgeDirectory( - id=gen_uuid(), - user_id=current_user.id, - name=req.name, - parent_id=req.parent_id, - knowledge_base_id=req.knowledge_base_id, - created_at=datetime.now(), - ) - db.add(directory) - db.commit() - db.refresh(directory) - - return { - "code": 0, - "data": { - "id": directory.id, "name": directory.name, - "parent_id": directory.parent_id, "knowledge_base_id": directory.knowledge_base_id, - "created_at": str(directory.created_at), - }, - "message": "创建成功", - } - - -# ============================================================ -# 知识库问答溯源 API -# ============================================================ - -@router.post("/qa") -def knowledge_qa( - req: KnowledgeQARequest, - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - """知识库问答：在用户的知识库中检索相关文档并生成回答，返回原文引用和来源路径""" - # 验证知识库归属 - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == req.base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - # 检索知识库中的文档 - documents = db.query(KnowledgeDocument).filter( - and_( - KnowledgeDocument.knowledge_base_id == req.base_id, - KnowledgeDocument.status == "completed", - ) - ).all() - - if not documents: - return { - "code": 0, - "data": { - "question": req.question, - "answer": "该知识库中暂无已完成的文档。请先上传并处理文档后再进行问答。", - "sources": [], - "knowledge_base_name": kb.name, - }, - "message": "ok", - } - - # 构建引用来源列表（基于文档关键词匹配模拟检索） - import re - sources = [] - relevant_docs = [] - question_lower = req.question.lower() - - for doc in documents: - filename_lower = doc.filename.lower() - # 简单的关键词匹配检索（实际生产环境应使用向量检索） - keywords = re.split(r'[\s\-\_\.\,\;\:\!\?\(\)\[\]\{\}]+', question_lower) - filename_words = re.split(r'[\s\-\_\.\,\;\:\!\?\(\)\[\]\{\}]+', filename_lower) - match_count = sum(1 for kw in keywords if kw and any(kw in fw for fw in filename_words)) - - if match_count > 0 or len(documents) <= 3: - relevant_docs.append(doc) - # 为每个匹配文档生成模拟原文引用 - sample_content = generate_document_excerpt(doc, req.question) - sources.append({ - "document_id": doc.id, - "document_name": doc.filename, - "file_type": doc.file_type or "", - "file_size": doc.file_size or 0, - "source_path": f"/knowledge/{kb.id}/documents/{doc.id}", - "excerpt": sample_content, - "relevance": min(0.99, max(0.3, match_count * 0.2 + 0.3)) if match_count > 0 else 0.3, - }) - - # 生成模拟回答 - kb_name = kb.name - answer_parts = [f"根据知识库「{kb_name}」中相关文档的内容，针对您的问题回复如下：\n"] - - if sources: - answer_parts.append(f"共检索到 {len(sources)} 篇相关文档，以下为引用分析：\n") - for i, src in enumerate(sources[:5], 1): - answer_parts.append(f"\n【引用 {i}】来源：《{src['document_name']}》") - answer_parts.append(f"原文摘录：\"{src['excerpt']}\"") - answer_parts.append(f"文档类型：{src['file_type'].upper()} | 文档大小：{format_file_size(src['file_size'])}") - answer_parts.append(f"\n---\n综合以上文档内容，我们建议您参考上述引用来源获取详细信息。") - else: - answer_parts.append("未找到与问题高度相关的文档内容，建议上传更多相关文档后重试。") - - final_answer = "\n".join(answer_parts) - - return { - "code": 0, - "data": { - "question": req.question, - "answer": final_answer, - "sources": sources, - "knowledge_base_name": kb_name, - "total_documents": len(documents), - }, - "message": "ok", - } - - -@router.get("/bases/{base_id}/search") -def search_knowledge( - base_id: str, - keyword: str = Query(min_length=1), - db: Session = Depends(get_db), - current_user: User = Depends(get_current_user), -): - """知识库文档检索：按关键词搜索文档并返回引用片段""" - kb = db.query(KnowledgeBase).filter( - and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) - ).first() - if not kb: - raise HTTPException(status_code=404, detail="知识库不存在") - - documents = db.query(KnowledgeDocument).filter( - and_( - KnowledgeDocument.knowledge_base_id == base_id, - KnowledgeDocument.status == "completed", - KnowledgeDocument.filename.ilike(f"%{keyword}%"), - ) - ).all() - - result = [] - for doc in documents: - result.append({ - "document_id": doc.id, - "document_name": doc.filename, - "file_type": doc.file_type or "", - "file_size": doc.file_size or 0, - "source_path": f"/knowledge/{kb.id}/documents/{doc.id}", - "excerpt": generate_document_excerpt(doc, keyword), - "created_at": str(doc.created_at) if doc.created_at else None, - }) - - return { - "code": 0, - "data": { - "keyword": keyword, - "knowledge_base_name": kb.name, - "results": result, - "total": len(result), - }, - "message": "ok", - } - - -def generate_document_excerpt(doc, query: str) -> str: - """根据文档类型生成模拟的原文摘录""" - excerpts = { - "pdf": f"本文档详细阐述了与「{query}」相关的核心内容。根据第3章第2节的分析，关键指标呈现明显的增长趋势，其中核心参数值在统计周期内提升了约27.3%。进一步的研究表明，该趋势与市场整体发展方向高度一致...", - "docx": f"关于「{query}」的讨论，文档中明确指出：\"从数据层面来看，该领域在过去四个季度中保持了平均15%的复合增长率。这一结论基于对超过2000个样本的综合分析，置信度达到95%以上。建议决策层关注以下三个关键方向...\"", - "xlsx": f"数据表格「{query}」相关统计：\n| 指标 | Q1 | Q2 | Q3 | Q4 |\n|------|----|----|----|----|\n| 用户增长 | 12% | 18% | 22% | 28% |\n| 营收增幅 | 8% | 14% | 19% | 25% |\n数据来源：内部经营分析系统", - "ppt": f"演示文稿第15页「{query}」关键要点：\n• 战略目标：实现年度增长30%\n• 核心路径：优化产品矩阵，拓展海外市场\n• 资源配置：预算投入约500万元\n• 时间节点：Q2完成试点，Q4全面推广\n详细方案见第15-22页", - "md": f"# {query}\n\n## 核心观点\n\n文档明确指出：*\"在当前的行业格局下，企业需要通过技术创新和模式变革来建立竞争壁垒。具体来说，可以从以下三个维度切入：数据驱动决策、自动化流程优化、以及个性化用户体验。\"*\n\n## 关键数据\n\n- 市场规模：预计2026年达到200亿\n- 年复合增长率：35%\n- 头部企业市占率：约40%\n\n> 详细分析参见完整文档。", - "txt": f"文档内容摘录：\n\n\"...针对「{query}」这一课题，我们进行了深入的调研和分析。研究发现，超过78%的受访者认为该领域的创新速度正在加快。同时，技术成熟度曲线显示，相关技术有望在未来12-18个月内进入主流应用阶段。建议前期投入研发资源约50人月，预期投资回报周期为18个月...\"", - } - file_type = (doc.file_type or "txt").lower() - return excerpts.get(file_type, excerpts["txt"]) - - -def format_file_size(size_bytes: int) -> str: - """格式化文件大小""" - if size_bytes >= 1024 * 1024: - return f"{size_bytes / (1024 * 1024):.1f} MB" - elif size_bytes >= 1024: - return f"{size_bytes / 1024:.1f} KB" - return f"{size_bytes} B" - - -# ============================================================ -# 知识图谱构建与查询 API -# ============================================================ - -class BuildGraphRequest(BaseModel): - base_id: str - +import asyncio +from fastapi import APIRouter, Depends, HTTPException, Query, UploadFile, File, Form +from sqlalchemy.orm import Session +from sqlalchemy import and_, desc, asc +from pydantic import BaseModel, Field +from typing import Optional, List +import uuid +import os +import json +import re +import logging +import httpx +from datetime import datetime + +from ..database import get_db +from ..models import ( + User, KnowledgeBase, KnowledgeDocument, KnowledgeDirectory, Asset, + KnowledgeGraphNode, KnowledgeGraphEdge, KnowledgeGraphBuildLog, + AIModel, ModelProvider, + gen_uuid +) +from ..auth import get_current_user +from ..middleware.log_service import LogService + +router = APIRouter() +logger = logging.getLogger("aiflow") + +# 知识库文件存储根目录 +KNOWLEDGE_STORAGE_DIR = os.environ.get("KNOWLEDGE_STORAGE_DIR", "/data/share/project/data/knowledge") + + +class CreateKnowledgeBaseRequest(BaseModel): + name: str = Field(min_length=1, max_length=200) + description: Optional[str] = None + icon: Optional[str] = None + + +class UpdateKnowledgeBaseRequest(BaseModel): + name: Optional[str] = None + description: Optional[str] = None + icon: Optional[str] = None + status: Optional[str] = None + + +class CreateDocumentRequest(BaseModel): + filename: str + file_size: Optional[int] = None + file_type: Optional[str] = None + + +class CreateDirectoryRequest(BaseModel): + name: str = Field(min_length=1, max_length=200) + parent_id: Optional[str] = None + knowledge_base_id: Optional[str] = None + + +class KnowledgeQARequest(BaseModel): + base_id: str + question: str = Field(min_length=1, max_length=5000) + + +@router.get("/bases") +def list_knowledge_bases( + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + bases = ( + db.query(KnowledgeBase) + .filter(KnowledgeBase.user_id == current_user.id) + .order_by(desc(KnowledgeBase.updated_at)) + .all() + ) + + # 如果用户没有知识库，自动创建两个默认知识库（每个用户只创建一次） + if not bases: + # 检查是否已存在名为"默认知识库"或"知识图谱库"的知识库（兜底检查） + existing_default = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.user_id == current_user.id, KnowledgeBase.name == "默认知识库") + ).first() + existing_graph = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.user_id == current_user.id, KnowledgeBase.name == "知识图谱库") + ).first() + + if not existing_default: + default_kb = KnowledgeBase( + id=gen_uuid(), + user_id=current_user.id, + name="默认知识库", + description="联网采集、直接上传的内容默认存储", + directory_path="/", + document_count=0, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + db.add(default_kb) + LogService.log_operation( + user_id=current_user.id, username=current_user.username, action="auto_create", + target="knowledge_base", target_id=default_kb.id, + detail="自动创建默认知识库", + ) + + if not existing_graph: + graph_kb = KnowledgeBase( + id=gen_uuid(), + user_id=current_user.id, + name="知识图谱库", + description="对话转知识图谱、构建知识图谱的默认存储", + directory_path="/", + document_count=0, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + db.add(graph_kb) + LogService.log_operation( + user_id=current_user.id, username=current_user.username, action="auto_create", + target="knowledge_base", target_id=graph_kb.id, + detail="自动创建知识图谱库", + ) + + db.commit() + + # 重新查询 + bases = ( + db.query(KnowledgeBase) + .filter(KnowledgeBase.user_id == current_user.id) + .order_by(desc(KnowledgeBase.updated_at)) + .all() + ) + + # 批量查询所有知识库的真实图谱节点数（从 knowledge_graph_nodes 表统计） + kb_ids = [kb.id for kb in bases] + graph_node_counts = {} + if kb_ids: + from sqlalchemy import func as sa_func + node_stats = db.query( + KnowledgeGraphNode.knowledge_base_id, + sa_func.count(KnowledgeGraphNode.id) + ).filter( + KnowledgeGraphNode.knowledge_base_id.in_(kb_ids) + ).group_by(KnowledgeGraphNode.knowledge_base_id).all() + graph_node_counts = {row[0]: row[1] for row in node_stats} + + result = [] + for kb in bases: + # 使用真实图谱数据，如果图谱表中有数据则用它，否则用旧字段（可能残留旧值） + real_node_count = graph_node_counts.get(kb.id) + if real_node_count is not None and real_node_count > 0: + display_graph_count = real_node_count + else: + display_graph_count = 0 # 清零旧残留值 + + result.append({ + "id": kb.id, + "name": kb.name, + "description": kb.description, + "document_count": kb.document_count or 0, + "knowledge_graph_count": display_graph_count, + "created_at": str(kb.created_at) if kb.created_at else None, + "updated_at": str(kb.updated_at) if kb.updated_at else None, + }) + + # 计算存储空间信息 + # 存储已用：所有知识库文档的真实文件大小总和 + if bases: + from sqlalchemy import func as sa_func + total_size_result = db.query(sa_func.sum(KnowledgeDocument.file_size)).filter( + KnowledgeDocument.knowledge_base_id.in_([kb.id for kb in bases]) + ).scalar() + storage_used = total_size_result or 0 + else: + storage_used = 0 + + # 根据用户会员等级计算存储限额 + membership_level = current_user.membership_level or 0 + storage_limits = {0: 1 * 1024 * 1024 * 1024, 2: 10 * 1024 * 1024 * 1024, 3: 100 * 1024 * 1024 * 1024} + storage_limit = storage_limits.get(membership_level, 1 * 1024 * 1024 * 1024) + + return { + "code": 0, + "data": { + "bases": result, + "storage_used": storage_used, + "storage_limit": storage_limit, + "membership_level": membership_level, + }, + "message": "ok", + } + + +@router.post("/bases") +def create_knowledge_base( + req: CreateKnowledgeBaseRequest, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + kb = KnowledgeBase( + id=gen_uuid(), + user_id=current_user.id, + name=req.name, + description=req.description, + directory_path="/", + document_count=0, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + db.add(kb) + db.commit() + db.refresh(kb) + + LogService.log_operation( + user_id=current_user.id, username=current_user.username, action="create", + target="knowledge_base", target_id=kb.id, + detail=f"创建知识库 {kb.name}", + ) + + return { + "code": 0, + "data": { + "id": kb.id, "name": kb.name, "description": kb.description, + "created_at": str(kb.created_at), + }, + "message": "创建成功", + } + + +@router.post("/bases/{base_id}/build-graph") +def build_knowledge_graph( + base_id: str, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + """构建知识图谱：调用LLM分析文档内容提取实体和关系，持久化到数据库""" + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + # 获取知识库中的已完成文档 + documents = db.query(KnowledgeDocument).filter( + and_( + KnowledgeDocument.knowledge_base_id == base_id, + KnowledgeDocument.status == "completed", + ) + ).all() + + if not documents: + raise HTTPException(status_code=400, detail="知识库中没有已完成的文档，请先上传文档并等待处理完成") + + # 获取LLM配置（获取默认启用的模型及其provider信息） + llm_config = _get_llm_config(db) + if not llm_config: + raise HTTPException(status_code=500, detail="没有可用的AI模型，请先在后台管理中配置模型") + + api_url, api_key, model_name, timeout = llm_config + + # 先清理旧图谱数据 + db.query(KnowledgeGraphEdge).filter( + KnowledgeGraphEdge.knowledge_base_id == base_id + ).delete() + db.query(KnowledgeGraphNode).filter( + KnowledgeGraphNode.knowledge_base_id == base_id + ).delete() + db.query(KnowledgeGraphBuildLog).filter( + KnowledgeGraphBuildLog.knowledge_base_id == base_id + ).delete() + db.flush() + + total_nodes = 0 + total_edges = 0 + + for doc in documents: + # 读取文档文件内容 + file_content = _read_document_content(doc) + + if not file_content or len(file_content.strip()) < 10: + logger.warning(f"文档 {doc.filename} 内容为空或过短，跳过图谱构建") + # 记录构建日志 + build_log = KnowledgeGraphBuildLog( + id=gen_uuid(), + knowledge_base_id=base_id, + document_id=doc.id, + node_count=0, + edge_count=0, + status="completed", + created_at=datetime.now(), + ) + db.add(build_log) + continue + + # 调用LLM提取实体和关系 + entities, relations = _extract_entities_with_llm( + api_url, api_key, model_name, timeout, doc.filename, file_content + ) + + # 将提取结果写入数据库 + doc_node_count, doc_edge_count = _persist_graph_data( + db, base_id, doc.id, entities, relations + ) + total_nodes += doc_node_count + total_edges += doc_edge_count + + # 记录构建日志 + build_log = KnowledgeGraphBuildLog( + id=gen_uuid(), + knowledge_base_id=base_id, + document_id=doc.id, + node_count=doc_node_count, + edge_count=doc_edge_count, + status="completed", + created_at=datetime.now(), + ) + db.add(build_log) + + db.commit() + + # 更新知识库的图谱计数 + kb.knowledge_graph_count = total_nodes + kb.updated_at = datetime.now() + db.commit() + + LogService.log_operation( + user_id=current_user.id, username=current_user.username, action="build_graph", + target="knowledge_base", target_id=base_id, + detail=f"为知识库 {kb.name} 构建知识图谱，共 {total_nodes} 节点、{total_edges} 边", + ) + + # 构建返回数据 + graph_data = _build_graph_response(db, base_id, total_nodes, total_edges, len(documents)) + + return { + "code": 0, + "data": graph_data, + "message": f"知识图谱构建完成，共 {total_nodes} 个节点、{total_edges} 条边", + } + + +def _get_llm_config(db: Session): + """获取LLM配置：优先默认模型，否则使用第一个启用的模型""" + logger.info("[图谱构建] 开始查找LLM配置...") + ai_model = db.query(AIModel).filter( + AIModel.is_default == True, + AIModel.is_enabled == True, + ).first() + + if not ai_model: + ai_model = db.query(AIModel).filter( + AIModel.is_enabled == True, + AIModel.endpoint.isnot(None), + ).first() + logger.info(f"[图谱构建] 无默认模型，尝试使用第一个启用模型: {ai_model.name if ai_model else 'None'}") + + if not ai_model: + logger.error("[图谱构建] 没有可用的AI模型！") + return None + + api_key = ai_model.api_key + api_url = ai_model.endpoint + + # 如果模型本身没有独立的api_key/endpoint，从provider获取 + if (not api_key or not api_url) and ai_model.provider_id: + provider = db.query(ModelProvider).filter(ModelProvider.id == ai_model.provider_id).first() + if provider: + if not api_key: + api_key = provider.api_key + if not api_url and provider.api_base_url: + api_url = f"{provider.api_base_url.rstrip('/')}/v1/chat/completions" + + if not api_url or not api_key: + logger.error(f"[图谱构建] LLM配置不完整: api_url={api_url}, api_key={'***' if api_key else 'None'}") + return None + + logger.info(f"[图谱构建] 找到LLM配置: model={ai_model.name}, url={api_url[:50]}..., timeout={ai_model.timeout or 60}") + return (api_url, api_key, ai_model.name, ai_model.timeout or 60) + + +def _read_document_content(doc: KnowledgeDocument) -> str: + """从磁盘读取文档的实际内容""" + storage_dir = os.environ.get("KNOWLEDGE_STORAGE_DIR", "/data/share/project/data/knowledge") + file_path = os.path.join(storage_dir, doc.knowledge_base_id, f"{doc.id}.{doc.file_type}") + + logger.info(f"[图谱构建] 读取文档: {doc.filename}, id={doc.id}, type={doc.file_type}, path={file_path}, status={doc.status}") + + if not os.path.exists(file_path): + # 尝试其他可能的路径 + alt_path = os.path.join(storage_dir, doc.knowledge_base_id, f"{doc.id}.{doc.file_type.lower()}") + if os.path.exists(alt_path): + file_path = alt_path + logger.info(f"[图谱构建] 使用备用路径: {file_path}") + else: + # 尝试直接读file_path字段 + if doc.file_path and os.path.exists(doc.file_path): + file_path = doc.file_path + logger.info(f"[图谱构建] 使用数据库file_path: {file_path}") + else: + logger.warning(f"[图谱构建] 文档文件不存在: {file_path} (备用: {alt_path}, 数据库file_path: {doc.file_path})") + return "" + + try: + with open(file_path, "r", encoding="utf-8") as f: + content = f.read() + return content + except Exception as e: + logger.error(f"读取文档文件失败: {file_path}, error={e}") + return "" + + +def _extract_entities_with_llm( + api_url: str, api_key: str, model_name: str, timeout: int, + doc_filename: str, file_content: str +): + """调用LLM分析文档内容，提取实体和关系""" + # 限制内容长度（防止token溢出） + max_content_len = 8000 + if len(file_content) > max_content_len: + truncated_content = file_content[:max_content_len] + "\n\n...(内容已截断)" + else: + truncated_content = file_content + + prompt = f"""分析以下文档内容，识别所有命名实体（人物、组织、地点、产品、技术术语等）和它们之间的关系。 + +文档标题：{doc_filename} + +文档内容： +{truncated_content} + +请以JSON格式返回，格式如下（只返回JSON，不要包含其他内容）： +{{"entities":[{{"name":"实体名","type":"entity/attribute/concept/event","description":"简介","attributes":{{"key":"value"}}}}],"relations":[{{"source":"源实体名","target":"目标实体名","relation":"关系类型","description":"关系描述"}}]}} + +要求： +1. 实体名称使用文档原文中的表述 +2. 实体类型只能从 entity（命名实体）、attribute（属性特征）、concept（核心概念）、event（事件）中选择 +3. 每个实体尽量给出简短描述 +4. 关系需要明确源实体和目标实体的名称 +5. 不要返回不存在的实体或关系""" + + try: + with httpx.Client(timeout=timeout) as client: + resp = client.post( + api_url, + headers={ + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + }, + json={ + "model": model_name, + "messages": [ + {"role": "system", "content": "你是一个专业的知识图谱构建助手，擅长从文档中提取实体和关系。你只返回严格符合要求的JSON格式数据。"}, + {"role": "user", "content": prompt}, + ], + "max_tokens": 4096, + "temperature": 0.3, + }, + ) + + if resp.status_code != 200: + logger.error(f"LLM API调用失败 (HTTP {resp.status_code}): {resp.text[:500]}") + return [], [] + + data = resp.json() + content = data.get("choices", [{}])[0].get("message", {}).get("content", "") + + # 记录token使用量（可选） + usage = data.get("usage", {}) + logger.info( + f"LLM图谱构建 token使用: prompt={usage.get('prompt_tokens', 0)}, " + f"completion={usage.get('completion_tokens', 0)}, " + f"total={usage.get('total_tokens', 0)}" + ) + + if not content: + logger.error("LLM返回空内容") + return [], [] + + # 解析LLM返回的JSON + entities, relations = _parse_llm_graph_response(content) + return entities, relations + + except httpx.TimeoutException: + logger.error(f"LLM API调用超时: {api_url}") + return [], [] + except httpx.ConnectError: + logger.error(f"LLM API连接失败: {api_url}") + return [], [] + except Exception as e: + logger.error(f"LLM API调用异常: {e}") + return [], [] + + +def _parse_llm_graph_response(content: str): + """解析LLM返回的图谱JSON，做容错处理""" + entities = [] + relations = [] + + # 尝试多种方式提取JSON + # 方式1: 直接解析 + try: + result = json.loads(content) + entities = result.get("entities", []) + relations = result.get("relations", []) + return entities, relations + except json.JSONDecodeError: + pass + + # 方式2: 提取```json ... ```代码块 + json_match = re.search(r'```(?:json)?\s*([\s\S]*?)```', content) + if json_match: + try: + result = json.loads(json_match.group(1)) + entities = result.get("entities", []) + relations = result.get("relations", []) + return entities, relations + except json.JSONDecodeError: + pass + + # 方式3: 提取最外层的{...} + brace_match = re.search(r'\{[\s\S]*\}', content) + if brace_match: + try: + result = json.loads(brace_match.group(0)) + entities = result.get("entities", []) + relations = result.get("relations", []) + return entities, relations + except json.JSONDecodeError: + pass + + logger.warning(f"无法解析LLM返回的JSON: {content[:500]}") + return [], [] + + +def _persist_graph_data(db: Session, base_id: str, doc_id: str, entities: list, relations: list): + """将LLM提取的实体和关系持久化到数据库，返回(节点数, 边数)""" + # 先创建节点，记录name到node_id的映射 + name_to_node_id = {} + node_count = 0 + + for entity in entities: + if not isinstance(entity, dict): + continue + name = entity.get("name", "").strip() + if not name: + continue + entity_type = entity.get("type", "entity") + if entity_type not in ("entity", "attribute", "concept", "event"): + entity_type = "entity" + description = entity.get("description", "") + attributes = entity.get("attributes", {}) + if not isinstance(attributes, dict): + attributes = {} + + node_id = gen_uuid() + node = KnowledgeGraphNode( + id=node_id, + knowledge_base_id=base_id, + source_doc_id=doc_id, + label=name[:500], + type=entity_type, + properties=attributes, + description=description[:65535] if description else None, + created_at=datetime.now(), + ) + db.add(node) + name_to_node_id[name] = node_id + node_count += 1 + + # 必须flush才能获取节点ID用于外键关联 + db.flush() + + # 创建边 + edge_count = 0 + for rel in relations: + if not isinstance(rel, dict): + continue + source_name = rel.get("source", "").strip() + target_name = rel.get("target", "").strip() + if not source_name or not target_name: + continue + source_node_id = name_to_node_id.get(source_name) + target_node_id = name_to_node_id.get(target_name) + if not source_node_id or not target_node_id: + # 如果引用了LLM返回中不存在的实体名，跳过 + logger.debug(f"关系引用了未注册实体: {source_name} -> {target_name}") + continue + + edge = KnowledgeGraphEdge( + id=gen_uuid(), + knowledge_base_id=base_id, + source_node_id=source_node_id, + target_node_id=target_node_id, + relation=rel.get("relation", "关联")[:200], + weight=float(rel.get("weight", 1.0)) if rel.get("weight") else 1.0, + description=rel.get("description", "")[:65535] if rel.get("description") else None, + created_at=datetime.now(), + ) + db.add(edge) + edge_count += 1 + + return node_count, edge_count + + +def _build_graph_response(db: Session, base_id: str, total_nodes: int, total_edges: int, doc_count: int): + """从数据库读取图谱数据并构建响应""" + nodes = db.query(KnowledgeGraphNode).filter( + KnowledgeGraphNode.knowledge_base_id == base_id + ).order_by(KnowledgeGraphNode.created_at.asc()).all() + + edges = db.query(KnowledgeGraphEdge).filter( + KnowledgeGraphEdge.knowledge_base_id == base_id + ).order_by(KnowledgeGraphEdge.created_at.asc()).all() + + graph_nodes = [] + for node in nodes: + graph_nodes.append({ + "id": node.id, + "label": node.label, + "type": node.type, + "size": 10 + (len(node.description) % 10 if node.description else 5), + "group": node.type, + "description": node.description, + "properties": node.properties or {}, + "source_doc_id": node.source_doc_id, + }) + + graph_edges = [] + for edge in edges: + graph_edges.append({ + "id": edge.id, + "source": edge.source_node_id, + "target": edge.target_node_id, + "label": edge.relation, + "weight": edge.weight, + "description": edge.description, + }) + + # 统计各类型节点数量 + entity_types = {} + for node in nodes: + t = node.type + entity_types[t] = entity_types.get(t, 0) + 1 + + return { + "nodes": graph_nodes, + "edges": graph_edges, + "stats": { + "total_nodes": total_nodes, + "total_edges": total_edges, + "document_count": doc_count, + "entity_types": entity_types, + }, + } + + +@router.get("/bases/{base_id}/graph") +def get_knowledge_graph( + base_id: str, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + """获取知识图谱数据：从持久化表中读取已构建的图谱""" + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + # 从数据库表读取图谱节点 + nodes = db.query(KnowledgeGraphNode).filter( + KnowledgeGraphNode.knowledge_base_id == base_id + ).order_by(KnowledgeGraphNode.created_at.asc()).all() + + if not nodes: + # 清理旧残留数据：如果 knowledge_graph_count > 0 但图谱表中无节点，清零旧值 + if (kb.knowledge_graph_count or 0) > 0: + kb.knowledge_graph_count = 0 + db.commit() + + # 检查是否有文档但尚未构建图谱 + doc_count = db.query(KnowledgeDocument).filter( + and_( + KnowledgeDocument.knowledge_base_id == base_id, + KnowledgeDocument.status == "completed", + ) + ).count() + + return { + "code": 0, + "data": { + "nodes": [], + "edges": [], + "stats": {"total_nodes": 0, "total_edges": 0, "document_count": doc_count}, + }, + "message": "暂无可用于构建图谱的文档，请先构建知识图谱", + } + + # 从数据库表读取图谱边 + edges = db.query(KnowledgeGraphEdge).filter( + KnowledgeGraphEdge.knowledge_base_id == base_id + ).order_by(KnowledgeGraphEdge.created_at.asc()).all() + + graph_nodes = [] + for node in nodes: + graph_nodes.append({ + "id": node.id, + "label": node.label, + "type": node.type, + "size": 10 + (len(node.description) % 10 if node.description else 5), + "group": node.type, + "description": node.description, + "properties": node.properties or {}, + "source_doc_id": node.source_doc_id, + }) + + graph_edges = [] + for edge in edges: + graph_edges.append({ + "id": edge.id, + "source": edge.source_node_id, + "target": edge.target_node_id, + "label": edge.relation, + "weight": edge.weight, + "description": edge.description, + }) + + # 统计各类型节点数量 + entity_types = {} + for node in nodes: + t = node.type + entity_types[t] = entity_types.get(t, 0) + 1 + + graph_data = { + "nodes": graph_nodes, + "edges": graph_edges, + "stats": { + "total_nodes": len(graph_nodes), + "total_edges": len(graph_edges), + "document_count": len(set(n.source_doc_id for n in nodes if n.source_doc_id)), + "entity_types": entity_types, + }, + } + + return {"code": 0, "data": graph_data, "message": "ok"} + +@router.get("/bases/{base_id}") +def get_knowledge_base( + base_id: str, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + return { + "code": 0, + "data": { + "id": kb.id, "name": kb.name, "description": kb.description, + "document_count": kb.document_count or 0, + "created_at": str(kb.created_at), "updated_at": str(kb.updated_at), + }, + "message": "ok", + } + + +@router.put("/bases/{base_id}") +def update_knowledge_base( + base_id: str, + req: UpdateKnowledgeBaseRequest, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + if req.name is not None: + kb.name = req.name + if req.description is not None: + kb.description = req.description + if req.icon is not None: + pass # icon not in model + if req.status is not None: + pass # status not in model + kb.updated_at = datetime.now() + + db.commit() + db.refresh(kb) + + return {"code": 0, "data": { + "id": kb.id, "name": kb.name, "description": kb.description, + }, "message": "更新成功"} + + +@router.delete("/bases/{base_id}") +def delete_knowledge_base( + base_id: str, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + # 删除知识库下的所有文档 + db.query(KnowledgeDocument).filter(KnowledgeDocument.knowledge_base_id == base_id).delete() + db.delete(kb) + db.commit() + + return {"code": 0, "message": "知识库已删除"} + + +@router.get("/bases/{base_id}/documents") +def get_documents( + base_id: str, + status: Optional[str] = Query(None), + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + # 先验证知识库归属 + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + query = db.query(KnowledgeDocument).filter(KnowledgeDocument.knowledge_base_id == base_id) + if status: + query = query.filter(KnowledgeDocument.status == status) + + documents = query.order_by(desc(KnowledgeDocument.created_at)).all() + + result = [] + for doc in documents: + result.append({ + "id": doc.id, + "base_id": doc.knowledge_base_id, + "filename": doc.filename, + "file_size": doc.file_size, + "file_type": doc.file_type, + "status": doc.status, + "chunk_count": doc.chunk_count or 0, + "created_at": str(doc.created_at) if doc.created_at else None, + "updated_at": str(doc.updated_at) if doc.updated_at else None, + }) + + return {"code": 0, "data": result, "message": "ok"} + + +@router.post("/bases/{base_id}/documents") +async def _index_kb_doc_to_pgvector(doc_id: str, kb_id: str, content: str): + """将文档内容异步索引到 pgvector，供语义向量检索使用（best-effort）。""" + try: + from app.pgvector_search import index_document + await index_document(doc_id, kb_id, content, metadata={"source": "kb_upload"}) + except Exception as e: + logger.warning(f"[kb] 向量索引失败 doc={doc_id}: {e}") + + +async def upload_document( + base_id: str, + file: Optional[UploadFile] = File(None), + filename: str = Form(...), + file_type: Optional[str] = Form(None), + file_size: Optional[int] = Form(None), + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + # 确定文件扩展名和存储路径 + file_ext = (file_type or "txt").lower() + if file_ext.startswith("."): + file_ext = file_ext[1:] + + # 创建知识库存储目录 + storage_dir = os.path.join(KNOWLEDGE_STORAGE_DIR, base_id) + os.makedirs(storage_dir, exist_ok=True) + + # 生成唯一文件名并存储 + doc_id = gen_uuid() + stored_filename = f"{doc_id}.{file_ext}" + stored_path = os.path.join(storage_dir, stored_filename) + + # 文件内容解析 + actual_size = file_size or 0 + content = "" + uploaded_content = b"" + + # 如果有真实文件上传，先保存并读取内容 + if file and file.filename: + try: + uploaded_content = await file.read() + actual_size = len(uploaded_content) + # 先写入原始文件 + with open(stored_path, "wb") as f: + f.write(uploaded_content) + # 对于文本类格式，尝试直接解析 + text_exts = {"txt", "md", "csv", "json", "yaml", "yml", "xml", "py", "js", "ts", "html", "css", "log"} + if file_ext in text_exts: + try: + content = uploaded_content.decode("utf-8") + except UnicodeDecodeError: + try: + content = uploaded_content.decode("gbk") + except UnicodeDecodeError: + content = f"[二进制文件: {filename}, 大小: {actual_size} bytes]" + else: + # 非文本格式：记录文件但内容用摘要代替 + content = _generate_document_content(filename, file_ext, actual_size) + logger.info(f"非文本文件 {filename}({file_ext}) 已保存到 {stored_path}, 使用模拟内容用于检索") + except Exception as e: + logger.error(f"文件读取失败: {e}") + content = _generate_document_content(filename, file_ext, actual_size) + else: + # 无文件上传（兼容旧 JSON 模式）：生成模拟内容 + actual_size = file_size or len(filename.encode('utf-8')) * 100 + content = _generate_document_content(filename, file_ext, actual_size) + # 写入模拟内容 + try: + with open(stored_path, "w", encoding="utf-8") as f: + f.write(content) + except Exception as e: + logger.error(f"知识库文档写入失败: {stored_path}, error={e}") + raise HTTPException(status_code=500, detail=f"文件存储失败: {str(e)}") + + # 确保文件已写入（防止重复写入） + if uploaded_content and file_ext not in {"txt", "md", "csv", "json", "yaml", "yml", "xml", "py", "js", "ts", "html", "css", "log"}: + pass # 已在上面 wb 写入 + elif not os.path.exists(stored_path): + try: + with open(stored_path, "w", encoding="utf-8") as f: + f.write(content) + except Exception as e: + logger.error(f"知识库文档写入失败: {stored_path}, error={e}") + raise HTTPException(status_code=500, detail=f"文件存储失败: {str(e)}") + + logger.info(f"知识库文档已处理: {stored_path}, ext={file_ext}, size={actual_size}, content_len={len(content)}") + + # 文本分块 + chunks = _split_content_into_chunks(content, chunk_size=500) + + doc = KnowledgeDocument( + id=doc_id, + knowledge_base_id=base_id, + filename=filename, + file_size=actual_size, + file_type=file_ext, + status="completed", + chunk_count=len(chunks), + created_at=datetime.now(), + updated_at=datetime.now(), + ) + db.add(doc) + + # 更新知识库统计 + kb.document_count = (kb.document_count or 0) + 1 + kb.updated_at = datetime.now() + + # 同步创建Asset记录 + _sync_asset_for_knowledge_doc(db, current_user.id, doc_id, filename, file_ext, actual_size) + + db.commit() + db.refresh(doc) + + # T1: 异步索引到 pgvector 语义库（best-effort，不阻塞上传响应） + try: + asyncio.create_task(_index_kb_doc_to_pgvector(doc.id, base_id, content)) + except Exception: + pass + + LogService.log_operation( + user_id=current_user.id, username=current_user.username, action="upload", + target="knowledge_document", target_id=doc.id, + detail=f"上传文档 {doc.filename} 到知识库 {kb.name}, 存储路径: {stored_path}, 分块数: {len(chunks)}", + ) + + return { + "code": 0, + "data": { + "id": doc.id, "base_id": doc.knowledge_base_id, "filename": doc.filename, + "file_size": doc.file_size, "file_type": doc.file_type, + "status": "completed", "chunk_count": doc.chunk_count, + "created_at": str(doc.created_at), + "stored_path": stored_path, + }, + "message": "上传成功", + } + + +def _generate_document_content(filename: str, file_ext: str, file_size: int) -> str: + """根据文件类型生成真实的模拟文档内容（文本解析后的结果）""" + name_base = filename.rsplit(".", 1)[0] + paragraphs = [] + + if file_ext in ("txt", "md"): + paragraphs = [ + f"# {name_base}", + f"", + f"## 概述", + f"本文档记录了关于 {name_base} 的详细内容。该文档由用户上传至知识库，" + f"经过文本解析引擎处理后，提取了以下关键信息。", + f"", + f"## 核心内容", + f"1. 主题分析：{name_base} 涉及的核心领域包括数据管理、知识工程和智能分析。", + f"2. 关键指标：通过对文档内容的深度解析，识别出 {max(3, file_size // 10000)} 个关键概念节点。", + f"3. 关联关系：文档中各主题之间存在 {max(2, file_size // 20000)} 条语义关联。", + f"", + f"## 详细分析", + f"在 {name_base} 的深入研究中，我们发现以下几点值得关注：", + f"- 数据结构层面的优化空间约为 15-20%", + f"- 知识图谱构建的准确率达到 87.3%", + f"- 检索效率相比传统方法提升了 2.3 倍", + f"", + f"## 结论与建议", + f"基于以上分析，建议在以下几个方面进行改进：", + f"1. 持续优化知识提取的准确度", + f"2. 扩展实体识别覆盖范围", + f"3. 增强跨文档关联分析能力", + f"", + f"---", + f"*本文档由 AIFlow 知识库系统自动解析生成，上载时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*", + ] + elif file_ext == "pdf": + paragraphs = [ + f"PDF文档: {name_base}", + f"", + f"【文档信息】", + f"标题：{name_base}", + f"页数估计：{max(1, file_size // 50000)} 页", + f"创建者：AIFlow 用户", + f"", + f"【提取的文本内容】", + f"第一章 引言", + f"本文档旨在阐述 {name_base} 相关的核心概念和实践方法。", + f"随着人工智能技术的不断发展，知识管理系统在企业数字化转型中扮演着越来越重要的角色。", + f"", + f"第二章 方法论", + f"本研究采用混合方法，结合定量分析与定性评估，对 {name_base} 进行全面探讨。", + f"数据来源包括：内部数据库查询结果、用户行为日志、以及与领域专家的深度访谈。", + f"", + f"第三章 结果分析", + f"通过对收集数据的系统分析，得出以下关键发现：", + f"1. 核心指标X的均值为 78.5，标准差为 12.3", + f"2. 关键因素Y与Z之间存在显著相关性（r=0.72, p<0.001）", + f"3. 85% 的受访者认为 {name_base} 是未来发展的关键方向", + f"", + f"第四章 结论", + f"综合以上分析，{name_base} 具有重要的研究价值和实践意义。", + ] + elif file_ext in ("docx", "doc"): + paragraphs = [ + f"Word文档: {name_base}", + f"", + f"{name_base} - 正式文档", + f"", + f"【摘要】", + f"本报告围绕 {name_base} 展开全面分析，涵盖市场趋势、技术演进和应用场景三个维度。", + f"通过系统性的研究和数据采集，为相关决策提供参考依据。", + f"", + f"【正文】", + f"一、背景介绍", + f"在当前的数字化转型浪潮中，{name_base} 已经成为组织提升竞争力的关键要素。", + f"根据Gartner的报告，到2026年将有超过60%的企业采用AI驱动的知识管理方案。", + f"", + f"二、技术架构", + f"系统采用微服务架构设计，包括数据采集层、处理引擎层、存储层和应用层四个核心模块。", + f"其中处理引擎层集成了NLP自然语言处理、实体识别和关系抽取等多种AI能力。", + f"", + f"三、实施计划", + f"第一阶段（1-2月）：完成基础架构搭建和数据迁移", + f"第二阶段（3-4月）：部署AI模型并开始增量学习", + f"第三阶段（5-6月）：全面上线并进行效果评估", + ] + elif file_ext in ("xlsx", "xls"): + paragraphs = [ + f"Excel表格: {name_base}", + f"", + f"【表格数据概览】", + f"| 指标类别 | Q1 | Q2 | Q3 | Q4 | 年度总计 |", + f"|---------|-----|-----|-----|-----|---------|", + f"| 用户增长 | 12% | 18% | 22% | 28% | 80% |", + f"| 营收增幅 | 8% | 14% | 19% | 25% | 66% |", + f"| 客户满意度 | 85 | 88 | 91 | 93 | 89.3 |", + f"| 市场份额 | 15% | 17% | 20% | 23% | -- |", + f"", + f"【关键指标分析】", + f"1. 用户增长率呈逐季上升趋势，Q4达到峰值28%", + f"2. 营收增长与用户增长高度正相关", + f"3. 客户满意度持续改善，Q4达到93分", + f"4. 市场份额稳步扩大，年度总增长约8个百分点", + f"", + f"【数据来源】内部经营分析系统，数据截止日期：{datetime.now().strftime('%Y-%m-%d')}", + ] + elif file_ext in ("pptx", "ppt"): + paragraphs = [ + f"PPT演示文稿: {name_base}", + f"", + f"【幻灯片结构】", + f"第1页：封面 - {name_base}", + f"第2页：目录与议程", + f"第3页：背景与现状分析", + f"第4页：核心挑战与机遇", + f"第5页：解决方案概述", + f"第6-8页：详细方案展开", + f"第9页：实施路线图", + f"第10页：总结与展望", + f"", + f"【关键要点提取】", + f"- 战略目标：实现 {name_base} 的全面数字化转型", + f"- 核心路径：AI赋能 + 数据驱动 + 敏捷迭代", + f"- 预期成果：效率提升30%、成本降低20%、用户满意度提升15%", + f"- 投入预算：约500万元（分三期投入）", + ] + else: + paragraphs = [ + f"文档: {name_base}", + f"", + f"文件类型：{file_ext.upper()}", + f"文件大小：{file_size} bytes", + f"上传时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", + f"", + f"【内容摘要】", + f"本文档类型为 {file_ext.upper()} 格式，包含与 {name_base} 相关的内容。", + f"系统已完成基本的文本解析和索引建立，支持关键词检索和语义匹配。", + ] + + return "\n".join(paragraphs) + + +def _split_content_into_chunks(content: str, chunk_size: int = 500) -> list: + """将文档内容分割为文本块（模拟embedding的chunking过程）""" + chunks = [] + paragraphs = content.split("\n\n") + current_chunk = "" + + for para in paragraphs: + if len(current_chunk) + len(para) + 2 <= chunk_size: + current_chunk = (current_chunk + "\n\n" + para).strip() + else: + if current_chunk: + chunks.append(current_chunk) + if len(para) > chunk_size: + # 长段落进一步切分 + for i in range(0, len(para), chunk_size): + chunks.append(para[i:i+chunk_size]) + current_chunk = "" + else: + current_chunk = para + + if current_chunk: + chunks.append(current_chunk) + + return chunks if chunks else [content[:chunk_size]] + + +def _sync_asset_for_knowledge_doc(db: Session, user_id: str, doc_id: str, filename: str, file_type: str, file_size: int): + """将知识库文档同步为资产记录（问题9：资产关联知识库文档用量）""" + try: + # 文件类型归一化为前端子分类标签（docx/doc->word 等），确保出现在「文件 > Word」子栏 + _ft_label = { + "pdf": "pdf", "docx": "word", "doc": "word", "xlsx": "excel", "xls": "excel", + "pptx": "ppt", "ppt": "ppt", "csv": "txt", "md": "md", "txt": "txt", + }.get((file_type or "").lower(), "other") + asset = Asset( + id=gen_uuid(), + user_id=user_id, + name=filename, + category="file", + description=f"知识库文档 (ID: {doc_id})", + file_type=_ft_label, + size=file_size, + status="active", + knowledge_doc_id=doc_id, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + db.add(asset) + logger.info(f"已同步知识库文档到资产: {filename} -> asset_id={asset.id}") + except Exception as e: + logger.warning(f"同步知识库文档到资产失败: {e}") + + +@router.delete("/bases/{base_id}/documents/{doc_id}") +def delete_document( + base_id: str, + doc_id: str, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + doc = db.query(KnowledgeDocument).filter( + and_(KnowledgeDocument.id == doc_id, KnowledgeDocument.knowledge_base_id == base_id) + ).first() + if not doc: + raise HTTPException(status_code=404, detail="文档不存在") + + # 更新知识库统计 + kb.document_count = max(0, (kb.document_count or 0) - 1) + pass # total_size removed + kb.updated_at = datetime.now() + + # 同步删除知识库文档对应的资产记录（双向同步） + try: + linked_asset = db.query(Asset).filter(Asset.knowledge_doc_id == doc_id).first() + if linked_asset: + linked_asset.status = "deleted" + linked_asset.updated_at = datetime.now() + except Exception as _e: + logger.warning(f"同步删除关联资产失败: {_e}") + + db.delete(doc) + db.commit() + + return {"code": 0, "message": "文档已删除"} + + +@router.get("/directories") +def list_directories( + base_id: Optional[str] = Query(None), + parent_id: Optional[str] = Query(None), + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + query = db.query(KnowledgeDirectory).filter( + KnowledgeDirectory.user_id == current_user.id + ) + + if base_id: + query = query.filter(KnowledgeDirectory.knowledge_base_id == base_id) + if parent_id is not None: + query = query.filter(KnowledgeDirectory.parent_id == parent_id) + else: + # 默认查根目录 + query = query.filter(KnowledgeDirectory.parent_id.is_(None)) + + directories = query.order_by(asc(KnowledgeDirectory.name)).all() + + result = [] + for d in directories: + result.append({ + "id": d.id, + "name": d.name, + "parent_id": d.parent_id, + "knowledge_base_id": d.knowledge_base_id, + "created_at": str(d.created_at) if d.created_at else None, + }) + + return {"code": 0, "data": result, "message": "ok"} + + +@router.post("/directories") +def create_directory( + req: CreateDirectoryRequest, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + directory = KnowledgeDirectory( + id=gen_uuid(), + user_id=current_user.id, + name=req.name, + parent_id=req.parent_id, + knowledge_base_id=req.knowledge_base_id, + created_at=datetime.now(), + ) + db.add(directory) + db.commit() + db.refresh(directory) + + return { + "code": 0, + "data": { + "id": directory.id, "name": directory.name, + "parent_id": directory.parent_id, "knowledge_base_id": directory.knowledge_base_id, + "created_at": str(directory.created_at), + }, + "message": "创建成功", + } + + +# ============================================================ +# 知识库问答溯源 API +# ============================================================ + +@router.post("/qa") +def knowledge_qa( + req: KnowledgeQARequest, + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + """知识库问答：在用户的知识库中检索相关文档并生成回答，返回原文引用和来源路径""" + # 验证知识库归属 + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == req.base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + # 检索知识库中的文档 + documents = db.query(KnowledgeDocument).filter( + and_( + KnowledgeDocument.knowledge_base_id == req.base_id, + KnowledgeDocument.status == "completed", + ) + ).all() + + if not documents: + return { + "code": 0, + "data": { + "question": req.question, + "answer": "该知识库中暂无已完成的文档。请先上传并处理文档后再进行问答。", + "sources": [], + "knowledge_base_name": kb.name, + }, + "message": "ok", + } + + # 构建引用来源列表（基于文档关键词匹配模拟检索） + import re + sources = [] + relevant_docs = [] + question_lower = req.question.lower() + + for doc in documents: + filename_lower = doc.filename.lower() + # 简单的关键词匹配检索（实际生产环境应使用向量检索） + keywords = re.split(r'[\s\-\_\.\,\;\:\!\?\(\)\[\]\{\}]+', question_lower) + filename_words = re.split(r'[\s\-\_\.\,\;\:\!\?\(\)\[\]\{\}]+', filename_lower) + match_count = sum(1 for kw in keywords if kw and any(kw in fw for fw in filename_words)) + + if match_count > 0 or len(documents) <= 3: + relevant_docs.append(doc) + # 为每个匹配文档生成模拟原文引用 + sample_content = generate_document_excerpt(doc, req.question) + sources.append({ + "document_id": doc.id, + "document_name": doc.filename, + "file_type": doc.file_type or "", + "file_size": doc.file_size or 0, + "source_path": f"/knowledge/{kb.id}/documents/{doc.id}", + "excerpt": sample_content, + "relevance": min(0.99, max(0.3, match_count * 0.2 + 0.3)) if match_count > 0 else 0.3, + }) + + # 生成模拟回答 + kb_name = kb.name + answer_parts = [f"根据知识库「{kb_name}」中相关文档的内容，针对您的问题回复如下：\n"] + + if sources: + answer_parts.append(f"共检索到 {len(sources)} 篇相关文档，以下为引用分析：\n") + for i, src in enumerate(sources[:5], 1): + answer_parts.append(f"\n【引用 {i}】来源：《{src['document_name']}》") + answer_parts.append(f"原文摘录：\"{src['excerpt']}\"") + answer_parts.append(f"文档类型：{src['file_type'].upper()} | 文档大小：{format_file_size(src['file_size'])}") + answer_parts.append(f"\n---\n综合以上文档内容，我们建议您参考上述引用来源获取详细信息。") + else: + answer_parts.append("未找到与问题高度相关的文档内容，建议上传更多相关文档后重试。") + + final_answer = "\n".join(answer_parts) + + return { + "code": 0, + "data": { + "question": req.question, + "answer": final_answer, + "sources": sources, + "knowledge_base_name": kb_name, + "total_documents": len(documents), + }, + "message": "ok", + } + + +@router.get("/bases/{base_id}/search") +def search_knowledge( + base_id: str, + keyword: str = Query(min_length=1), + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + """知识库文档检索：按关键词搜索文档并返回引用片段""" + kb = db.query(KnowledgeBase).filter( + and_(KnowledgeBase.id == base_id, KnowledgeBase.user_id == current_user.id) + ).first() + if not kb: + raise HTTPException(status_code=404, detail="知识库不存在") + + documents = db.query(KnowledgeDocument).filter( + and_( + KnowledgeDocument.knowledge_base_id == base_id, + KnowledgeDocument.status == "completed", + KnowledgeDocument.filename.ilike(f"%{keyword}%"), + ) + ).all() + + result = [] + for doc in documents: + result.append({ + "document_id": doc.id, + "document_name": doc.filename, + "file_type": doc.file_type or "", + "file_size": doc.file_size or 0, + "source_path": f"/knowledge/{kb.id}/documents/{doc.id}", + "excerpt": generate_document_excerpt(doc, keyword), + "created_at": str(doc.created_at) if doc.created_at else None, + }) + + return { + "code": 0, + "data": { + "keyword": keyword, + "knowledge_base_name": kb.name, + "results": result, + "total": len(result), + }, + "message": "ok", + } + + +@router.get("/tree-search") +async def tree_search_endpoint( + query: str = Query(min_length=1, description="检索查询"), + kb_id: Optional[str] = Query(default=None, description="限定知识库ID"), + scope: Optional[str] = Query( + default=None, + description="作用域: doc: | dir: | base:", + ), + top_k: int = Query(default=5, ge=1, le=20), + db: Session = Depends(get_db), + current_user: User = Depends(get_current_user), +): + """TreeSearch 结构化文档两级检索。 + + Level 1: 返回用户可见的知识库 -> 目录 -> 文档结构树(tree)。 + Level 2: 在作用域内做分片段(chunk)级精准检索，命中带完整面包屑与片段摘要。 + """ + from app.tree_search import build_kb_tree, tree_search + + # Level 1: 结构树 + tree = build_kb_tree(db, user_id=current_user.id, kb_id=kb_id) + + # Level 2: 分片段检索 + scope_doc = scope_dir = scope_base = None + if scope: + if scope.startswith("doc:"): + scope_doc = scope[len("doc:"):] + elif scope.startswith("dir:"): + scope_dir = scope[len("dir:"):] + elif scope.startswith("base:"): + scope_base = scope[len("base:"):] + + hits = await tree_search( + db, query, + user_id=current_user.id, + kb_ids=[kb_id] if kb_id else None, + scope_base=scope_base, scope_dir=scope_dir, scope_doc=scope_doc, + top_k=top_k, mode="hybrid", + ) + + return { + "code": 0, + "data": { + "query": query, + "tree": tree, + "hits": hits, + "total": len(hits), + }, + "message": "ok", + } + + +def generate_document_excerpt(doc, query: str) -> str: + """根据文档类型生成模拟的原文摘录""" + excerpts = { + "pdf": f"本文档详细阐述了与「{query}」相关的核心内容。根据第3章第2节的分析，关键指标呈现明显的增长趋势，其中核心参数值在统计周期内提升了约27.3%。进一步的研究表明，该趋势与市场整体发展方向高度一致...", + "docx": f"关于「{query}」的讨论，文档中明确指出：\"从数据层面来看，该领域在过去四个季度中保持了平均15%的复合增长率。这一结论基于对超过2000个样本的综合分析，置信度达到95%以上。建议决策层关注以下三个关键方向...\"", + "xlsx": f"数据表格「{query}」相关统计：\n| 指标 | Q1 | Q2 | Q3 | Q4 |\n|------|----|----|----|----|\n| 用户增长 | 12% | 18% | 22% | 28% |\n| 营收增幅 | 8% | 14% | 19% | 25% |\n数据来源：内部经营分析系统", + "ppt": f"演示文稿第15页「{query}」关键要点：\n• 战略目标：实现年度增长30%\n• 核心路径：优化产品矩阵，拓展海外市场\n• 资源配置：预算投入约500万元\n• 时间节点：Q2完成试点，Q4全面推广\n详细方案见第15-22页", + "md": f"# {query}\n\n## 核心观点\n\n文档明确指出：*\"在当前的行业格局下，企业需要通过技术创新和模式变革来建立竞争壁垒。具体来说，可以从以下三个维度切入：数据驱动决策、自动化流程优化、以及个性化用户体验。\"*\n\n## 关键数据\n\n- 市场规模：预计2026年达到200亿\n- 年复合增长率：35%\n- 头部企业市占率：约40%\n\n> 详细分析参见完整文档。", + "txt": f"文档内容摘录：\n\n\"...针对「{query}」这一课题，我们进行了深入的调研和分析。研究发现，超过78%的受访者认为该领域的创新速度正在加快。同时，技术成熟度曲线显示，相关技术有望在未来12-18个月内进入主流应用阶段。建议前期投入研发资源约50人月，预期投资回报周期为18个月...\"", + } + file_type = (doc.file_type or "txt").lower() + return excerpts.get(file_type, excerpts["txt"]) + + +def format_file_size(size_bytes: int) -> str: + """格式化文件大小""" + if size_bytes >= 1024 * 1024: + return f"{size_bytes / (1024 * 1024):.1f} MB" + elif size_bytes >= 1024: + return f"{size_bytes / 1024:.1f} KB" + return f"{size_bytes} B" + + +# ============================================================ +# 知识图谱构建与查询 API +# ============================================================ + +class BuildGraphRequest(BaseModel): + base_id: str + diff --git a/backend/python/app/kb_router.py b/backend/python/app/kb_router.py index f822749..5632bdc 100644 --- a/backend/python/app/kb_router.py +++ b/backend/python/app/kb_router.py @@ -1,175 +1,204 @@ -""" -知识库路由层 -============ -将检索请求智能路由到本地BM25、向量检索、或外部KB服务(Dify)。 -根据知识库配置自动选择最优检索策略。 - -使用方式: - from app.kb_router import kb_router_search - results = await kb_router_search(query, kb_id, top_k=5) -""" - -import logging -from typing import List, Dict, Optional - -logger = logging.getLogger("aiflow.kb_router") - - -async def _get_external_services(db_session, enabled_only: bool = True) -> List[Dict]: - """获取已配置的外部KB服务（仅 dify）""" - from sqlalchemy import text - try: - query = "SELECT * FROM external_kb_services" - if enabled_only: - query += " WHERE is_enabled = 1" - query += " ORDER BY created_at DESC" - result = db_session.execute(text(query)) - columns = result.keys() - return [dict(zip(columns, row)) for row in result.fetchall()] - except Exception as e: - logger.error(f"Failed to get external services: {e}") - return [] - - -async def _search_dify(service: Dict, query: str, top_k: int = 5) -> List[Dict]: - """检索 Dify 知识库（走 app.dify_kb 客户端）""" - try: - from app.dify_kb import dify_retrieve - api_url = service.get("api_url", "") - api_key = service.get("api_key", "") - dataset_name = None - cfg = service.get("config") - if isinstance(cfg, str): - try: - import json - cfg = json.loads(cfg) - except Exception: - cfg = None - if isinstance(cfg, dict): - dataset_name = cfg.get("dataset_name") - hits = dify_retrieve(api_url, api_key, query, top_k=top_k, dataset_name=dataset_name) - return hits - except Exception as e: - logger.error(f"Dify search error: {e}") - return [] - - -async def _search_external(service: Dict, query: str, top_k: int = 5) -> List[Dict]: - """路由到对应的外部KB服务（当前仅支持 dify）""" - service_type = service.get("service_type", "") - if service_type == "dify": - return await _search_dify(service, query, top_k) - else: - logger.warning(f"Unknown external KB type: {service_type}") - return [] - - -async def kb_router_search( - query: str, - kb_id: Optional[str] = None, - top_k: int = 5, - bm25_results: Optional[List[Dict]] = None, - db_session=None, - use_vector: bool = True, - use_external: bool = True, -) -> List[Dict]: - """ - 知识库路由层 — 混合检索 - - 检索策略: - 1. 本地 BM25 检索 (MySQL, 已有) - 2. 向量语义检索 (PostgreSQL+pgvector, 新增) - 3. 外部KB检索 (Dify, 按配置) - - 融合方式: Reciprocal Rank Fusion (RRF) - - Args: - query: 查询文本 - kb_id: 知识库ID - top_k: 返回结果数 - bm25_results: BM25检索结果 - db_session: 数据库会话(用于获取外部KB配置) - use_vector: 是否使用向量检索 - use_external: 是否使用外部KB检索 - - Returns: - 融合后的检索结果列表 - """ - all_results = [] - - # 1. BM25 结果 - if bm25_results: - for rank, item in enumerate(bm25_results): - all_results.append({ - "content": item.get("content", item.get("text", "")), - "score": item.get("score", 1.0 / (60 + rank + 1)), - "source": "local_bm25", - "rank": rank + 1, - "metadata": item.get("metadata", {}), - }) - - # 2. 向量检索 - if use_vector: - try: - from app.pgvector_search import vector_search - vec_results = await vector_search(query, kb_id, top_k) - for rank, item in enumerate(vec_results): - all_results.append({ - "content": item.get("content", ""), - "score": item.get("score", 0.5), - "source": "vector_pg", - "rank": rank + 1, - "metadata": item.get("metadata", {}), - }) - except ImportError: - logger.debug("pgvector_search not available, skipping vector search") - except Exception as e: - logger.warning(f"Vector search failed: {e}") - - # 3. 外部KB检索 (Dify) - if use_external and db_session: - services = await _get_external_services(db_session) - for service in services: - if service.get("status") == "connected": - ext_results = await _search_external(service, query, top_k) - for rank, item in enumerate(ext_results): - all_results.append({ - "content": item.get("content", ""), - "score": item.get("score", 0.5), - "source": item.get("source", "external_dify"), - "rank": rank + 1, - "metadata": item.get("metadata", {}), - }) - - if not all_results: - return [] - - # RRF 融合 - rrf_k = 60 - score_map = {} - for item in all_results: - content_key = item["content"][:200] # 用内容前200字符作为去重键 - rrf_score = 1.0 / (rrf_k + item.get("rank", 1)) - - if content_key in score_map: - score_map[content_key]["score"] += rrf_score - score_map[content_key]["sources"].append(item["source"]) - else: - score_map[content_key] = { - "content": item["content"], - "score": rrf_score, - "sources": [item["source"]], - "metadata": item.get("metadata", {}), - } - - # 按融合分数排序 - merged = sorted(score_map.values(), key=lambda x: x["score"], reverse=True) - - # 标准化分数到 0-1 - if merged: - max_score = merged[0]["score"] - for item in merged: - item["score"] = item["score"] / max_score if max_score > 0 else 0 - item["sources"] = list(set(item["sources"])) - - return merged[:top_k] +""" +知识库路由层 +============ +将检索请求智能路由到本地BM25、向量检索、或外部KB服务(Dify)。 +根据知识库配置自动选择最优检索策略。 + +使用方式: + from app.kb_router import kb_router_search + results = await kb_router_search(query, kb_id, top_k=5) +""" + +import logging +from typing import List, Dict, Optional + +logger = logging.getLogger("aiflow.kb_router") + + +async def _get_external_services(db_session, enabled_only: bool = True) -> List[Dict]: + """获取已配置的外部KB服务（仅 dify）""" + from sqlalchemy import text + try: + query = "SELECT * FROM external_kb_services" + if enabled_only: + query += " WHERE is_enabled = 1" + query += " ORDER BY created_at DESC" + result = db_session.execute(text(query)) + columns = result.keys() + return [dict(zip(columns, row)) for row in result.fetchall()] + except Exception as e: + logger.error(f"Failed to get external services: {e}") + return [] + + +async def _search_dify(service: Dict, query: str, top_k: int = 5) -> List[Dict]: + """检索 Dify 知识库（走 app.dify_kb 客户端）""" + try: + from app.dify_kb import dify_retrieve + api_url = service.get("api_url", "") + api_key = service.get("api_key", "") + dataset_name = None + cfg = service.get("config") + if isinstance(cfg, str): + try: + import json + cfg = json.loads(cfg) + except Exception: + cfg = None + if isinstance(cfg, dict): + dataset_name = cfg.get("dataset_name") + hits = dify_retrieve(api_url, api_key, query, top_k=top_k, dataset_name=dataset_name) + return hits + except Exception as e: + logger.error(f"Dify search error: {e}") + return [] + + +async def _search_external(service: Dict, query: str, top_k: int = 5) -> List[Dict]: + """路由到对应的外部KB服务（当前仅支持 dify）""" + service_type = service.get("service_type", "") + if service_type == "dify": + return await _search_dify(service, query, top_k) + else: + logger.warning(f"Unknown external KB type: {service_type}") + return [] + + +async def kb_router_search( + query: str, + kb_id: Optional[str] = None, + kb_ids: Optional[List[str]] = None, + top_k: int = 5, + bm25_results: Optional[List[Dict]] = None, + db_session=None, + use_vector: bool = True, + use_external: bool = True, + use_tree: bool = True, +) -> List[Dict]: + """ + 知识库路由层 — 混合检索 + + 检索策略: + 1. 本地 BM25 检索 (MySQL, 已有) + 2. 向量语义检索 (PostgreSQL+pgvector, 新增) + 3. 结构化分片段检索 (TreeSearch, 新增: 知识库目录树 + chunk 级精准命中) + 4. 外部KB检索 (Dify, 按配置) + + 融合方式: Reciprocal Rank Fusion (RRF) + + Args: + query: 查询文本 + kb_id: 知识库ID + top_k: 返回结果数 + bm25_results: BM25检索结果 + db_session: 数据库会话(用于获取外部KB配置与分片段检索) + use_vector: 是否使用向量检索 + use_external: 是否使用外部KB检索 + use_tree: 是否使用 TreeSearch 结构化分片段检索 + + Returns: + 融合后的检索结果列表 + """ + all_results = [] + + # 1. BM25 结果 + if bm25_results: + for rank, item in enumerate(bm25_results): + all_results.append({ + "content": item.get("content", item.get("text", "")), + "score": item.get("score", 1.0 / (60 + rank + 1)), + "source": "local_bm25", + "rank": rank + 1, + "metadata": item.get("metadata", {}), + }) + + # 2. 向量检索 + if use_vector: + try: + from app.pgvector_search import vector_search + kb_targets = kb_ids if kb_ids else ([kb_id] if kb_id else [None]) + vec_results = [] + for _kid in kb_targets: + try: + vec_results.extend(await vector_search(query, _kid, top_k)) + except Exception as e: + logger.warning(f"Vector search failed for kb={_kid}: {e}") + for rank, item in enumerate(vec_results): + all_results.append({ + "content": item.get("content", ""), + "score": item.get("score", 0.5), + "source": "vector_pg", + "rank": rank + 1, + "metadata": item.get("metadata", {}), + }) + except ImportError: + logger.debug("pgvector_search not available, skipping vector search") + except Exception as e: + logger.warning(f"Vector search failed: {e}") + + # 2.5 结构化分片段检索 (TreeSearch, Level 2) + if use_tree and db_session is not None: + try: + from app.tree_search import tree_search + kb_targets = kb_ids if kb_ids else ([kb_id] if kb_id else None) + tree_hits = await tree_search( + db_session, query, kb_ids=kb_targets, top_k=top_k, mode="hybrid", + ) + for rank, item in enumerate(tree_hits): + all_results.append({ + "content": item.get("content", ""), + "score": item.get("score", 0.5), + "source": "tree_chunk", + "rank": rank + 1, + "metadata": item.get("metadata", {}), + }) + except Exception as e: + logger.warning(f"TreeSearch failed: {e}") + + # 3. 外部KB检索 (Dify) + if use_external and db_session: + services = await _get_external_services(db_session) + for service in services: + if service.get("status") == "connected": + ext_results = await _search_external(service, query, top_k) + for rank, item in enumerate(ext_results): + all_results.append({ + "content": item.get("content", ""), + "score": item.get("score", 0.5), + "source": item.get("source", "external_dify"), + "rank": rank + 1, + "metadata": item.get("metadata", {}), + }) + + if not all_results: + return [] + + # RRF 融合 + rrf_k = 60 + score_map = {} + for item in all_results: + content_key = item["content"][:200] # 用内容前200字符作为去重键 + rrf_score = 1.0 / (rrf_k + item.get("rank", 1)) + + if content_key in score_map: + score_map[content_key]["score"] += rrf_score + score_map[content_key]["sources"].append(item["source"]) + else: + score_map[content_key] = { + "content": item["content"], + "score": rrf_score, + "sources": [item["source"]], + "metadata": item.get("metadata", {}), + } + + # 按融合分数排序 + merged = sorted(score_map.values(), key=lambda x: x["score"], reverse=True) + + # 标准化分数到 0-1 + if merged: + max_score = merged[0]["score"] + for item in merged: + item["score"] = item["score"] / max_score if max_score > 0 else 0 + item["sources"] = list(set(item["sources"])) + + return merged[:top_k] diff --git a/backend/python/app/tree_search.py b/backend/python/app/tree_search.py new file mode 100644 index 0000000..f517eed --- /dev/null +++ b/backend/python/app/tree_search.py @@ -0,0 +1,405 @@ +""" +TreeSearch — 结构化文档两级检索 +================================ + +为知识库提供有结构的「两级检索」能力（零额外基础设施成本，复用 MySQL 已存的分段数据）： + + Level 1 (build_kb_tree) : 知识库 -> 目录 -> 文档 的嵌套结构树，用于浏览/定位/面包屑。 + Level 2 (tree_search) : 在给定作用域（整个 KB / 某目录 / 某文档）内做分段(chunk)级检索， + BM25 关键词排序 + 可选 pgvector 语义向量 boost，返回带完整 + 面包屑(KB名 -> 目录路径 -> 文件名 -> 段序号/标题)的精准命中。 + +设计要点： + * Level 2 直接消费 knowledge_document_chunks 表（content + chunk_metadata.title/keywords）， + 因此召回的是「文档内部的精确片段」，而非整篇文档——这正是 RAG 想要的细粒度上下文。 + * 与 kb_router_search 融合：tree_search 作为 RRF 的一个新来源 (source="tree_chunk") 接入， + 聊天 KB 上下文自动获得结构化分片段命中。 + * 向量 boost 为可选且优雅降级：当 pgvector 无对应数据时，自动退化为纯 BM25，不影响主链路。 +""" + +import logging +import math +from typing import List, Dict, Optional + +logger = logging.getLogger("aiflow.tree_search") + + +# --------------------------------------------------------------------------- +# 中文友好的分词（jieba 优先，否则按字符 bigram + 英文正则兜底） +# --------------------------------------------------------------------------- +def _tokenize(text: str) -> List[str]: + if not text: + return [] + toks: List[str] = [] + try: + import re + # 英文/数字 token + for m in re.findall(r"[a-zA-Z0-9_]+", text.lower()): + if len(m) >= 1: + toks.append(m) + # 中文按 jieba 或字符切分 + cjk = "".join(re.findall(r"[一-鿿]", text)) + if cjk: + try: + import jieba + for w in jieba.lcut(cjk): + w = w.strip().lower() + if w: + toks.append(w) + except Exception: + # 字符级 bigram（兼顾召回与精度） + for i in range(len(cjk)): + toks.append(cjk[i]) + if i + 1 < len(cjk): + toks.append(cjk[i:i + 2]) + except Exception: + toks = [t for t in text.lower().split() if t] + return toks + + +def _chunk_display_text(chunk) -> str: + """拼装分段的检索文本：正文 + 标题 + 关键词（来自 chunk_metadata）。""" + parts = [] + content = getattr(chunk, "content", None) or "" + parts.append(content) + md = getattr(chunk, "chunk_metadata", None) or {} + if isinstance(md, dict): + title = md.get("title") or md.get("heading") or "" + keywords = md.get("keywords") or [] + if title: + parts.append(title) + if isinstance(keywords, list) and keywords: + parts.append(" ".join(str(k) for k in keywords)) + elif isinstance(keywords, str) and keywords: + parts.append(keywords) + return "\n".join(p for p in parts if p).strip() + + +def _chunk_title(chunk) -> str: + md = getattr(chunk, "chunk_metadata", None) or {} + if isinstance(md, dict): + return md.get("title") or md.get("heading") or "" + return "" + + +# --------------------------------------------------------------------------- +# BM25（在作用域内的分段语料上即时计算 IDF） +# --------------------------------------------------------------------------- +def _bm25_rank(query_tokens: List[str], corpus: List[Dict], k1: float = 1.5, b: float = 0.75) -> List[Dict]: + """corpus: [{"id":..., "tokens":[...], "chunk":, "text":str}] -> 返回带 bm25 分的列表""" + if not query_tokens or not corpus: + return [] + n = len(corpus) + df = {} + for qt in set(query_tokens): + df[qt] = sum(1 for d in corpus if qt in d["tokens"]) + idf = {qt: math.log((n - df[qt] + 0.5) / (df[qt] + 0.5) + 1) for qt in query_tokens} + avgdl = sum(len(d["tokens"]) for d in corpus) / max(n, 1) + scored = [] + for d in corpus: + dl = len(d["tokens"]) + score = 0.0 + for qt in set(query_tokens): + if qt in d["tokens"]: + f = d["tokens"].count(qt) + score += idf[qt] * (f * (k1 + 1)) / (f + k1 * (1 - b + b * (dl / max(avgdl, 1)))) + if score > 0: + scored.append({**d, "bm25": score}) + scored.sort(key=lambda x: x["bm25"], reverse=True) + return scored + + +# --------------------------------------------------------------------------- +# Level 1：知识库目录树 +# --------------------------------------------------------------------------- +def build_kb_tree(db, user_id: Optional[str] = None, kb_id: Optional[str] = None) -> List[Dict]: + """构造 知识库 -> 目录(可嵌套) -> 文档 的结构树。""" + from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument + + q = db.query(KnowledgeBase) + if user_id: + q = q.filter(KnowledgeBase.user_id == user_id) + if kb_id: + q = q.filter(KnowledgeBase.id == kb_id) + bases = q.all() + + tree = [] + for kb in bases: + dirs = db.query(KnowledgeDirectory).filter( + KnowledgeDirectory.knowledge_base_id == kb.id + ).all() + docs = db.query(KnowledgeDocument).filter( + KnowledgeDocument.knowledge_base_id == kb.id, + KnowledgeDocument.status == "completed", + ).all() + dir_by_id = {d.id: d for d in dirs} + node_by_id = {} + + def make_dir_node(d): + return { + "id": d.id, + "name": d.name, + "type": "directory", + "children": [], + "documents": [], + } + + for d in dirs: + node_by_id[d.id] = make_dir_node(d) + root_dirs = [] + for d in dirs: + node = node_by_id[d.id] + parent = d.parent_id + if parent and parent in node_by_id: + node_by_id[parent]["children"].append(node) + else: + root_dirs.append(node) + for doc in docs: + item = { + "id": doc.id, + "name": doc.filename, + "type": "document", + "file_type": doc.file_type, + "chunk_count": doc.chunk_count or 0, + } + if doc.directory_id and doc.directory_id in node_by_id: + node_by_id[doc.directory_id]["documents"].append(item) + else: + root_dirs.append(item) # 顶层文档 + tree.append({ + "id": kb.id, + "name": kb.name, + "type": "knowledge_base", + "directories": root_dirs, + "document_count": len(docs), + }) + return tree + + +# --------------------------------------------------------------------------- +# 面包屑：KB名 -> 目录路径 -> 文件名 -> 段标题/序号 +# --------------------------------------------------------------------------- +def _breadcrumb(db, kb_id: str, directory_id: Optional[str], doc_id: str, chunk_index: int, title: str) -> Dict: + from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument + kb = db.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first() + kb_name = kb.name if kb else kb_id + path = [] + cur = directory_id + guard = 0 + while cur and guard < 20: + d = db.query(KnowledgeDirectory).filter(KnowledgeDirectory.id == cur).first() + if not d: + break + path.insert(0, d.name) + cur = d.parent_id + guard += 1 + doc = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first() + doc_name = doc.filename if doc else doc_id + return { + "knowledge_base": kb_name, + "directory_path": "/".join(path) if path else "/", + "document": doc_name, + "chunk_index": chunk_index, + "chunk_title": title, + "label": f"{kb_name}/{('/'.join(path) if path else '')}/{doc_name}" + (f"#段{chunk_index}" + (f"·{title}" if title else "")), + } + + +# --------------------------------------------------------------------------- +# Level 2：结构化分片段检索 +# --------------------------------------------------------------------------- +async def tree_search( + db, + query: str, + user_id: Optional[str] = None, + kb_ids: Optional[List[str]] = None, + scope_base: Optional[str] = None, + scope_dir: Optional[str] = None, + scope_doc: Optional[str] = None, + top_k: int = 5, + mode: str = "hybrid", + max_tokens: int = 1200, +) -> List[Dict]: + """ + 两级检索的 Level 2：在作用域内做分片段精准检索。 + + Args: + query : 查询文本 + kb_ids : 限定知识库范围（多个） + scope_base : 限定单一知识库 + scope_dir : 限定单一目录（含子目录） + scope_doc : 限定单一文档 + top_k : 返回命中数 + mode : "bm25" | "vector" | "hybrid"(默认) + Returns: + [{ + "content": <分段文本>, + "score": 0-1, + "source": "tree_chunk", + "metadata": { "breadcrumb": {...}, "snippet": ..., "chunk_index": int, "title": str } + }, ...] + """ + from app.models import KnowledgeBase, KnowledgeDirectory, KnowledgeDocument, KnowledgeDocumentChunk + + if not query or not query.strip(): + return [] + + # 1) 解析作用域 -> 目标文档 id 集合 + target_kb_ids = list(kb_ids or []) + if scope_base and scope_base not in target_kb_ids: + target_kb_ids.append(scope_base) + + doc_scope = set() + if scope_doc: + doc_scope.add(scope_doc) + if scope_dir: + # 收集该目录及其所有子目录下的文档 + dirs = db.query(KnowledgeDirectory).all() + dir_by_id = {d.id: d for d in dirs} + sub = set() + + def collect(did): + sub.add(did) + for d in dirs: + if d.parent_id == did: + collect(d.id) + collect(scope_dir) + docs = db.query(KnowledgeDocument).filter(KnowledgeDocument.directory_id.in_(sub)).all() + for d in docs: + doc_scope.add(d.id) + if not doc_scope and not target_kb_ids and not user_id: + return [] + + # 2) 取分段语料 + q = db.query(KnowledgeDocumentChunk) + if doc_scope: + q = q.filter(KnowledgeDocumentChunk.document_id.in_(doc_scope)) + else: + if target_kb_ids: + q = q.filter(KnowledgeDocumentChunk.knowledge_base_id.in_(target_kb_ids)) + elif user_id: + kb_ids_u = [r[0] for r in db.query(KnowledgeBase.id).filter(KnowledgeBase.user_id == user_id).all()] + if not kb_ids_u: + return [] + q = q.filter(KnowledgeDocumentChunk.knowledge_base_id.in_(kb_ids_u)) + chunks = q.order_by(KnowledgeDocumentChunk.chunk_index).limit(500).all() + if not chunks: + return [] + + query_tokens = _tokenize(query) + corpus = [] + for c in chunks: + text = _chunk_display_text(c) + if len(text.strip()) < 8: + continue + corpus.append({ + "id": c.id, + "tokens": _tokenize(text), + "chunk": c, + "text": text, + }) + + # 3) BM25 排序 + bm25_hits = _bm25_rank(query_tokens, corpus) if mode in ("bm25", "hybrid") else [] + + # 4) 可选向量 boost + vector_boost = {} + if mode in ("vector", "hybrid"): + try: + from app.pgvector_search import vector_search + for kid in (target_kb_ids or [None]): + try: + vres = await vector_search(query, kid, top_k=top_k * 3) + except Exception as e: + logger.warning(f"[tree_search] vector_search failed: {e}") + continue + for vr in vres: + vcontent = (vr.get("content") or "").strip() + if vcontent: + vector_boost[vcontent] = max(vector_boost.get(vcontent, 0.0), float(vr.get("score", 0.0))) + except Exception as e: + logger.debug(f"[tree_search] vector unavailable: {e}") + + # 5) 融合 + merged_map = {} + for rank, h in enumerate(bm25_hits): + key = h["text"][:300] + boost = 0.0 + for vc, vs in vector_boost.items(): + if vc and vc in h["text"]: + boost = max(boost, vs) + merged_map[key] = { + "chunk": h["chunk"], + "bm25": h["bm25"], + "vec": boost, + "rank": rank + 1, + } + # 纯向量模式：把仅在向量里命中的也纳入 + if mode == "vector": + for vc, vs in vector_boost.items(): + if vc[:300] not in merged_map: + # 找到对应 chunk + for c in corpus: + if vc in c["text"]: + merged_map[vc[:300]] = {"chunk": c["chunk"], "bm25": 0.0, "vec": vs, "rank": 999} + break + + if not merged_map: + return [] + + # 综合分数（BM25 主导，向量做 0~1 归一 boost） + max_bm25 = max((m["bm25"] for m in merged_map.values()), default=1.0) or 1.0 + items = [] + for key, m in merged_map.items(): + c = m["chunk"] + norm_bm25 = (m["bm25"] / max_bm25) if max_bm25 > 0 else 0.0 + norm_vec = min(m["vec"], 1.0) + combined = norm_bm25 * 0.7 + norm_vec * 0.3 + title = _chunk_title(c) + breadcrumb = _breadcrumb(db, c.knowledge_base_id, c.document_id and _doc_dir(db, c.document_id), c.document_id, c.chunk_index, title) + snippet = _make_snippet(c.content or "", query) + items.append({ + "content": c.content or "", + "score": round(combined, 4), + "source": "tree_chunk", + "metadata": { + "breadcrumb": breadcrumb, + "snippet": snippet, + "chunk_index": c.chunk_index, + "title": title, + "document_id": c.document_id, + "knowledge_base_id": c.knowledge_base_id, + }, + }) + items.sort(key=lambda x: x["score"], reverse=True) + return items[:top_k] + + +def _doc_dir(db, doc_id: str): + from app.models import KnowledgeDocument + d = db.query(KnowledgeDocument).filter(KnowledgeDocument.id == doc_id).first() + return d.directory_id if d else None + + +def _make_snippet(content: str, query: str, win: int = 80) -> str: + if not content: + return "" + q = query.strip() + idx = content.find(q) + if idx < 0: + # 尝试首个 token + toks = _tokenize(q) + for t in toks: + idx = content.find(t) + if idx >= 0: + break + if idx < 0: + return content[:win * 2] + start = max(0, idx - win) + end = min(len(content), idx + len(q) + win) + snip = content[start:end] + if start > 0: + snip = "..." + snip + if end < len(content): + snip = snip + "..." + return snip