dataflow_agent/agentroles/data_agents/operator_qa_agent.py

"""
OperatorQA Agent - 算子问答 Agent
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
生成时间: 2025-12-01 15:05:06
生成位置: dataflow_agent/agentroles/common_agents/operator_qa_agent.py

本文件实现了基于 Agentic RAG 的算子问答 Agent，支持：
1. 自然语言查询算子功能
2. 查询算子参数含义
3. 查看算子源码
4. 多轮对话

架构设计：
- OperatorRAGService: 封装所有 RAG 相关逻辑，便于后续升级为 Agentic RAG
- OperatorQAAgent: Agent 核心类，负责对话管理和 LLM 交互
"""

from __future__ import annotations

import os
from typing import Any, Dict, List, Optional, Union

from dataflow_agent.state import MainState
from dataflow_agent.toolkits.tool_manager import ToolManager
from dataflow_agent.logger import get_logger
from dataflow_agent.agentroles.cores.base_agent import BaseAgent
from dataflow_agent.agentroles.cores.registry import register
from dataflow_agent.graphbuilder.message_history import AdvancedMessageHistory

log = get_logger(__name__)


# ==============================================================================
# OperatorRAGService - RAG 服务类（解耦设计，便于后续升级为 Agentic RAG）
# ==============================================================================
class OperatorRAGService:
    """
    算子 RAG 检索服务
    
    封装所有 RAG 相关逻辑，与 Agent 解耦，便于后续升级为更复杂的 Agentic RAG。
    
    功能：
    - search(query): 向量检索相关算子
    - get_operator_info(names): 获取算子基本信息
    - get_operator_source(name): 获取算子源码
    - get_operator_params(name): 获取参数详情
    """
    
    def __init__(
        self,
        ops_json_path: Optional[str] = None,
        faiss_index_path: Optional[str] = None,
        embedding_model: str = "text-embedding-3-small",
        embedding_api_url: str = "http://123.129.219.111:3000/v1/embeddings",
        api_key: Optional[str] = None,
    ):
        """
        初始化 RAG 服务
        
        Args:
            ops_json_path: 算子 JSON 文件路径，默认使用项目内置路径
            faiss_index_path: FAISS 索引文件路径，默认使用项目内置路径
            embedding_model: embedding 模型名称
            embedding_api_url: embedding API 地址
            api_key: API 密钥
        """
        from dataflow_agent.utils import get_project_root
        
        project_root = get_project_root()
        self.ops_json_path = ops_json_path or str(
            project_root / "dataflow_agent/toolkits/resources/ops.json"
        )
        self.faiss_index_path = faiss_index_path or str(
            project_root / "dataflow_agent/resources/faiss_cache/all_ops.index"
        )
        self.embedding_model = embedding_model
        self.embedding_api_url = embedding_api_url
        self.api_key = api_key or os.getenv("DF_API_KEY")
        
        self._searcher = None  # 懒加载
        
    def _get_searcher(self):
        """获取或创建 RAG 检索器（懒加载）"""
        if self._searcher is None:
            from dataflow_agent.toolkits.optool.op_tools import RAGOperatorSearch
            self._searcher = RAGOperatorSearch(
                ops_json_path=self.ops_json_path,
                faiss_index_path=self.faiss_index_path,
                model_name=self.embedding_model,
                base_url=self.embedding_api_url,
                api_key=self.api_key,
            )
        return self._searcher
    
    def search(
        self,
        query: Union[str, List[str]],
        top_k: int = 5,
    ) -> Union[List[str], List[List[str]]]:
        """
        向量检索相关算子
        
        Args:
            query: 查询字符串或查询列表
            top_k: 返回 top-k 个结果
            
        Returns:
            算子名称列表
        """
        try:
            searcher = self._get_searcher()
            return searcher.search(query, top_k=top_k)
        except Exception as e:
            log.error(f"RAG 检索失败: {e}")
            return [] if isinstance(query, str) else [[] for _ in query]
    
    def get_operator_info(self, names: List[str]) -> str:
        """
        获取算子基本信息（名称、描述、分类）
        
        Args:
            names: 算子名称列表
            
        Returns:
            JSON 格式的算子信息字符串
        """
        from dataflow_agent.toolkits.optool.op_tools import get_operators_info_by_names
        try:
            return get_operators_info_by_names(names)
        except Exception as e:
            log.error(f"获取算子信息失败: {e}")
            return "[]"
    
    def get_operator_source(self, name: str) -> str:
        """
        获取算子源码
        
        Args:
            name: 算子名称
            
        Returns:
            源码字符串
        """
        from dataflow_agent.toolkits.optool.op_tools import get_operator_source_by_name
        try:
            return get_operator_source_by_name(name)
        except Exception as e:
            log.error(f"获取算子源码失败: {e}")
            return f"# 无法获取算子 '{name}' 的源码: {e}"
    
    def get_operator_params(self, name: str) -> Dict[str, Any]:
        """
        获取算子参数详情（init 和 run 参数）
        
        Args:
            name: 算子名称
            
        Returns:
            包含 init 和 run 参数的字典
        """
        import json
        try:
            info_str = self.get_operator_info([name])
            info_list = json.loads(info_str)
            if info_list and len(info_list) > 0:
                op_info = info_list[0]
                return op_info.get("parameter", {"init": [], "run": []})
            return {"init": [], "run": []}
        except Exception as e:
            log.error(f"获取算子参数失败: {e}")
            return {"init": [], "run": [], "error": str(e)}
    
    def search_and_get_info(
        self,
        query: str,
        top_k: int = 5,
    ) -> Dict[str, Any]:
        """
        综合检索：先向量检索，再获取详细信息
        
        Args:
            query: 查询字符串
            top_k: 返回 top-k 个结果
            
        Returns:
            包含检索结果和详细信息的字典
        """
        import json
        
        # 1. 向量检索
        matched_names = self.search(query, top_k=top_k)
        
        # 2. 获取详细信息
        if matched_names:
            info_str = self.get_operator_info(matched_names)
            try:
                info_list = json.loads(info_str)
            except:
                info_list = []
        else:
            info_list = []
        
        return {
            "query": query,
            "matched_operators": matched_names,
            "operator_details": info_list,
        }


