""" 知识检索服务 — 从知识库中语义检索相关经验,注入到 Agent 执行上下文 """ from __future__ import annotations import logging from typing import Any, Dict, List, Optional from sqlalchemy.orm import Session from sqlalchemy import desc, func from app.core.database import SessionLocal logger = logging.getLogger(__name__) class KnowledgeRetriever: """知识检索器 — 执行前从知识库检索相关经验""" def __init__(self, top_k: int = 3, min_score: float = 0.3): self.top_k = top_k self.min_score = min_score def retrieve(self, query: str, category: Optional[str] = None) -> List[Dict[str, Any]]: """基于关键词+标签检索相关知识条目(后续可升级为 embedding 语义检索)。""" from app.models.knowledge_entry import KnowledgeEntry db: Optional[Session] = None try: db = SessionLocal() q = db.query(KnowledgeEntry).filter( KnowledgeEntry.is_active == True, ) if category: q = q.filter(KnowledgeEntry.category == category) # 关键词匹配:在 title / situation / solution / tags 中搜索 keywords = [w for w in query[:100].split() if len(w) >= 2] if keywords: conditions = [] for kw in keywords[:10]: like_pat = f"%{kw}%" conditions.append(KnowledgeEntry.title.like(like_pat)) conditions.append(KnowledgeEntry.situation.like(like_pat)) conditions.append(KnowledgeEntry.solution.like(like_pat)) from sqlalchemy import or_ q = q.filter(or_(*conditions)) entries = ( q.order_by(desc(KnowledgeEntry.confidence), desc(KnowledgeEntry.retrieval_count)) .limit(self.top_k) .all() ) results = [] for e in entries: # 更新检索计数 e.retrieval_count = (e.retrieval_count or 0) + 1 results.append({ "id": str(e.id), "title": e.title, "category": e.category, "tags": e.tags or [], "situation": e.situation, "solution": e.solution, "caveats": e.caveats, "confidence": e.confidence, }) db.commit() return results except Exception as e: logger.warning("知识检索失败: %s", e) if db: try: db.rollback() except Exception: pass return [] finally: if db: try: db.close() except Exception: pass def format_for_prompt(self, entries: List[Dict[str, Any]]) -> str: """将检索到的知识条目格式化为 Prompt 注入文本。""" if not entries: return "" parts = ["\n## 相关知识库经验(请参考以下经验来更好地完成任务)\n"] for i, e in enumerate(entries, 1): parts.append(f"### {i}. {e['title']} (置信度: {e['confidence']:.0%})") if e.get("situation"): parts.append(f"**适用场景**: {e['situation']}") if e.get("solution"): parts.append(f"**方案**: {e['solution']}") if e.get("caveats"): parts.append(f"**注意**: {e['caveats']}") parts.append("") return "\n".join(parts) def inject_knowledge(self, system_prompt: str, user_input: str) -> str: """检索相关知识并注入到 system prompt 中。""" entries = self.retrieve(user_input) knowledge_text = self.format_for_prompt(entries) if knowledge_text: return system_prompt + knowledge_text return system_prompt knowledge_retriever = KnowledgeRetriever()