61 lines
2.7 KiB
Python
61 lines
2.7 KiB
Python
|
# sub_agents.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 tools import human_approval
|
||
|
|
from db_utils import load_latest_report, extract_and_clean_json
|
||
|
|
from terminal_utils import print_colorful_break
|
||
|
|
|
||
|
|
|
||
|
|
async def fixer_agent_node(state: AgentState):
|
||
|
|
print_colorful_break("FIXER")
|
||
|
|
logger = state['logger']
|
||
|
|
logger.info("--- Running Fixer Agent ---")
|
||
|
|
all_tools = state['mcp_tools']
|
||
|
|
model = state['model']
|
||
|
|
timestamp = state['timestamp']
|
||
|
|
|
||
|
|
fixer_tools = [t for t in all_tools if t.name in ["graph_update_entity", "documents_delete_entity", "graph_update_relation", "documents_delete_relation", "graph_entity_exists"]] + [load_latest_report, human_approval]
|
||
|
|
fixer_prompt = '''Your goal is to correct data quality issues.
|
||
|
|
1. **Load Auditor's Report**: Load `auditor_report.json`.
|
||
|
|
2. **Create a Plan**: Create a step-by-step plan to correct the issues.
|
||
|
|
3. **Get Human Approval**: Use `human_approval` to get your plan approved by calling the tool with the plan.
|
||
|
|
4. **Execute**: Execute the approved plan.
|
||
|
|
5. **Save Report**: Save a report of your actions to `fixer_report.json`.'''
|
||
|
|
|
||
|
|
prompt = ChatPromptTemplate.from_template(fixer_prompt)
|
||
|
|
agent_executor = create_openai_tools_agent(model, fixer_tools, prompt)
|
||
|
|
|
||
|
|
task_input = "Your task is to fix issues from the auditor's report. Begin now."
|
||
|
|
|
||
|
|
result = await agent_executor.ainvoke({"input": task_input, "timestamp": timestamp})
|
||
|
|
|
||
|
|
logger.info(f"Fixer Agent finished with output: {result['output']}")
|
||
|
|
return {"messages": state['messages'] + [AIMessage(content=result['output'])]}
|
||
|
|
|
||
|
|
def save_fixer_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 Fixer 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('fixer_report_id', 'unknown_id')
|
||
|
|
save_fixer_report({"fixer_report": json.dumps(report_json)})
|
||
|
|
status = f"Successfully saved fixer report with ID {report_json.get('report_id')}"
|
||
|
|
logger.info(status)
|
||
|
|
except (ValueError, KeyError) as e:
|
||
|
|
status = f"Error processing or saving fixer report: {e}"
|
||
|
|
logger.error(status, exc_info=True)
|
||
|
|
|
||
|
|
return {"messages": state['messages'] + [AIMessage(content=status)]}
|