ailine/backend/app/tools/__init__.py

"""
Agent Tools - 所有工具统一定义
"""

from langchain_core.tools import tool
from backend.app.logger import info

# ========== RAG ==========

_rag_pipeline = None


def _get_rag_pipeline():
    global _rag_pipeline
    if _rag_pipeline is None:
        from backend.app.rag.pipeline import RAGPipeline
        _rag_pipeline = RAGPipeline(
            num_queries=3,
            rerank_top_n=5,
            use_rerank=True,
            return_parent_docs=True,
        )
    return _rag_pipeline


@tool
async def rag_search(query: str) -> str:
    """
    检索知识库获取相关信息

    Returns:
        包含检索结果和置信度的结构化回复，格式：
        - 内容：检索到的相关信息
        - 置信度评估：基于向量相似度、重排分数、LLM判断的综合评分
    """
    info(f"[Tool] rag_search: {query[:30]}...")
    try:
        pipeline = _get_rag_pipeline()
        # 使用带置信度的检索
        result = await pipeline.aretrieve_with_confidence(query, original_query=query)

        if not result.content:
            return "【RAG检索结果】\n未在知识库中找到相关内容。\n置信度：0.0\n建议：可尝试联网搜索获取信息。"

        # 构建包含置信度的回复
        confidence_desc = "高"
        if result.confidence < 0.4:
            confidence_desc = "低"
        elif result.confidence < 0.6:
            confidence_desc = "中"

        response = f"""【RAG检索结果】
{result.content}

【置信度评估】
- 综合置信度：{result.confidence:.2f}（{confidence_desc}）
- 向量相似度：{result.scores['embedding']:.2f}
- 重排分数：{result.scores['rerank']:.2f}
- LLM评估：{result.scores['llm']:.2f}

{'✅ 检索结果可信，可直接使用' if result.is_useful else '⚠️ 检索结果置信度较低，可能需要联网搜索补充'}"""

        info(f"[Tool] rag_search 完成: confidence={result.confidence:.3f}, is_useful={result.is_useful}")
        return response

    except Exception as e:
        info(f"[Tool] rag_search 失败: {e}")
        return f"【RAG检索失败】\n错误：{str(e)}\n建议：请稍后重试或使用联网搜索"


# ========== 联网搜索 ==========

@tool
def web_search(query: str) -> str:
    """联网搜索获取最新信息"""
    info(f"[Tool] web_search: {query[:30]}...")
    try:
        from backend.app.core.web_search import web_search as search_fn
        return search_fn(query, max_results=5)
    except Exception as e:
        info(f"[Tool] web_search 失败: {e}")
        return f"联网搜索失败: {str(e)}"


# ========== 子图工具 ==========

async def _call_subgraph(builder_fn, state_cls, query: str) -> str:
    """通用子图调用"""
    try:
        graph = builder_fn().compile()
        state = state_cls(user_query=query)
        result = await graph.ainvoke(state)
        return result.get("final_result", "执行完成")
    except Exception as e:
        info(f"[Tool] 子图调用失败: {e}")
        return f"执行失败: {str(e)}"


@tool
async def contact_lookup(query: str) -> str:
    """查询通讯录"""
    from backend.app.subgraphs.contact.graph import build_contact_subgraph
    from backend.app.subgraphs.contact.state import ContactState
    return await _call_subgraph(build_contact_subgraph, ContactState, query)


@tool
async def dictionary_lookup(word: str) -> str:
    """查询词典/翻译"""
    from backend.app.subgraphs.dictionary.graph import build_dictionary_subgraph
    from backend.app.subgraphs.dictionary.state import DictionaryState
    return await _call_subgraph(build_dictionary_subgraph, DictionaryState, word)


@tool
async def news_analysis(topic: str) -> str:
    """分析新闻热点"""
    from backend.app.subgraphs.news_analysis.graph import build_news_analysis_subgraph
    from backend.app.subgraphs.news_analysis.state import NewsAnalysisState
    return await _call_subgraph(build_news_analysis_subgraph, NewsAnalysisState, topic)


# ========== 导出 ==========

ALL_TOOLS = [rag_search, web_search, contact_lookup, dictionary_lookup, news_analysis]