diff --git a/src/dynavec/agents/__init__.py b/src/dynavec/agents/__init__.py index 4e00e4c..750a278 100644 --- a/src/dynavec/agents/__init__.py +++ b/src/dynavec/agents/__init__.py @@ -24,7 +24,6 @@ ) from .graph import ( AgentGraph, - AgentState, BaseNode, Edge, FlowGraph, @@ -40,36 +39,24 @@ append_reducer, default_reducer, ) +from .handoff import Handoff, HandoffRequest from .planner import Planner from .react import ReActAgent from .registry import ToolRegistry, default_registry +from .state import AgentState +from .supervisor import AgentLike, Supervisor __all__ = [ - # Existing agent primitives "AgentResult", + "AgentLike", "AgentStep", "AgentTool", - "Plan", - "PlanStep", - "Planner", - "ReActAgent", - "tool", - # Tool registry & connectors - "ToolRegistry", - "create_http_tool", - "create_python_code_tool", - "default_registry", - "mcp_tools_from_session", - "mcp_tools_from_session_async", - # Checkpointer - "Checkpoint", - "BaseCheckpointer", - "MemoryCheckpointer", - "DynamoDBCheckpointer", - # Graph orchestration engine (#277 / #278) "AgentGraph", "AgentState", "BaseNode", + "BaseCheckpointer", + "Checkpoint", + "DynamoDBCheckpointer", "Edge", "FlowGraph", "FlowState", @@ -79,8 +66,23 @@ "GraphResult", "GraphState", "GraphValidationError", + "Handoff", + "HandoffRequest", "MaxIterationsExceededError", + "MemoryCheckpointer", "NodeTrace", + "Plan", + "PlanStep", + "Planner", + "ReActAgent", + "Supervisor", + "ToolRegistry", "append_reducer", + "create_http_tool", + "create_python_code_tool", "default_reducer", + "default_registry", + "mcp_tools_from_session", + "mcp_tools_from_session_async", + "tool", ] diff --git a/src/dynavec/agents/handoff.py b/src/dynavec/agents/handoff.py new file mode 100644 index 0000000..49c57e8 --- /dev/null +++ b/src/dynavec/agents/handoff.py @@ -0,0 +1,55 @@ +"""Handoff primitives for multi-agent workflows.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from .base import AgentResult +from .state import AgentState + +if TYPE_CHECKING: + from .supervisor import AgentLike + + +@dataclass(frozen=True) +class HandoffRequest: + """Request graph execution to continue at a named specialist.""" + + target: str + goal: str + + +class Handoff: + """Pass workflow state through or request transfer to a named specialist.""" + + def __call__( + self, + target: str | AgentState, + goal: str | None = None, + ) -> HandoffRequest | AgentState: + """Preserve legacy state passthrough and support named handoff requests.""" + if isinstance(target, AgentState): + if goal is not None: + raise TypeError("A goal cannot be supplied when passing through state.") + return target + + if goal is None: + raise TypeError("A goal is required for a named handoff.") + + if not target.strip(): + raise ValueError("Handoff target must not be empty.") + if not goal.strip(): + raise ValueError("Handoff goal must not be empty.") + + return HandoffRequest(target=target, goal=goal) + + def transfer( + self, + target: AgentLike, + goal: str, + *, + state: AgentState, + ) -> AgentResult: + """Invoke a target specialist with the existing workflow state.""" + return target.run(goal, state=state) diff --git a/src/dynavec/agents/react.py b/src/dynavec/agents/react.py index 3bd9546..d8be15c 100644 --- a/src/dynavec/agents/react.py +++ b/src/dynavec/agents/react.py @@ -2,12 +2,14 @@ from __future__ import annotations +import json from collections.abc import Callable, Sequence from typing import TYPE_CHECKING, Any from ..chat.base import ChatModel, Message, Tool, ToolCall from .base import AgentResult, AgentStep, AgentTool from .registry import ToolRegistry +from .state import AgentState if TYPE_CHECKING: from ..checkpoint import BaseCheckpointer @@ -93,17 +95,25 @@ def __init__( self._chat_tools.append(t.to_chat_tool()) elif tools: for raw_tool in tools: - agent_tool = ( - raw_tool if isinstance(raw_tool, AgentTool) else AgentTool(raw_tool) - ) + agent_tool = raw_tool if isinstance(raw_tool, AgentTool) else AgentTool(raw_tool) self._tools_map[agent_tool.name] = agent_tool self._chat_tools.append(agent_tool.to_chat_tool()) - def _init_messages(self, goal: str) -> list[Message]: + def _init_messages( + self, + goal: str, + state: AgentState | None = None, + ) -> list[Message]: messages: list[Message] = [] if self.system_prompt: messages.append(Message(role="system", content=self.system_prompt)) - messages.append(Message(role="user", content=goal)) + + user_content = goal + if state is not None and state.data: + shared_context = json.dumps(state.data, default=str, sort_keys=True) + user_content = f"{goal}\n\n" f"Shared workflow context (JSON):\n{shared_context}" + + messages.append(Message(role="user", content=user_content)) return messages def _dispatch_tool_call(self, tc: ToolCall) -> str: @@ -118,9 +128,7 @@ async def _adispatch_tool_call(self, tc: ToolCall) -> str: return await tool_instance.aexecute(tc.arguments) return f"Error: Tool {tc.name!r} is not registered in available tools." - def _process_tool_calls( - self, tool_calls: list[ToolCall], messages: list[Message] - ) -> list[str]: + def _process_tool_calls(self, tool_calls: list[ToolCall], messages: list[Message]) -> list[str]: step_observations: list[str] = [] for tc in tool_calls: obs = self._dispatch_tool_call(tc) @@ -150,12 +158,9 @@ async def _aprocess_tool_calls( ) return step_observations - def _finalize_result( - self, steps: list[AgentStep], total_tool_calls: int - ) -> AgentResult: + def _finalize_result(self, steps: list[AgentStep], total_tool_calls: int) -> AgentResult: last_output = ( - steps[-1].thought - or (steps[-1].observations[-1] if steps[-1].observations else "") + steps[-1].thought or (steps[-1].observations[-1] if steps[-1].observations else "") if steps else "" ) @@ -174,6 +179,7 @@ def run( *, thread_id: str | None = None, resume_from: str | None = None, + state: AgentState | None = None, **kwargs: Any, ) -> AgentResult: """Execute the ReAct loop synchronously until goal completion or max_steps.""" @@ -184,17 +190,27 @@ def run( if self.checkpointer is not None and thread_id is not None: existing_cp = self.checkpointer.get(thread_id, resume_from) + if existing_cp is not None: - state = existing_cp.state - if "messages" in state: - messages = [_dict_to_message(m) for m in state["messages"]] - if "steps" in state: - steps = [_dict_to_step(s) for s in state["steps"]] - total_tool_calls = int(state.get("total_tool_calls", 0)) + checkpoint_state = existing_cp.state + if checkpoint_state is not None: + if "messages" in checkpoint_state: + messages = [ + _dict_to_message(m) + for m in checkpoint_state["messages"] + ] + if "steps" in checkpoint_state: + steps = [ + _dict_to_step(s) + for s in checkpoint_state["steps"] + ] + total_tool_calls = int( + checkpoint_state.get("total_tool_calls", 0) + ) start_step = existing_cp.step + 1 if not messages: - messages = self._init_messages(goal) + messages = self._init_messages(goal, state) if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, @@ -296,6 +312,7 @@ async def arun( *, thread_id: str | None = None, resume_from: str | None = None, + state: AgentState | None = None, **kwargs: Any, ) -> AgentResult: """Execute the ReAct loop asynchronously.""" @@ -307,16 +324,25 @@ async def arun( if self.checkpointer is not None and thread_id is not None: existing_cp = self.checkpointer.get(thread_id, resume_from) if existing_cp is not None: - state = existing_cp.state - if "messages" in state: - messages = [_dict_to_message(m) for m in state["messages"]] - if "steps" in state: - steps = [_dict_to_step(s) for s in state["steps"]] - total_tool_calls = int(state.get("total_tool_calls", 0)) + checkpoint_state = existing_cp.state + if checkpoint_state is not None: + if "messages" in checkpoint_state: + messages = [ + _dict_to_message(m) + for m in checkpoint_state["messages"] + ] + if "steps" in checkpoint_state: + steps = [ + _dict_to_step(s) + for s in checkpoint_state["steps"] + ] + total_tool_calls = int( + checkpoint_state.get("total_tool_calls", 0) + ) start_step = existing_cp.step + 1 if not messages: - messages = self._init_messages(goal) + messages = self._init_messages(goal, state) if self.checkpointer is not None and thread_id is not None: self.checkpointer.put( thread_id, @@ -371,9 +397,7 @@ async def arun( ) total_tool_calls += len(msg.tool_calls) - step_observations = await self._aprocess_tool_calls( - list(msg.tool_calls), messages - ) + step_observations = await self._aprocess_tool_calls(list(msg.tool_calls), messages) step = AgentStep( step_number=step_idx, diff --git a/src/dynavec/agents/state.py b/src/dynavec/agents/state.py new file mode 100644 index 0000000..54ae44d --- /dev/null +++ b/src/dynavec/agents/state.py @@ -0,0 +1,13 @@ +"""Shared state primitives for multi-agent workflows.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class AgentState: + """Shared state passed between agents in a workflow.""" + + data: dict[str, Any] = field(default_factory=dict) diff --git a/src/dynavec/agents/supervisor.py b/src/dynavec/agents/supervisor.py new file mode 100644 index 0000000..06cb5eb --- /dev/null +++ b/src/dynavec/agents/supervisor.py @@ -0,0 +1,135 @@ +"""Supervisor primitives for multi-agent workflows.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Protocol, Union + +from .base import AgentResult +from .handoff import HandoffRequest +from .state import AgentState + + +class AgentLike(Protocol): + """Interface required by a supervisor-managed agent.""" + + def run( + self, + goal: str, + *, + state: AgentState, + ) -> AgentResult: + """Execute an agent using shared workflow state.""" + ... + + +Router = Callable[[str, AgentState], Union[str, list[str]]] + + +class Supervisor: + """Delegate a goal to specialists using shared workflow state.""" + + def __init__( + self, + agents: list[AgentLike] | dict[str, AgentLike], + *, + router: Router | None = None, + ) -> None: + self.agents: dict[str, AgentLike] | list[AgentLike] + if isinstance(agents, dict): + self.agents = dict(agents) + else: + self.agents = list(agents) + + self.router = router + + if router is not None and not isinstance(self.agents, dict): + raise TypeError("A router requires agents to be provided as a named dictionary.") + + if isinstance(self.agents, dict) and not self.agents: + raise ValueError("At least one agent must be registered.") + if isinstance(self.agents, list) and not self.agents: + raise ValueError("At least one agent must be registered.") + + def run(self, goal: str, *, state: AgentState) -> list[AgentResult]: + """Run all agents sequentially or route to one selected specialist.""" + if self.router is not None: + selected = self.router(goal, state) + agent_names = [selected] if isinstance(selected, str) else selected + + if not agent_names: + raise ValueError("Router must select at least one agent.") + + unknown = [name for name in agent_names if name not in self.agents] + if unknown: + raise ValueError( + f"Router selected unknown agent(s) {unknown!r}. " + f"Registered agents: {list(self.agents)}" + ) + + if self.router is not None: + if not isinstance(self.agents, dict): + raise TypeError( + "A router requires agents to be provided as a named dictionary." + ) + + selected = self.router(goal, state) + agent_names = [selected] if isinstance(selected, str) else selected + + if not agent_names: + raise ValueError("Router must select at least one agent.") + + unknown = [name for name in agent_names if name not in self.agents] + if unknown: + raise ValueError( + f"Router selected unknown agent(s) {unknown!r}. " + f"Registered agents: {list(self.agents)}" + ) + + return [ + self.agents[name].run(goal, state=state) + for name in agent_names + ] + + agents = self.agents.values() if isinstance(self.agents, dict) else self.agents + return [agent.run(goal, state=state) for agent in agents] + + def run_and_aggregate( + self, + goal: str, + *, + state: AgentState, + ) -> AgentResult: + """Run the selected or registered agents and combine their outputs.""" + results = self.run(goal, state=state) + combined_output = "\n\n".join(result.output for result in results if result.output.strip()) + + return AgentResult( + output=combined_output, + steps=[step for result in results for step in result.steps], + finished=all(result.finished for result in results), + termination_reason=( + "completed" if all(result.finished for result in results) else "incomplete" + ), + total_steps=sum(result.total_steps for result in results), + tool_calls_count=sum(result.tool_calls_count for result in results), + ) + + def handoff( + self, + request: HandoffRequest, + *, + state: AgentState, + ) -> AgentResult: + """Transfer execution to the named specialist with shared state.""" + if not isinstance(self.agents, dict): + raise TypeError("Handoffs require agents to be registered by name.") + + agent = self.agents.get(request.target) + if agent is None: + raise ValueError( + f"Unknown handoff target {request.target!r}. " + f"Registered agents: {list(self.agents)}" + ) + + return agent.run(request.goal, state=state) diff --git a/tests/test_agent_state.py b/tests/test_agent_state.py new file mode 100644 index 0000000..78db07a --- /dev/null +++ b/tests/test_agent_state.py @@ -0,0 +1,20 @@ +from dynavec.agents import AgentState + + +def test_agent_state_stores_shared_data(): + state = AgentState() + + state.data["customer_id"] = "C123" + state.data["decline_code"] = "51" + + assert state.data["customer_id"] == "C123" + assert state.data["decline_code"] == "51" + + +def test_agent_state_has_independent_data(): + first = AgentState() + second = AgentState() + + first.data["customer_id"] = "C123" + + assert second.data == {} diff --git a/tests/test_agents.py b/tests/test_agents.py index 4ee91cb..9ca3124 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -9,6 +9,7 @@ import pytest from dynavec.agents import ( + AgentState, AgentTool, Plan, Planner, @@ -43,9 +44,7 @@ def invoke( if self.call_count >= len(self.responses): # Default response if script exhausted return ChatResult( - message=Message( - role="assistant", content="Scripted responses exhausted." - ), + message=Message(role="assistant", content="Scripted responses exhausted."), finish_reason="stop", ) @@ -343,3 +342,67 @@ async def test_planner_async_aplan() -> None: assert len(plan.steps) == 2 assert plan.steps[0].description == "Step A" assert plan.steps[1].description == "Step B" + + +def test_react_agent_accepts_shared_state(): + state = AgentState() + + agent = ReActAgent( + model=ScriptedChatModel( + [ + "Transaction analysis completed.", + ] + ) + ) + + result = agent.run( + "Analyze transaction TX123.", + state=state, + ) + + assert result.output == "Transaction analysis completed." + + +def test_react_agent_includes_shared_state_in_messages() -> None: + """Shared workflow context should be visible to the model.""" + state = AgentState( + data={ + "transaction_id": "TX123", + "decline_code": "51", + } + ) + fake_model = ScriptedChatModel(["Transaction analysis completed."]) + agent = ReActAgent(model=fake_model) + + result = agent.run("Analyze transaction TX123.", state=state) + + assert result.output == "Transaction analysis completed." + sent_messages = fake_model.received_messages[0] + all_content = "\n".join(message.content or "" for message in sent_messages) + + assert "TX123" in all_content + assert "51" in all_content + + +@pytest.mark.asyncio +async def test_react_agent_async_includes_shared_state_in_messages() -> None: + state = AgentState( + data={ + "transaction_id": "TX123", + "decline_code": "51", + } + ) + fake_model = ScriptedChatModel(["Transaction analysis completed."]) + agent = ReActAgent(model=fake_model) + + result = await agent.arun( + "Analyze transaction TX123.", + state=state, + ) + + assert result.output == "Transaction analysis completed." + sent_messages = fake_model.received_messages[0] + all_content = "\n".join(message.content or "" for message in sent_messages) + + assert "TX123" in all_content + assert "51" in all_content diff --git a/tests/test_handoff.py b/tests/test_handoff.py new file mode 100644 index 0000000..1f98d88 --- /dev/null +++ b/tests/test_handoff.py @@ -0,0 +1,56 @@ +"""Handoff primitives for multi-agent workflows.""" + +from __future__ import annotations + +from dynavec.agents import AgentResult, AgentState, Handoff +from dynavec.agents.handoff import HandoffRequest + + +def test_handoff_creates_named_routing_request(): + handoff = Handoff() + + request = handoff("policy", "Review transaction TX123.") + + assert isinstance(request, HandoffRequest) + assert request.target == "policy" + assert request.goal == "Review transaction TX123." + + +def test_handoff_rejects_empty_target_or_goal(): + handoff = Handoff() + + try: + handoff("", "Review transaction.") + except ValueError as exc: + assert "target" in str(exc) + else: + raise AssertionError("Expected an empty target to be rejected.") + + try: + handoff("policy", "") + except ValueError as exc: + assert "goal" in str(exc) + else: + raise AssertionError("Expected an empty goal to be rejected.") + + +def test_handoff_invokes_target_agent_with_shared_state(): + class PolicyAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + assert goal == "Review transaction TX123." + assert state.data["decline_code"] == "51" + state.data["policy_result"] = "Manual review required." + return AgentResult(output=state.data["policy_result"]) + + state = AgentState(data={"transaction_id": "TX123", "decline_code": "51"}) + handoff = Handoff() + target = PolicyAgent() + + result = handoff.transfer( + target, + "Review transaction TX123.", + state=state, + ) + + assert result.output == "Manual review required." + assert state.data["policy_result"] == "Manual review required." diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py new file mode 100644 index 0000000..9b98406 --- /dev/null +++ b/tests/test_supervisor.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +from typing import Any + +import pytest + +from dynavec.agents import ( + AgentResult, + AgentState, + Handoff, + ReActAgent, + Supervisor, +) +from dynavec.chat.base import ChatChunk, ChatModel, ChatResult, Message, Tool + + +class ScriptedChatModel(ChatModel): + """Deterministic chat model for supervisor tests.""" + + def __init__(self, responses: list[str]) -> None: + self.responses = responses + self.call_count = 0 + + def invoke( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: Any, + ) -> ChatResult: + response = self.responses[self.call_count] + self.call_count += 1 + + return ChatResult( + message=Message( + role="assistant", + content=response, + ), + finish_reason="stop", + ) + + def stream( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: Any, + ): + result = self.invoke(messages, tools, **kwargs) + + yield ChatChunk( + content=result.message.content, + tool_calls=result.message.tool_calls, + ) + + +class TransactionAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + state.data["transaction_id"] = "TX123" + state.data["decline_code"] = "51" + + return AgentResult( + output="Transaction declined with code 51.", + finished=True, + termination_reason="completed", + total_steps=0, + tool_calls_count=0, + ) + + +class PolicyAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + decline_code = state.data["decline_code"] + + return AgentResult( + output=f"Policy analysis for decline code {decline_code}.", + finished=True, + termination_reason="completed", + total_steps=0, + tool_calls_count=0, + ) + + +def test_supervisor_delegates_to_multiple_agents_with_shared_state(): + state = AgentState() + + supervisor = Supervisor( + agents=[ + TransactionAgent(), + PolicyAgent(), + ], + ) + + results = supervisor.run( + "Analyze the declined transaction.", + state=state, + ) + + assert len(results) == 2 + + assert results[0].output == "Transaction declined with code 51." + assert results[1].output == "Policy analysis for decline code 51." + + assert state.data["transaction_id"] == "TX123" + assert state.data["decline_code"] == "51" + + +def test_supervisor_runs_react_agents(): + state = AgentState() + + first_agent = ReActAgent(model=ScriptedChatModel(["Transaction analysis completed."])) + + second_agent = ReActAgent(model=ScriptedChatModel(["Policy analysis completed."])) + + supervisor = Supervisor( + agents=[first_agent, second_agent], + ) + + results = supervisor.run( + "Analyze transaction TX123.", + state=state, + ) + + assert [result.output for result in results] == [ + "Transaction analysis completed.", + "Policy analysis completed.", + ] + + +def test_supervisor_preserves_state_between_specialists(): + state = AgentState() + + class FirstAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + state.data["transaction_id"] = "TX123" + state.data["decline_code"] = "51" + + return AgentResult( + output="Transaction analyzed.", + finished=True, + termination_reason="completed", + total_steps=0, + tool_calls_count=0, + ) + + class SecondAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + assert state.data["transaction_id"] == "TX123" + assert state.data["decline_code"] == "51" + + return AgentResult( + output="Policy analyzed.", + finished=True, + termination_reason="completed", + total_steps=0, + tool_calls_count=0, + ) + + supervisor = Supervisor( + agents=[FirstAgent(), SecondAgent()], + ) + + results = supervisor.run( + "Analyze the declined transaction.", + state=state, + ) + + assert [result.output for result in results] == [ + "Transaction analyzed.", + "Policy analyzed.", + ] + + +def test_supervisor_routes_only_to_selected_specialist(): + calls: list[str] = [] + + class NamedAgent: + def __init__(self, name: str, output: str) -> None: + self.name = name + self.output = output + + def run(self, goal: str, *, state: AgentState) -> AgentResult: + calls.append(self.name) + return AgentResult(output=self.output) + + agents = { + "transaction": NamedAgent("transaction", "Declined with code 51."), + "policy": NamedAgent("policy", "Manual review required."), + } + + supervisor = Supervisor( + agents=agents, + router=lambda goal, state: "transaction", + ) + + results = supervisor.run("Analyze transaction TX123.", state=AgentState()) + + assert calls == ["transaction"] + assert len(results) == 1 + assert results[0].output == "Declined with code 51." + + +def test_supervisor_rejects_unknown_agent_from_router(): + class NamedAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + return AgentResult(output="Transaction analyzed.") + + supervisor = Supervisor( + agents={"transaction": NamedAgent()}, + router=lambda goal, state: "unknown", + ) + + with pytest.raises(ValueError, match="Router selected unknown agent"): + supervisor.run("Analyze transaction TX123.", state=AgentState()) + + +def test_supervisor_aggregates_only_routed_specialist(): + calls: list[str] = [] + + class NamedAgent: + def __init__(self, name: str, output: str) -> None: + self.name = name + self.output = output + + def run(self, goal: str, *, state: AgentState) -> AgentResult: + calls.append(self.name) + return AgentResult(output=self.output) + + supervisor = Supervisor( + agents={ + "transaction": NamedAgent("transaction", "Declined with code 51."), + "policy": NamedAgent("policy", "Manual review required."), + }, + router=lambda goal, state: "policy", + ) + + result = supervisor.run_and_aggregate( + "Review this transaction.", + state=AgentState(), + ) + + assert calls == ["policy"] + assert result.output == "Manual review required." + + +def test_supervisor_handoff_routes_to_named_specialist_with_shared_state(): + state = AgentState(data={"transaction_id": "TX123"}) + calls = [] + + class FakeAgent: + def __init__(self, output): + self.output = output + + def run(self, goal, *, state): + calls.append((goal, state)) + state.data["policy_review"] = "decline_code_51" + return AgentResult(output=self.output) + + transaction_agent = FakeAgent("Transaction identified") + policy_agent = FakeAgent("Policy review completed") + + supervisor = Supervisor( + { + "transaction": transaction_agent, + "policy": policy_agent, + } + ) + + request = Handoff()( + "policy", + "Review transaction TX123 using the transaction findings.", + ) + + result = supervisor.handoff(request, state=state) + + assert result.output == "Policy review completed" + assert len(calls) == 1 + assert calls[0][0] == ("Review transaction TX123 using the transaction findings.") + assert calls[0][1] is state + assert state.data["policy_review"] == "decline_code_51" diff --git a/tests/test_supervisor_routing.py b/tests/test_supervisor_routing.py new file mode 100644 index 0000000..9d25223 --- /dev/null +++ b/tests/test_supervisor_routing.py @@ -0,0 +1,152 @@ +from dynavec.agents import AgentGraph, AgentResult, AgentState, GraphState, Supervisor + + +def test_supervisor_routes_only_to_selected_specialist(): + calls: list[str] = [] + + def route(state: GraphState) -> str: + return state["selected_specialist"] + + def transaction_agent(state: GraphState) -> dict: + calls.append("transaction") + return {"transaction_result": "Declined with code 51"} + + def policy_agent(state: GraphState) -> dict: + calls.append("policy") + return {"policy_result": "Review required"} + + graph = ( + AgentGraph("supervisor-routing") + .add_node("supervisor", lambda state: {}) + .add_node("transaction", transaction_agent) + .add_node("policy", policy_agent) + .add_conditional_edges( + "supervisor", + route, + branches={ + "transaction": "transaction", + "policy": "policy", + }, + ) + .set_entry("supervisor") + ) + + result = graph.run({"selected_specialist": "transaction"}) + + assert result.success + assert calls == ["transaction"] + assert result.state["transaction_result"] == "Declined with code 51" + assert "policy_result" not in result.state + + +def test_handoff_preserves_findings_for_next_specialist(): + def transaction_agent(state: GraphState) -> dict: + return { + "transaction_result": "Declined", + "decline_code": "51", + } + + def policy_agent(state: GraphState) -> dict: + assert state["decline_code"] == "51" + return {"policy_result": "Manual review required"} + + graph = ( + AgentGraph("supervisor-handoff") + .add_node("transaction", transaction_agent) + .add_node("policy", policy_agent) + .add_edge("transaction", "policy") + .set_entry("transaction") + ) + + result = graph.run({"customer_query": "Why was my transaction declined?"}) + + assert result.success + assert result.state["decline_code"] == "51" + assert result.state["policy_result"] == "Manual review required" + + +def test_supervisor_routes_to_multiple_specialists_with_shared_state(): + state = AgentState() + calls = [] + + class TransactionAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + calls.append("transaction") + state.data["transaction_id"] = "TX123" + state.data["decline_code"] = "51" + return AgentResult(output="Transaction analyzed.") + + class PolicyAgent: + def run(self, goal: str, *, state: AgentState) -> AgentResult: + calls.append("policy") + assert state.data["transaction_id"] == "TX123" + assert state.data["decline_code"] == "51" + return AgentResult(output="Policy reviewed.") + + supervisor = Supervisor( + agents={ + "transaction": TransactionAgent(), + "policy": PolicyAgent(), + }, + router=lambda goal, state: ["transaction", "policy"], + ) + + results = supervisor.run( + "Review transaction TX123.", + state=state, + ) + + assert calls == ["transaction", "policy"] + assert [result.output for result in results] == [ + "Transaction analyzed.", + "Policy reviewed.", + ] + + +def test_handoff_request_routes_to_named_specialist_with_shared_state(): + from dynavec.agents import Handoff + + handoff = Handoff() + request = handoff("policy", "Review transaction TX123.") + + calls = [] + + def supervisor_node(state: GraphState) -> dict: + return { + "handoff_target": request.target, + "handoff_goal": request.goal, + "decline_code": "51", + } + + def transaction_agent(state: GraphState) -> dict: + calls.append("transaction") + return {"unexpected": True} + + def policy_agent(state: GraphState) -> dict: + calls.append("policy") + assert state["handoff_goal"] == "Review transaction TX123." + assert state["decline_code"] == "51" + return {"policy_result": "Manual review required."} + + graph = ( + AgentGraph("handoff-request-routing") + .add_node("supervisor", supervisor_node) + .add_node("transaction", transaction_agent) + .add_node("policy", policy_agent) + .add_conditional_edges( + "supervisor", + lambda state: state["handoff_target"], + branches={ + "transaction": "transaction", + "policy": "policy", + }, + ) + .set_entry("supervisor") + ) + + result = graph.run() + + assert result.success + assert calls == ["policy"] + assert result.state["policy_result"] == "Manual review required." + assert "unexpected" not in result.state