aiagent/backend/app/services/knowledge_retriever.py

"""
知识检索服务 — 从知识库中语义检索相关经验，注入到 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()