# ==============================================================================
# OperatorQAAgent - Agent 核心类
# ==============================================================================
@register("operator_qa")
class OperatorQAAgent(BaseAgent):
    """
    算子问答 Agent
    
    支持功能：
    1. 自然语言查询算子功能（"我想过滤掉缺失值用哪个算子？"）
    2. 查询特定算子做什么（"df.filter_by 是干嘛的？"）
    3. 查询算子参数含义（"算子的 run 函数里面参数是什么意思？"）
    4. 查看算子源码
    5. 多轮对话
    """
    
    def __init__(
        self,
        tool_manager: Optional[ToolManager] = None,
        rag_service: Optional[OperatorRAGService] = None,
        **kwargs,
    ):
        """
        初始化 OperatorQAAgent
        
        Args:
            tool_manager: 工具管理器
            rag_service: RAG 服务实例，如果不传则自动创建
            **kwargs: 传递给 BaseAgent 的其他参数
        """
        # 启用多轮对话支持
        kwargs.setdefault("ignore_history", False)
        super().__init__(tool_manager=tool_manager, **kwargs)
        
        # 组合 RAG 服务
        self.rag_service = rag_service or OperatorRAGService()
    
    # ---------- 工厂方法 ----------
    @classmethod
    def create(
        cls,
        tool_manager: Optional[ToolManager] = None,
        rag_service: Optional[OperatorRAGService] = None,
        **kwargs,
    ):
        return cls(tool_manager=tool_manager, rag_service=rag_service, **kwargs)

    # ---------- 基本配置 ----------
    @property
    def role_name(self) -> str:
        return "operator_qa"

    @property
    def system_prompt_template_name(self) -> str:
        return "system_prompt_for_operator_qa"

    @property
    def task_prompt_template_name(self) -> str:
        return "task_prompt_for_operator_qa"

    # ---------- Prompt 参数 ----------
    def get_task_prompt_params(self, pre_tool_results: Dict[str, Any]) -> Dict[str, Any]:
        """
        根据前置工具结果构造 prompt 参数
        
        RAG 检索已改为后置工具，由 LLM 自主决定是否调用。
        对话历史由 BaseAgent 的 AdvancedMessageHistory 自动管理。
        
        Args:
            pre_tool_results: 前置工具执行结果
            
        Returns:
            prompt 参数字典
        """
        return {
            "user_query": pre_tool_results.get("user_query", ""),
        }

    def get_default_pre_tool_results(self) -> Dict[str, Any]:
        """默认前置工具结果"""
        return {
            "user_query": "",
        }

    # ---------- 结果写回 ----------
    def update_state_result(
        self,
        state: MainState,
        result: Dict[str, Any],
        pre_tool_results: Dict[str, Any],
    ):
        """将推理结果写回 MainState"""
        # 存储 QA 结果
        state.operator_qa_result = result
        super().update_state_result(state, result, pre_tool_results)
    
    # ---------- RAG 相关方法（供外部调用） ----------
    def search_operators(self, query: str, top_k: int = 5) -> List[str]:
        """检索相关算子"""
        return self.rag_service.search(query, top_k=top_k)
    
    def get_operator_info(self, names: List[str]) -> str:
        """获取算子信息"""
        return self.rag_service.get_operator_info(names)
    
    def get_operator_source(self, name: str) -> str:
        """获取算子源码"""
        return self.rag_service.get_operator_source(name)
    
    def get_operator_params(self, name: str) -> Dict[str, Any]:
        """获取算子参数"""
        return self.rag_service.get_operator_params(name)



