wx-agent/tools/function_tool.py

"""
工具函数定义
包含RAG、GraphRAG、VLM、ASR、TTS等工具的function call封装
"""

import os
import json
import re
import httpx
import asyncio
from typing import Optional, Dict, Any, List
from langchain_core.tools import tool
from dotenv import load_dotenv

from config import RAG_CONFIG, ASR_CONFIG
from modelsAPI.model_api import OpenaiAPI
from utils.function_tracker import track_function_calls


load_dotenv()


from .graph_tools import (
    graph_rag_search,
    fault_graph_rag_search,
    operation_graph_rag_search,
)

from .rag_tools import (
    rag_search_tool,
    file_to_text_tool,
)

from .vlm_tools import (
    vlm_image_to_text,
    image_rag_describe,
)

from .fault_statistics import fault_type_statistics_tool, fault_frequency_statistics_tool, spare_parts_statistics_tool


@tool
def detect_input_type(raw_input: Dict[str, Any]) -> Dict[str, Any]:
    """
    判断输入类型工具（异步版本）：严格对齐 MainAgentState 字段。
    识别：image, audio/asr, file, query (text)
    """

    if not isinstance(raw_input, dict):
        return {
            "has_image": False, "has_audio": False, "has_file": False, "has_query": False,
            "input_type": "unknown", "error": "raw_input 必须是 dict"
        }

    has_image = False
    img_v = raw_input.get("image")
    if img_v:
        has_image = True

    has_audio = False
    aud_v = raw_input.get("audio")
    if aud_v:
        has_audio = True

    has_file = False
    file_v = raw_input.get("file_text")
    if file_v:
        has_file = True

    has_query = False
    query_v = raw_input.get("query")
    if isinstance(query_v, str) and query_v.strip():
        has_query = True

    return {
        "has_image": has_image,
        "has_audio": has_audio,
        "has_file": has_file,
        "has_query": has_query,
    }


async def asr_audio_to_text_async(
    audio_path: str,
    model: Optional[str] = None,
    language: Optional[str] = None
) -> Dict[str, Any]:
    """
    ASR语音转文本工具（异步实现）
    将音频文件转换为文本内容

    Args:
        audio_path: 音频文件路径（如：audio/test.wav）
        model: ASR模型名称（可选，默认从配置读取，如"base"）
        language: 语言代码（可选，默认从配置读取，如"zh"表示中文）

    Returns:
        包含success和text字段的字典
    """
    if not os.path.exists(audio_path):
        return {
            "success": False,
            "error": f"音频文件不存在: {audio_path}",
            "text": ""
        }

    api_base = ASR_CONFIG["api_base"]
    asr_model = model or ASR_CONFIG["model"]
    asr_language = language or ASR_CONFIG["language"]
    timeout = ASR_CONFIG["timeout"]

    api_endpoint = f"{api_base}/audio/transcriptions"

    try:
        with open(audio_path, "rb") as f:
            audio_data = f.read()

        files = {
            "file": (os.path.basename(audio_path), audio_data, "audio/wav")
        }
        data = {
            "model": asr_model,
            "language": asr_language,
        }

        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(api_endpoint, files=files, data=data)

        if response.status_code == 200:
            result = response.json()
            text = result.get("text", "")
            return {
                "success": True,
                "text": text if text else "",
                "error": None
            }
        else:
            return {
                "success": False,
                "error": f"HTTP {response.status_code}: {response.text}",
                "text": ""
            }

    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "text": ""
        }


@tool
async def asr_audio_to_text(
    audio_path: str,
    model: Optional[str] = None,
    language: Optional[str] = None
) -> str:
    """
    ASR语音转文本工具（异步版本）
    将音频文件转换为文本内容

    Args:
        audio_path: 音频文件路径（如：audio/test.wav）
        model: ASR模型名称（可选，默认从配置读取，如"base"）
        language: 语言代码（可选，默认从配置读取，如"zh"表示中文）

    Returns:
        识别出的文本内容
    """
    result = await asr_audio_to_text_async(audio_path, model, language)

    if not result.get("success", False):
        return f"语音识别失败: {result.get('error', '未知错误')}"

    return result.get("text", "")


@tool
async def extract_device_and_fault(query: str) -> Dict[str, Any]:
    """
    从用户查询中提取设备名称和故障现象（使用大模型）

    Args:
        query: 用户查询文本

    Returns:
        包含device和fault字段的字典
    """
    if not query or not query.strip():
        return {"device": "", "fault": ""}

    prompt = f"""从以下用户查询中提取设备名称和故障现象：
查询：{query}

请以JSON格式返回：
{{"device": "设备名称", "fault": "故障现象"}}

如果无法提取，请返回空字符串。"""

    try:
        result_text = await OpenaiAPI.open_api_chat_without_thinking(
            query=prompt,
            json_output=True
        )

        result_text = result_text.strip()
        result_text = re.sub(r'^```json\s*', '', result_text, flags=re.IGNORECASE)
        result_text = re.sub(r'^```\s*', '', result_text)
        result_text = re.sub(r'```$', '', result_text)

        json_match = re.search(r'\{.*\}', result_text, re.DOTALL)
        if json_match:
            result = json.loads(json_match.group(0))
            return {
                "device": result.get("device", ""),
                "fault": result.get("fault", "")
            }

        return {"device": "", "fault": ""}

    except Exception as e:
        print(f"提取设备和故障失败: {str(e)}")
        return {"device": "", "fault": ""}


# 工具列表导出
ALL_TOOLS = [
    detect_input_type,
    graph_rag_search,
    fault_graph_rag_search,
    operation_graph_rag_search,
    rag_search_tool,
    vlm_image_to_text,
    asr_audio_to_text,
    extract_device_and_fault,
    image_rag_describe,
    file_to_text_tool,
    fault_type_statistics_tool,
    fault_frequency_statistics_tool,
    spare_parts_statistics_tool
]

# 按类别分组的工具
TOOL_CATEGORIES = {
    "input_detection": [detect_input_type],
    "rag": [graph_rag_search, rag_search_tool, fault_graph_rag_search, operation_graph_rag_search],
    "vlm": [vlm_image_to_text],
    "asr": [asr_audio_to_text],
    "extraction": [extract_device_and_fault],
    "image_rag": [image_rag_describe],
    "file_process": [file_to_text_tool],
    "statistics": [fault_type_statistics_tool, fault_frequency_statistics_tool, spare_parts_statistics_tool]
}