| """ |
| dataRetrievalAgentGraph.py - Data Retrieval Agent Graph Builder |
| """ |
|
|
| from langgraph.graph import StateGraph, START, END |
| from src.llms.groqllm import GroqLLM |
| from src.states.dataRetrievalAgentState import DataRetrievalAgentState |
| from src.nodes.dataRetrievalAgentNode import DataRetrievalAgentNode |
|
|
|
|
| class DataRetrievalAgentGraph(DataRetrievalAgentNode): |
| def __init__(self, llm): |
| super().__init__(llm) |
| self.llm = llm |
|
|
| def prepare_worker_tasks(self, state: DataRetrievalAgentState) -> dict: |
| tasks = state.generated_tasks |
| initial_states = [{"generated_tasks": [task]} for task in tasks] |
| return {"tasks_for_workers": initial_states} |
|
|
| def create_worker_graph(self): |
| worker_graph_builder = StateGraph(DataRetrievalAgentState) |
|
|
| worker_graph_builder.add_node("worker_agent", self.worker_agent_node) |
| worker_graph_builder.add_node("tool_node", self.tool_node) |
|
|
| worker_graph_builder.set_entry_point("worker_agent") |
| worker_graph_builder.add_edge("worker_agent", "tool_node") |
| worker_graph_builder.add_edge("tool_node", END) |
|
|
| return worker_graph_builder.compile() |
|
|
| def aggregate_results(self, state: DataRetrievalAgentState) -> dict: |
| worker_outputs = getattr(state, "worker", []) |
| new_results = [] |
|
|
| if isinstance(worker_outputs, list): |
| for output in worker_outputs: |
| if "worker_results" in output and output["worker_results"]: |
| new_results.extend(output["worker_results"]) |
|
|
| return {"worker_results": new_results, "latest_worker_results": new_results} |
|
|
| def format_output(self, state: DataRetrievalAgentState) -> dict: |
| classified_events = state.classified_buffer |
| insights = [] |
|
|
| for event in classified_events: |
| insights.append( |
| { |
| "source_event_id": event.event_id, |
| "domain": event.target_agent, |
| "severity": "medium", |
| "summary": event.content_summary, |
| "risk_score": event.confidence_score, |
| } |
| ) |
|
|
| print(f"[DATA RETRIEVAL] Formatted {len(insights)} insights for parent graph") |
| return {"domain_insights": insights} |
|
|
| def build_data_retrieval_agent_graph(self): |
| worker_graph = self.create_worker_graph() |
| workflow = StateGraph(DataRetrievalAgentState) |
|
|
| workflow.add_node("master_delegator", self.master_agent_node) |
| workflow.add_node("prepare_worker_tasks", self.prepare_worker_tasks) |
| workflow.add_node( |
| "worker", |
| lambda state: { |
| "worker": worker_graph.map().invoke(state.tasks_for_workers) |
| }, |
| ) |
| workflow.add_node("aggregate_results", self.aggregate_results) |
| workflow.add_node("classifier_agent", self.classifier_agent_node) |
| workflow.add_node("format_output", self.format_output) |
|
|
| workflow.set_entry_point("master_delegator") |
| workflow.add_edge("master_delegator", "prepare_worker_tasks") |
| workflow.add_edge("prepare_worker_tasks", "worker") |
| workflow.add_edge("worker", "aggregate_results") |
| workflow.add_edge("aggregate_results", "classifier_agent") |
| workflow.add_edge("classifier_agent", "format_output") |
| workflow.add_edge("format_output", END) |
|
|
| return workflow.compile() |
|
|
|
|
| llm = GroqLLM().get_llm() |
| graph_builder = DataRetrievalAgentGraph(llm) |
| graph = graph_builder.build_data_retrieval_agent_graph() |
|
|
