79 lines
3.7 KiB
Python
79 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
from agents.control_room import AgentControlRoomService
|
|
from control_plane.agents.models import AgentVersion
|
|
from graph.native_runtime import GraphExecutionContext
|
|
from graph.registry import NodeHandlerRegistry, NodeResult
|
|
from graph.spec import ExecutionGraphSpec, GraphEdgeSpec, GraphNodeSpec
|
|
|
|
|
|
def agent_investigation_graph_v1() -> ExecutionGraphSpec:
|
|
nodes = ["prepare", "gather_agent_evidence", "smart_investigation", "improvement_candidates", "complete"]
|
|
spec = ExecutionGraphSpec(
|
|
name="agent_investigation",
|
|
version=1,
|
|
graph_type="AGENT_INVESTIGATION",
|
|
entry="prepare",
|
|
nodes={node: GraphNodeSpec(node, node if node == "complete" else f"agent_{node}") for node in nodes},
|
|
edges=[GraphEdgeSpec(nodes[index], nodes[index + 1], "success") for index in range(len(nodes) - 1)],
|
|
terminal_nodes=["complete"],
|
|
metadata={"description": "Agent investigation workflow using Control Room telemetry and Progeny Smart Investigation."},
|
|
)
|
|
spec.validate()
|
|
return spec
|
|
|
|
|
|
class AgentInvestigationNode:
|
|
idempotent = True
|
|
replay_safe = True
|
|
destructive = False
|
|
|
|
def __init__(self, service: AgentControlRoomService, node_type: str) -> None:
|
|
self.service = service
|
|
self.node_type = node_type
|
|
|
|
def version(self, context: GraphExecutionContext) -> AgentVersion:
|
|
return AgentVersion.objects.get(id=context.graph_run.metadata["agent_version_id"])
|
|
|
|
|
|
class AgentSimpleNode(AgentInvestigationNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
return NodeResult("COMPLETE", "success")
|
|
|
|
|
|
class AgentGatherEvidenceNode(AgentInvestigationNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
version = self.version(context)
|
|
evidence = {"performance": self.service.get_agent_performance(version.id), "health": self.service.get_agent_health(version.id), "usage": self.service.get_agent_usage(version.id)}
|
|
metadata = dict(context.graph_run.metadata)
|
|
metadata["agent_evidence"] = evidence
|
|
context.graph_run.metadata = metadata
|
|
context.graph_run.save(update_fields=["metadata", "updated_at"])
|
|
return NodeResult("COMPLETE", "success", evidence)
|
|
|
|
|
|
class AgentSmartInvestigationNode(AgentInvestigationNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
investigation = self.service.investigate_agent(self.version(context))
|
|
metadata = dict(context.graph_run.metadata)
|
|
metadata["progeny_investigation_id"] = str(investigation.id)
|
|
context.graph_run.metadata = metadata
|
|
context.graph_run.save(update_fields=["metadata", "updated_at"])
|
|
return NodeResult("COMPLETE", "success", {"investigation_id": str(investigation.id)})
|
|
|
|
|
|
class AgentImprovementCandidateNode(AgentInvestigationNode):
|
|
def run(self, context: GraphExecutionContext) -> NodeResult:
|
|
from control_plane.agents.models import ProgenyInvestigation
|
|
|
|
investigation_id = context.graph_run.metadata.get("progeny_investigation_id")
|
|
investigation = ProgenyInvestigation.objects.get(id=investigation_id)
|
|
candidate = self.service.create_improvement_candidate(investigation)
|
|
return NodeResult("COMPLETE", "success", {"improvement_candidate_id": str(candidate.id)})
|
|
|
|
|
|
def agent_investigation_registry(service: AgentControlRoomService) -> NodeHandlerRegistry:
|
|
registry = NodeHandlerRegistry()
|
|
for handler in [AgentSimpleNode(service, "agent_prepare"), AgentGatherEvidenceNode(service, "agent_gather_agent_evidence"), AgentSmartInvestigationNode(service, "agent_smart_investigation"), AgentImprovementCandidateNode(service, "agent_improvement_candidates")]:
|
|
registry.register(handler)
|
|
return registry
|