# ==============================================================================
# Helper APIs
# ==============================================================================
async def operator_qa(
    state: MainState,
    model_name: Optional[str] = None,
    tool_manager: Optional[ToolManager] = None,
    rag_service: Optional[OperatorRAGService] = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    tool_mode: str = "auto",
    react_mode: bool = False,
    react_max_retries: int = 3,
    parser_type: str = "json",
    parser_config: Optional[Dict[str, Any]] = None,
    use_agent: bool = False,
    **kwargs,
) -> MainState:
    """
    OperatorQA 的异步入口
    
    Args:
        state: 主状态对象
        model_name: 模型名称
        tool_manager: 工具管理器实例
        rag_service: RAG 服务实例
        temperature: 采样温度
        max_tokens: 最大生成 token 数
        tool_mode: 工具调用模式
        react_mode: 是否启用 ReAct 推理模式
        react_max_retries: ReAct 模式下最大重试次数
        parser_type: 解析器类型
        parser_config: 解析器配置字典
        use_agent: 是否使用 agent 模式
        **kwargs: 其他传递给 execute 的参数
        
    Returns:
        更新后的 MainState 对象
    """
    agent = OperatorQAAgent(
        tool_manager=tool_manager,
        rag_service=rag_service,
        model_name=model_name,
        temperature=temperature,
        max_tokens=max_tokens,
        tool_mode=tool_mode,
        react_mode=react_mode,
        react_max_retries=react_max_retries,
        parser_type=parser_type,
        parser_config=parser_config,
    )
    return await agent.execute(state, use_agent=use_agent, **kwargs)


def create_operator_qa_agent(
    tool_manager: Optional[ToolManager] = None,
    rag_service: Optional[OperatorRAGService] = None,
    model_name: Optional[str] = None,
    temperature: float = 0.0,
    max_tokens: int = 4096,
    tool_mode: str = "auto",
    react_mode: bool = False,
    react_max_retries: int = 3,
    parser_type: str = "json",
    parser_config: Optional[Dict[str, Any]] = None,
    message_history: Optional["AdvancedMessageHistory"] = None,
    **kwargs,
) -> OperatorQAAgent:
    """
    创建 OperatorQAAgent 实例
    
    Args:
        tool_manager: 工具管理器
        rag_service: RAG 服务实例
        model_name: 模型名称
        temperature: 采样温度
        max_tokens: 最大 token 数
        tool_mode: 工具模式
        react_mode: ReAct 模式
        react_max_retries: 最大重试次数
        parser_type: 解析器类型
        parser_config: 解析器配置
        **kwargs: 其他参数
        
    Returns:
        OperatorQAAgent 实例
    """
    return OperatorQAAgent.create(
        tool_manager=tool_manager,
        rag_service=rag_service,
        model_name=model_name,
        temperature=temperature,
        max_tokens=max_tokens,
        tool_mode=tool_mode,
        react_mode=react_mode,
        react_max_retries=react_max_retries,
        parser_type=parser_type,
        parser_config=parser_config,
        message_history=message_history,
        **kwargs,
    )