103 lines
4.1 KiB
Python
103 lines
4.1 KiB
Python
|
# knowledge_agent.py
|
||
|
|
|
||
|
|
import json
|
||
|
|
from langchain_mcp_adapters.client import MultiServerMCPClient
|
||
|
|
from langgraph.graph import StateGraph, END
|
||
|
|
from sub_agents.analyst import analyst_agent_node, save_analyst_report_node
|
||
|
|
from sub_agents.researcher import researcher_agent_node
|
||
|
|
from sub_agents.curator import curator_agent_node
|
||
|
|
from sub_agents.auditor import auditor_agent_node, save_auditor_report_node
|
||
|
|
from sub_agents.fixer import fixer_agent_node, save_fixer_report_node
|
||
|
|
from sub_agents.advisor import advisor_agent_node, save_advisor_report_node
|
||
|
|
|
||
|
|
from state import AgentState
|
||
|
|
|
||
|
|
async def get_mcp_tools():
|
||
|
|
"""Initializes the MCP client and fetches the available tools."""
|
||
|
|
with open('mcp.json', 'r') as f:
|
||
|
|
mcp_server_config = json.load(f)
|
||
|
|
|
||
|
|
mcp_client = MultiServerMCPClient(mcp_server_config)
|
||
|
|
tools = await mcp_client.get_tools()
|
||
|
|
print(f"Successfully loaded {len(tools)} tools from MCP server.")
|
||
|
|
return tools
|
||
|
|
|
||
|
|
def create_knowledge_agent_graph(task: str, all_tools: list):
|
||
|
|
"""Creates the Knowledge Agent as a LangGraph StateGraph."""
|
||
|
|
|
||
|
|
workflow = StateGraph(AgentState)
|
||
|
|
|
||
|
|
# Define the workflow based on the task
|
||
|
|
|
||
|
|
if task == "maintenance":
|
||
|
|
# Full workflow with loops for each agent
|
||
|
|
workflow.add_node("analyst", analyst_agent_node)
|
||
|
|
workflow.add_node("save_analyst_report", save_analyst_report_node)
|
||
|
|
workflow.add_node("researcher", researcher_agent_node)
|
||
|
|
workflow.add_node("curator", curator_agent_node)
|
||
|
|
workflow.add_node("auditor", auditor_agent_node)
|
||
|
|
workflow.add_node("save_auditor_report", save_auditor_report_node)
|
||
|
|
workflow.add_node("fixer", fixer_agent_node)
|
||
|
|
workflow.add_node("save_fixer_report", save_fixer_report_node)
|
||
|
|
workflow.add_node("advisor", advisor_agent_node)
|
||
|
|
workflow.add_node("save_advisor_report", save_advisor_report_node)
|
||
|
|
|
||
|
|
workflow.set_entry_point("analyst")
|
||
|
|
workflow.add_edge("analyst", "save_analyst_report")
|
||
|
|
workflow.add_edge("save_analyst_report", "researcher")
|
||
|
|
workflow.add_edge("researcher", "curator")
|
||
|
|
workflow.add_edge("curator", "auditor")
|
||
|
|
workflow.add_edge("auditor", "save_auditor_report")
|
||
|
|
workflow.add_edge("save_auditor_report", "fixer")
|
||
|
|
workflow.add_edge("fixer", "save_fixer_report")
|
||
|
|
workflow.add_edge("save_fixer_report", "advisor")
|
||
|
|
workflow.add_edge("advisor", "save_advisor_report")
|
||
|
|
workflow.add_edge("save_advisor_report", END)
|
||
|
|
|
||
|
|
elif task == "analyze":
|
||
|
|
workflow.add_node("analyst", analyst_agent_node)
|
||
|
|
workflow.add_node("save_analyst_report", save_analyst_report_node)
|
||
|
|
|
||
|
|
workflow.set_entry_point("analyst")
|
||
|
|
workflow.add_edge("analyst", "save_analyst_report")
|
||
|
|
workflow.add_edge("save_analyst_report", END)
|
||
|
|
|
||
|
|
elif task == "research":
|
||
|
|
workflow.add_node("researcher", researcher_agent_node)
|
||
|
|
|
||
|
|
workflow.set_entry_point("researcher")
|
||
|
|
workflow.add_edge("researcher", END)
|
||
|
|
|
||
|
|
elif task == "curate":
|
||
|
|
workflow.add_node("curator", curator_agent_node)
|
||
|
|
|
||
|
|
workflow.set_entry_point("curator")
|
||
|
|
workflow.add_edge("curator", END)
|
||
|
|
|
||
|
|
elif task == "audit":
|
||
|
|
workflow.add_node("auditor", auditor_agent_node)
|
||
|
|
workflow.add_node("save_auditor_report", save_auditor_report_node)
|
||
|
|
|
||
|
|
workflow.set_entry_point("auditor")
|
||
|
|
workflow.add_edge("auditor", "save_auditor_report")
|
||
|
|
workflow.add_edge("save_auditor_report", END)
|
||
|
|
|
||
|
|
elif task == "fix":
|
||
|
|
workflow.add_node("fixer", fixer_agent_node)
|
||
|
|
workflow.add_node("save_fixer_report", save_fixer_report_node)
|
||
|
|
|
||
|
|
workflow.set_entry_point("fixer")
|
||
|
|
workflow.add_edge("fixer", "save_fixer_report")
|
||
|
|
workflow.add_edge("save_fixer_report", END)
|
||
|
|
|
||
|
|
elif task == "advise":
|
||
|
|
workflow.add_node("advisor", advisor_agent_node)
|
||
|
|
workflow.add_node("save_advisor_report", save_advisor_report_node)
|
||
|
|
|
||
|
|
workflow.set_entry_point("advisor")
|
||
|
|
workflow.add_edge("advisor", "save_advisor_report")
|
||
|
|
workflow.add_edge("save_advisor_report", END)
|
||
|
|
|
||
|
|
# Compile the graph
|
||
|
|
app = workflow.compile()
|
||
|
|
return app
|