참고소스 수정본

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

@@ -0,0 +1,56 @@
+# auditor.py
+from langchain.agents import create_openai_tools_agent, AgentExecutor
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.tools import tool, ToolException
+from langchain_core.messages import AIMessage
+import json
+import os
+import re
+from state import AgentState
+from db_utils import extract_and_clean_json
+from terminal_utils import print_colorful_break
+
+async def auditor_agent_node(state: AgentState):
+    print_colorful_break("AUDITOR")
+    logger = state['logger']
+    logger.info("--- Running Auditor Agent ---")
+    all_tools = state['mcp_tools']
+    model = state['model']
+    timestamp = state['timestamp']
+    
+    auditor_tools = [t for t in all_tools if t.name in ["graphs_get", "query"]]
+    auditor_prompt = '''Your goal is to review the LightRAG knowledge base for data quality issues.
+1.  **Identify issues**: Scan the graph for duplicates, irregular normalization, etc.
+2.  **Generate Report**: Create a report of your findings.
+3.  **Save Report**: Use `save_report` to save the findings to `auditor_report.json`.'''
+
+    prompt = ChatPromptTemplate.from_template(auditor_prompt)
+    agent_executor = create_openai_tools_agent(model, auditor_tools, prompt)
+
+    task_input = "Your task is to audit the knowledge base. Begin now."
+
+    result = await agent_executor.ainvoke({"input": task_input, "timestamp": timestamp})
+    
+    logger.info(f"Auditor Agent finished with output: {result['output']}")
+    return {"messages": state['messages'] + [AIMessage(content=result['output'])]}
+
+def save_auditor_report_node(state: AgentState):
+    """Saves the final report from the last AI message."""
+    logger = state['logger']
+    final_message_from_agent = state['messages'][-1]
+
+    status = f"--- Saving Auditor Report ---\n{final_message_from_agent.content}"
+    logger.info(status)
+
+    try:
+        report_json = extract_and_clean_json(final_message_from_agent.content)
+        if 'report_id' not in report_json:
+            report_json['report_id'] = state.get('auditor_report_id', 'unknown_id')
+        save_auditor_report({"auditor_report": json.dumps(report_json)})
+        status = f"Successfully saved auditor report with ID {report_json.get('report_id')}"
+        logger.info(status)
+    except (ValueError, KeyError) as e:
+        status = f"Error processing or saving auditor report: {e}"
+        logger.error(status, exc_info=True)
+
+    return {"messages": state['messages'] + [AIMessage(content=status)]}