참고소스 수정본

참고/knowledge_agent-main/sub_agents/analyst.py

@@ -0,0 +1,78 @@
+# sub_agents/analyst.py
+from langchain.agents import create_openai_tools_agent, AgentExecutor
+from langchain_core.prompts import ChatPromptTemplate
+from state import AgentState
+from db_utils import save_analyst_report, extract_and_clean_json
+from terminal_utils import print_colorful_break
+
+async def analyst_agent_node(state: AgentState):
+    """Runs the analyst agent and returns its raw output and the new report ID."""
+    print_colorful_break("ANALYST")
+    logger = state['logger']
+    timestamp = state['timestamp']
+    report_id = f"ana_{timestamp.replace('-', '').replace(':', '').replace('T', '_').split('.')[0]}"
+    status = f"Initialized analyst report with ID: {report_id}"
+    logger.info(status)
+
+    with open("prompts/analyst_prompt.txt", "r") as f:
+        analyst_prompt_template = f.read()
+        
+    analyst_prompt = ChatPromptTemplate.from_template(analyst_prompt_template)    
+    analyst_tools = [t for t in state['mcp_tools'] if t.name in ["query", "graphs_get", "graph_labels", "google_search", "fetch"]]
+
+    status = f"Attempting to invoke analyst agent executor with tools: {analyst_tools}"
+    logger.info(status)
+    try:
+        agent_runnable = create_openai_tools_agent(state['model'], analyst_tools, analyst_prompt)
+        executor = AgentExecutor(agent=agent_runnable, tools=analyst_tools, verbose=True)
+    except Exception as e:
+        status = f"Failed to create agent executor: {e}"
+        logger.error(status, exc_info=True)
+        return {"status": status}
+
+
+    task = state['messages'][0].content
+    status = f"Attempting to run agent executor with input: {task}"
+    logger.info(status)
+    try:
+        analyst_result = await executor.ainvoke({
+            "input": task,
+            "analyst_report_id": report_id
+        })
+        raw_report = analyst_result.get('output', '')
+        status = f"Analyst agent completed.\nRaw output: {raw_report}"
+        logger.info(status)
+
+    except Exception as e:
+        status = f"Analyst agent failed: {e}"
+        logger.error(status, exc_info=True)
+        raw_report = f"Error in Analyst Agent: {e}"
+
+    return {
+        "analyst_report_id": report_id,
+        "analyst_report": raw_report,
+        "status": status
+    }
+
+def save_analyst_report_node(state: AgentState):
+    """Saves the final report and updates the main status field."""
+    logger = state['logger']
+    raw_report_content = state.get("analyst_report")
+
+    report_id = state.get("analyst_report_id")
+    status = f"--- Saving Analyst Report: {report_id} ---"
+    logger.info(status)
+
+    try:
+        report_json = extract_and_clean_json(raw_report_content)
+        if 'report_id' not in report_json:
+            report_json['report_id'] = report_id
+        
+        save_analyst_report(report_json)
+        status = f"Successfully saved analyst report with ID {report_json.get('report_id')}"
+        logger.info(status)
+    except (ValueError, KeyError) as e:
+        status = f"Error processing or saving analyst report: {e}"
+        logger.error(status)
+    
+    return {"status": status}