triage/agent_system/controller/agent.py

from typing import Dict, Any, List
from agent_system.base import BaseAgent
from agent_system.controller.prompt import ControllerPrompt
from agent_system.controller.response_model import ControllerDecision, TaskInfo


class TaskController(BaseAgent):
    """
    任务控制器智能体
    
    负责根据患者的临床信息（现病史、既往史、主诉）从未完成的任务列表中
    选择最合适的下一步任务，并提供具体的执行指导建议。
    
    核心功能:
    1. 分析患者的临床信息和病情特征
    2. 评估待执行任务的优先级和重要性
    3. 选择最适合当前情况的任务
    4. 提供针对性的任务执行指导
    
    Attributes:
        model_type (str): 使用的大语言模型类型，默认为 gpt-oss:latest
        llm_config (dict): LLM模型配置参数
    """
    
    def __init__(self, model_type: str = "gpt-oss:latest", llm_config: dict = None):
        """
        初始化任务控制器智能体
        
        Args:
            model_type (str): 大语言模型类型，默认使用 gpt-oss:latest
            llm_config (dict): LLM模型的配置参数，如果为None则使用默认配置
        """
        super().__init__(
            model_type=model_type,
            description="医疗任务控制器，负责任务选择和执行指导",
            instructions=ControllerPrompt.instructions,
            response_model=ControllerDecision,
            llm_config=llm_config or {},
            structured_outputs=True,
            markdown=False,
            use_cache=False
        )
    
    def run(self, 
            pending_tasks: List[Dict[str, str]], 
            chief_complaint: str, 
            hpi_content: str = "", 
            ph_content: str = "") -> ControllerDecision:
        """
        执行任务控制决策
        
        基于患者的临床信息和待执行的任务列表，选择最合适的任务
        并提供具体的执行指导建议。
        
        Args:
            pending_tasks (List[Dict[str, str]]): 待执行的任务列表，每个任务包含name、priority、description字段
            chief_complaint (str): 患者主诉
            hpi_content (str, optional): 现病史内容，默认为空字符串
            ph_content (str, optional): 既往史内容，默认为空字符串
            
        Returns:
            ControllerDecision: 包含任务选择决策和指导建议的结构化数据，包括：
                - selected_task: 选择的任务信息
                - specific_guidance: 针对选定任务的具体指导建议
                
        Raises:
            Exception: 当LLM调用失败时，返回包含默认信息的ControllerDecision
        """
        try:
            # 构建决策提示词
            prompt = self._build_decision_prompt(
                pending_tasks, chief_complaint, hpi_content, ph_content
            )
            
            # 调用基类的run方法执行LLM推理
            result = super().run(prompt)
            
            # 确保返回正确的类型并进行类型转换
            return self._ensure_result_type(result)
            
        except Exception as e:
            # 当决策失败时记录错误并返回默认结果
            print(f"任务控制决策失败: {str(e)}")
            return self._get_fallback_result(pending_tasks)
    
    def _ensure_result_type(self, result: Any) -> ControllerDecision:
        """
        确保返回结果为正确的类型
        
        Args:
            result (Any): LLM返回的原始结果
            
        Returns:
            ControllerDecision: 转换后的结构化结果
        """
        if isinstance(result, ControllerDecision):
            return result
        elif isinstance(result, dict):
            return ControllerDecision(**result)
        else:
            # 如果类型不匹配，返回默认结果
            return self._get_fallback_result([])
    
    def _get_fallback_result(self, pending_tasks: List[Dict[str, str]]) -> ControllerDecision:
        """
        生成决策失败时的默认结果
        
        Args:
            pending_tasks (List[Dict[str, str]]): 待执行的任务列表
            
        Returns:
            ControllerDecision: 包含默认任务选择的结果
        """
        # 如果有待执行任务，选择第一个作为默认任务
        if pending_tasks:
            default_task = pending_tasks[0]
            selected_task_info = TaskInfo(
                task_name=default_task.get("name", "未知任务"),
                priority=default_task.get("priority", "中"),
                description=default_task.get("description", "任务描述不可用")
            )
        else:
            selected_task_info = TaskInfo(
                task_name="基本信息收集",
                priority="中",
                description="收集患者的基本临床信息"
            )
        
        return ControllerDecision(
            selected_task=selected_task_info,
            specific_guidance="由于系统异常，建议按照标准临床流程进行患者评估，重点关注患者的主要症状和病史信息，并人工审核患者情况。"
        )
    
    def _build_decision_prompt(self, 
                              pending_tasks: List[Dict[str, str]], 
                              chief_complaint: str, 
                              hpi_content: str, 
                              ph_content: str) -> str:
        """
        构建任务控制决策的提示词模板
        
        根据待执行任务列表和患者临床信息，构建简洁高效的决策提示词，
        引导LLM进行专业的任务选择和指导建议。
        
        Args:
            pending_tasks (List[Dict[str, str]]): 待执行的任务列表
            chief_complaint (str): 患者主诉
            hpi_content (str): 现病史内容
            ph_content (str): 既往史内容
            
        Returns:
            str: 精简的决策提示词
        """
        # 格式化待执行任务列表
        tasks_display = ""
        for i, task in enumerate(pending_tasks, 1):
            task_name = task.get("name", "未知任务")
            task_priority = task.get("priority", "未设定")
            task_desc = task.get("description", "无描述")
            tasks_display += f"{i}. 任务名称: {task_name}\n   优先级: {task_priority}\n   描述: {task_desc}\n\n"
        
        if not tasks_display.strip():
            tasks_display = "当前没有待执行的任务。"
        
        # 确保临床信息的合理显示
        hpi_display = hpi_content.strip() if hpi_content.strip() else "暂无现病史信息"
        ph_display = ph_content.strip() if ph_content.strip() else "暂无既往史信息"
        
        # 从prompt类获取示例输出格式
        example_output = ControllerPrompt.get_example_output()
        
        prompt = f"""患者临床信息：
主诉: {chief_complaint}
现病史: {hpi_display}
既往史: {ph_display}

待执行任务列表：
{tasks_display}

请根据患者的临床信息分析病情特征，从上述任务列表中选择最合适的下一步任务，并提供具体的执行指导建议。

输出格式示例：
{example_output}

请严格按照上述JSON格式输出。
输出内容为:"""
        
        return prompt
    
    def select_optimal_task(self, 
                           tasks: List[Dict[str, str]], 
                           patient_info: Dict[str, str]) -> ControllerDecision:
        """
        基于患者信息选择最优任务的便捷接口
        
        这是一个专门用于任务选择的简化接口，接受结构化的患者信息。
        
        Args:
            tasks (List[Dict[str, str]]): 待执行的任务列表
            patient_info (Dict[str, str]): 患者信息字典，包含chief_complaint、hpi、ph等字段
            
        Returns:
            ControllerDecision: 任务选择决策结果
        """
        chief_complaint = patient_info.get("chief_complaint", "")
        hpi_content = patient_info.get("hpi", "")
        ph_content = patient_info.get("ph", "")
        
        return self.run(
            pending_tasks=tasks,
            chief_complaint=chief_complaint,
            hpi_content=hpi_content,
            ph_content=ph_content
        )
    
    def get_task_guidance(self, result: ControllerDecision) -> Dict[str, Any]:
        """
        获取任务执行指导的结构化信息
        
        Args:
            result (ControllerDecision): 控制器决策结果
            
        Returns:
            Dict[str, Any]: 包含任务指导信息的字典
        """
        return {
            "task_name": result.selected_task.task_name,
            "priority": result.selected_task.priority,
            "guidance": result.specific_guidance
        }