From fc355efe0ae0827600bbc756ea573069019cecd8 Mon Sep 17 00:00:00 2001 From: Vardhman Gupta Date: Sat, 3 Oct 2026 09:18:58 +0530 Subject: [PATCH] feat: agent primitives (ReAct, Planner) for #288 --- src/dynavec/agents/__init__.py | 25 +++ src/dynavec/agents/base.py | 192 ++++++++++++++++++ src/dynavec/agents/planner.py | 124 ++++++++++++ src/dynavec/agents/react.py | 192 ++++++++++++++++++ tests/test_agents.py | 345 +++++++++++++++++++++++++++++++++ 5 files changed, 878 insertions(+) create mode 100644 src/dynavec/agents/__init__.py create mode 100644 src/dynavec/agents/base.py create mode 100644 src/dynavec/agents/planner.py create mode 100644 src/dynavec/agents/react.py create mode 100644 tests/test_agents.py diff --git a/src/dynavec/agents/__init__.py b/src/dynavec/agents/__init__.py new file mode 100644 index 0000000..c0110cf --- /dev/null +++ b/src/dynavec/agents/__init__.py @@ -0,0 +1,25 @@ +"""Agent primitives and execution loops for dynaflow.""" + +from __future__ import annotations + +from .base import ( + AgentResult, + AgentStep, + AgentTool, + Plan, + PlanStep, + tool, +) +from .planner import Planner +from .react import ReActAgent + +__all__ = [ + "AgentResult", + "AgentStep", + "AgentTool", + "Plan", + "PlanStep", + "Planner", + "ReActAgent", + "tool", +] diff --git a/src/dynavec/agents/base.py b/src/dynavec/agents/base.py new file mode 100644 index 0000000..9fea9b2 --- /dev/null +++ b/src/dynavec/agents/base.py @@ -0,0 +1,192 @@ +"""Core data models and tool primitives for dynaflow agents.""" + +from __future__ import annotations + +import inspect +import json +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +from ..chat.base import Tool, ToolCall + + +@dataclass +class AgentStep: + """A single execution step in an agent's reasoning loop.""" + + step_number: int + thought: str | None = None + tool_calls: list[ToolCall] = field(default_factory=list) + observations: list[str] = field(default_factory=list) + + +@dataclass +class AgentResult: + """The final result of an agent run.""" + + output: str + steps: list[AgentStep] = field(default_factory=list) + finished: bool = True + termination_reason: str = "completed" # "completed", "max_steps_reached", "error" + total_steps: int = 0 + tool_calls_count: int = 0 + + +@dataclass +class PlanStep: + """A single step in a decomposed task plan.""" + + step_number: int + description: str + tool_hint: str | None = None + + +@dataclass +class Plan: + """An ordered decomposition of a goal into actionable steps.""" + + goal: str + steps: list[PlanStep] = field(default_factory=list) + + +def _python_type_to_json_type(py_type: Any) -> str: + """Map standard Python types and type annotation strings to JSON Schema data types.""" + if py_type in (str, "str", "string"): + return "string" + if py_type in (int, "int", "integer"): + return "integer" + if py_type in (float, "float", "number"): + return "number" + if py_type in (bool, "bool", "boolean"): + return "boolean" + if py_type in (list, tuple, set, "list", "tuple", "set") or ( + isinstance(py_type, str) and py_type.startswith(("list[", "Sequence[", "tuple[")) + ): + return "array" + if py_type in (dict, Any, "dict", "dict[str, Any]", "Mapping") or ( + isinstance(py_type, str) and py_type.startswith("dict[") + ): + return "object" + return "string" + + +def _generate_json_schema(fn: Callable[..., Any]) -> dict[str, Any]: + """Generate a JSON schema parameters dictionary from a function signature.""" + sig = inspect.signature(fn) + properties: dict[str, Any] = {} + required: list[str] = [] + + for param_name, param in sig.parameters.items(): + if param_name in ("self", "cls"): + continue + + param_type = param.annotation + json_type = ( + _python_type_to_json_type(param_type) + if param_type is not inspect.Parameter.empty + else "string" + ) + + properties[param_name] = { + "type": json_type, + "description": f"Parameter '{param_name}'", + } + + if param.default is inspect.Parameter.empty: + required.append(param_name) + + schema: dict[str, Any] = { + "type": "object", + "properties": properties, + } + if required: + schema["required"] = required + return schema + + +class AgentTool: + """An executable tool wrapped with a typed JSON schema for agent use.""" + + def __init__( + self, + fn: Callable[..., Any], + name: str | None = None, + description: str | None = None, + parameters: dict[str, Any] | None = None, + ) -> None: + self.fn = fn + self.name = name or fn.__name__ + self.description = description or (fn.__doc__ or f"Execute {self.name}").strip() + self.parameters = ( + parameters if parameters is not None else _generate_json_schema(fn) + ) + + def to_chat_tool(self) -> Tool: + """Convert to a dynavec.chat.Tool schema.""" + return Tool( + name=self.name, + description=self.description, + parameters=self.parameters, + ) + + def execute(self, arguments: dict[str, Any] | str | None = None) -> str: + """Execute the wrapped function and return a string observation.""" + parsed_args: dict[str, Any] = {} + if isinstance(arguments, str): + if arguments.strip(): + try: + loaded = json.loads(arguments) + if isinstance(loaded, dict): + parsed_args = loaded + else: + parsed_args = {"input": loaded} + except Exception: + parsed_args = {"input": arguments} + elif isinstance(arguments, dict): + parsed_args = arguments + + try: + # Check if function accepts kwargs or positional + sig = inspect.signature(self.fn) + params = sig.parameters + + if not params: + result = self.fn() + elif len(params) == 1 and list(params.keys())[0] not in parsed_args: + # If single param expected and keys don't match, pass the first val or raw dict + first_val = ( + next(iter(parsed_args.values())) if parsed_args else arguments + ) + result = self.fn(first_val) + else: + # Filter only valid keyword arguments + valid_args = {k: v for k, v in parsed_args.items() if k in params} + result = self.fn(**valid_args) + + if isinstance(result, str): + return result + return json.dumps(result, ensure_ascii=False) + except Exception as exc: # noqa: BLE001 + return f"Error executing tool {self.name!r}: {exc}" + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + return self.fn(*args, **kwargs) + + +def tool( + name: str | None = None, + description: str | None = None, + parameters: dict[str, Any] | None = None, +) -> Callable[[Callable[..., Any]], AgentTool]: + """Decorator to convert a standard Python function into an AgentTool.""" + + def decorator(fn: Callable[..., Any]) -> AgentTool: + return AgentTool( + fn=fn, + name=name, + description=description, + parameters=parameters, + ) + + return decorator diff --git a/src/dynavec/agents/planner.py b/src/dynavec/agents/planner.py new file mode 100644 index 0000000..3c1a1b5 --- /dev/null +++ b/src/dynavec/agents/planner.py @@ -0,0 +1,124 @@ +"""Planner primitive for decomposing goals into structured action plans.""" + +from __future__ import annotations + +import json +import re +from typing import Any + +from ..chat.base import ChatModel, Message +from .base import Plan, PlanStep + +DEFAULT_PLANNER_SYSTEM_PROMPT = """You are an expert task planning agent. +Your job is to break down complex goals into an ordered, clear sequence of actionable steps. +Output your plan strictly as a JSON object matching this schema: +{ + "steps": [ + { + "step_number": 1, + "description": "Description of what needs to be done in this step", + "tool_hint": "Optional name of the tool or action suited for this step" + } + ] +} +Do not include any conversational filler or markdown other than the valid JSON.""" + + +def _extract_json_block(text: str) -> str: + """Extract JSON content from markdown code fences or plain text.""" + trimmed = text.strip() + match = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", trimmed, re.IGNORECASE) + if match: + return match.group(1).strip() + return trimmed + + +class Planner: + """Decomposes goals into structured, ordered execution plans.""" + + def __init__( + self, + model: ChatModel, + system_prompt: str | None = None, + ) -> None: + self.model = model + self.system_prompt = ( + system_prompt if system_prompt is not None else DEFAULT_PLANNER_SYSTEM_PROMPT + ) + + def _parse_plan(self, goal: str, response_text: str) -> Plan: + """Parse raw model output into a Plan object.""" + cleaned = _extract_json_block(response_text) + try: + data = json.loads(cleaned) + except Exception: + # Fallback: if json parsing fails, split lines into steps + lines = [ + line.strip() + for line in response_text.splitlines() + if line.strip() and not line.startswith("```") + ] + steps = [ + PlanStep(step_number=idx, description=line) + for idx, line in enumerate(lines, start=1) + ] + return Plan(goal=goal, steps=steps) + + steps_data = data.get("steps", []) if isinstance(data, dict) else data + plan_steps: list[PlanStep] = [] + + if isinstance(steps_data, list): + for idx, item in enumerate(steps_data, start=1): + if isinstance(item, dict): + step_num = item.get("step_number", idx) + desc = item.get("description", str(item)) + tool_hint = item.get("tool_hint") + plan_steps.append( + PlanStep( + step_number=step_num, + description=desc, + tool_hint=tool_hint, + ) + ) + elif isinstance(item, str): + plan_steps.append(PlanStep(step_number=idx, description=item)) + + return Plan(goal=goal, steps=plan_steps) + + def plan( + self, + goal: str, + context: str | None = None, + **kwargs: Any, + ) -> Plan: + """Decompose a goal into an ordered Plan synchronously.""" + messages: list[Message] = [] + if self.system_prompt: + messages.append(Message(role="system", content=self.system_prompt)) + + user_content = f"Goal: {goal}" + if context: + user_content += f"\n\nContext:\n{context}" + messages.append(Message(role="user", content=user_content)) + + res = self.model.invoke(messages, **kwargs) + return self._parse_plan(goal, res.message.content or "") + + async def aplan( + self, + goal: str, + context: str | None = None, + **kwargs: Any, + ) -> Plan: + """Decompose a goal into an ordered Plan asynchronously.""" + messages: list[Message] = [] + if self.system_prompt: + messages.append(Message(role="system", content=self.system_prompt)) + + user_content = f"Goal: {goal}" + if context: + user_content += f"\n\nContext:\n{context}" + messages.append(Message(role="user", content=user_content)) + + res = await self.model.ainvoke(messages, **kwargs) + return self._parse_plan(goal, res.message.content or "") diff --git a/src/dynavec/agents/react.py b/src/dynavec/agents/react.py new file mode 100644 index 0000000..67b83ee --- /dev/null +++ b/src/dynavec/agents/react.py @@ -0,0 +1,192 @@ +"""ReAct (Reason + Act) tool-calling agent implementation.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Sequence +from typing import Any + +from ..chat.base import ChatModel, Message, Tool +from .base import AgentResult, AgentStep, AgentTool + +DEFAULT_REACT_SYSTEM_PROMPT = """You are a helpful and precise reasoning agent. +You solve tasks step-by-step using a ReAct (Reason + Act) approach. +When you need to look up information or perform calculations, invoke the appropriate tools. +When you have the final answer to the user's goal, provide the clear and complete answer directly without calling any tools.""" + + +class ReActAgent: + """A ReAct-style agent that alternates between thinking, calling tools, and observing results.""" + + def __init__( + self, + model: ChatModel, + tools: Sequence[AgentTool | Callable[..., Any]] | None = None, + system_prompt: str | None = None, + max_steps: int = 10, + ) -> None: + self.model = model + self.system_prompt = ( + system_prompt if system_prompt is not None else DEFAULT_REACT_SYSTEM_PROMPT + ) + self.max_steps = max(1, max_steps) + + self._tools_map: dict[str, AgentTool] = {} + self._chat_tools: list[Tool] = [] + + if tools: + for t in tools: + agent_tool = t if isinstance(t, AgentTool) else AgentTool(t) + self._tools_map[agent_tool.name] = agent_tool + self._chat_tools.append(agent_tool.to_chat_tool()) + + def run(self, goal: str, **kwargs: Any) -> AgentResult: + """Execute the ReAct loop synchronously until goal completion or max_steps.""" + messages: list[Message] = [] + if self.system_prompt: + messages.append(Message(role="system", content=self.system_prompt)) + messages.append(Message(role="user", content=goal)) + + steps: list[AgentStep] = [] + total_tool_calls = 0 + + for step_idx in range(1, self.max_steps + 1): + chat_res = self.model.invoke( + messages, + tools=self._chat_tools if self._chat_tools else None, + **kwargs, + ) + msg = chat_res.message + messages.append(msg) + + if not msg.tool_calls: + # Final response reached without further tool calls + step = AgentStep( + step_number=step_idx, + thought=msg.content, + tool_calls=[], + observations=[], + ) + steps.append(step) + return AgentResult( + output=msg.content or "", + steps=steps, + finished=True, + termination_reason="completed", + total_steps=step_idx, + tool_calls_count=total_tool_calls, + ) + + # Execute tool calls + step_observations: list[str] = [] + for tc in msg.tool_calls: + total_tool_calls += 1 + tool_instance = self._tools_map.get(tc.name) + if tool_instance is not None: + obs = tool_instance.execute(tc.arguments) + else: + obs = f"Error: Tool {tc.name!r} is not registered in available tools." + + step_observations.append(obs) + messages.append( + Message( + role="tool", + content=obs, + tool_call_id=tc.id, + ) + ) + + step = AgentStep( + step_number=step_idx, + thought=msg.content, + tool_calls=list(msg.tool_calls), + observations=step_observations, + ) + steps.append(step) + + # Reached max steps without completing + last_output = steps[-1].thought or ( + steps[-1].observations[-1] if steps[-1].observations else "" + ) + return AgentResult( + output=last_output, + steps=steps, + finished=False, + termination_reason="max_steps_reached", + total_steps=self.max_steps, + tool_calls_count=total_tool_calls, + ) + + async def arun(self, goal: str, **kwargs: Any) -> AgentResult: + """Execute the ReAct loop asynchronously.""" + messages: list[Message] = [] + if self.system_prompt: + messages.append(Message(role="system", content=self.system_prompt)) + messages.append(Message(role="user", content=goal)) + + steps: list[AgentStep] = [] + total_tool_calls = 0 + + for step_idx in range(1, self.max_steps + 1): + chat_res = await self.model.ainvoke( + messages, + tools=self._chat_tools if self._chat_tools else None, + **kwargs, + ) + msg = chat_res.message + messages.append(msg) + + if not msg.tool_calls: + step = AgentStep( + step_number=step_idx, + thought=msg.content, + tool_calls=[], + observations=[], + ) + steps.append(step) + return AgentResult( + output=msg.content or "", + steps=steps, + finished=True, + termination_reason="completed", + total_steps=step_idx, + tool_calls_count=total_tool_calls, + ) + + step_observations: list[str] = [] + for tc in msg.tool_calls: + total_tool_calls += 1 + tool_instance = self._tools_map.get(tc.name) + if tool_instance is not None: + obs = await asyncio.to_thread(tool_instance.execute, tc.arguments) + else: + obs = f"Error: Tool {tc.name!r} is not registered in available tools." + + step_observations.append(obs) + messages.append( + Message( + role="tool", + content=obs, + tool_call_id=tc.id, + ) + ) + + step = AgentStep( + step_number=step_idx, + thought=msg.content, + tool_calls=list(msg.tool_calls), + observations=step_observations, + ) + steps.append(step) + + last_output = steps[-1].thought or ( + steps[-1].observations[-1] if steps[-1].observations else "" + ) + return AgentResult( + output=last_output, + steps=steps, + finished=False, + termination_reason="max_steps_reached", + total_steps=self.max_steps, + tool_calls_count=total_tool_calls, + ) diff --git a/tests/test_agents.py b/tests/test_agents.py new file mode 100644 index 0000000..4ee91cb --- /dev/null +++ b/tests/test_agents.py @@ -0,0 +1,345 @@ +"""Unit tests for agent primitives (ReActAgent, Planner, AgentTool).""" + +from __future__ import annotations + +import json +from collections.abc import Iterator +from typing import Any + +import pytest + +from dynavec.agents import ( + AgentTool, + Plan, + Planner, + ReActAgent, + tool, +) +from dynavec.chat.base import ( + ChatChunk, + ChatModel, + ChatResult, + Message, + Tool, + ToolCall, +) + + +class ScriptedChatModel(ChatModel): + """A deterministic test fake for ChatModel that replays configured responses.""" + + def __init__(self, responses: list[ChatResult | Message | str]) -> None: + self.responses = list(responses) + self.call_count = 0 + self.received_messages: list[list[Message]] = [] + + def invoke( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: Any, + ) -> ChatResult: + self.received_messages.append(list(messages)) + if self.call_count >= len(self.responses): + # Default response if script exhausted + return ChatResult( + message=Message( + role="assistant", content="Scripted responses exhausted." + ), + finish_reason="stop", + ) + + resp = self.responses[self.call_count] + self.call_count += 1 + + if isinstance(resp, ChatResult): + return resp + if isinstance(resp, Message): + return ChatResult(message=resp, finish_reason="stop") + return ChatResult( + message=Message(role="assistant", content=resp), + finish_reason="stop", + ) + + def stream( + self, + messages: list[Message], + tools: list[Tool] | None = None, + **kwargs: Any, + ) -> Iterator[ChatChunk]: + res = self.invoke(messages, tools, **kwargs) + yield ChatChunk(content=res.message.content, tool_calls=res.message.tool_calls) + + +def test_tool_decorator_and_schema_generation() -> None: + """Verify @tool extracts docstrings, names, parameter types, and executes correctly.""" + + @tool(name="custom_add", description="Add two integers.") + def add(a: int, b: int = 0) -> int: + """Add two numbers.""" + return a + b + + assert isinstance(add, AgentTool) + assert add.name == "custom_add" + assert add.description == "Add two integers." + assert add.parameters["type"] == "object" + assert add.parameters["properties"]["a"]["type"] == "integer" + assert add.parameters["properties"]["b"]["type"] == "integer" + assert "a" in add.parameters.get("required", []) + assert "b" not in add.parameters.get("required", []) + + # Test execution with dict + assert add.execute({"a": 5, "b": 10}) == "15" + # Test execution with JSON string + assert add.execute('{"a": 20, "b": 22}') == "42" + # Test direct call + assert add(3, 4) == 7 + + +def test_react_agent_two_tool_task() -> None: + """Acceptance criteria: Agent solves a 2-tool sequential task via the ReAct loop.""" + + # Define tools + @tool() + def get_user_id(username: str) -> str: + """Retrieve user ID by username.""" + if username == "alice": + return "user_123" + return "not_found" + + @tool() + def get_account_balance(user_id: str) -> float: + """Fetch account balance for a user ID.""" + if user_id == "user_123": + return 2500.50 + return 0.0 + + # Scripted model responses: + # 1. Call get_user_id("alice") + # 2. Call get_account_balance("user_123") + # 3. Final answer + step1_msg = Message( + role="assistant", + content="I will look up Alice's user ID.", + tool_calls=[ + ToolCall( + id="call_1", + name="get_user_id", + arguments=json.dumps({"username": "alice"}), + ) + ], + ) + step2_msg = Message( + role="assistant", + content="I found the user ID. Now I will check the balance.", + tool_calls=[ + ToolCall( + id="call_2", + name="get_account_balance", + arguments=json.dumps({"user_id": "user_123"}), + ) + ], + ) + step3_msg = Message( + role="assistant", + content="Alice's account balance is $2,500.50.", + tool_calls=[], + ) + + fake_model = ScriptedChatModel([step1_msg, step2_msg, step3_msg]) + agent = ReActAgent( + model=fake_model, + tools=[get_user_id, get_account_balance], + max_steps=5, + ) + + result = agent.run("What is Alice's balance?") + + assert result.finished is True + assert result.termination_reason == "completed" + assert result.total_steps == 3 + assert result.tool_calls_count == 2 + assert "2,500.50" in result.output + assert len(result.steps) == 3 + + # Validate step observations + assert result.steps[0].observations == ["user_123"] + assert result.steps[1].observations == ["2500.5"] + assert result.steps[2].observations == [] + + +def test_react_agent_max_steps_budget_guard() -> None: + """Acceptance criteria: Agent terminates cleanly when max_steps budget is exhausted.""" + + # Infinite loop model that keeps asking for tools + infinite_tool_msg = Message( + role="assistant", + content="Still thinking...", + tool_calls=[ + ToolCall( + id="call_loop", + name="dummy_tool", + arguments='{"val": "x"}', + ) + ], + ) + fake_model = ScriptedChatModel([infinite_tool_msg] * 10) + + @tool() + def dummy_tool(val: str) -> str: + """A dummy test tool.""" + return f"result_{val}" + + agent = ReActAgent( + model=fake_model, + tools=[dummy_tool], + max_steps=3, + ) + + result = agent.run("Run indefinitely") + + assert result.finished is False + assert result.termination_reason == "max_steps_reached" + assert result.total_steps == 3 + assert result.tool_calls_count == 3 + + +def test_react_agent_tool_error_handling() -> None: + """Verify tool execution exceptions are caught and passed back as observations.""" + + @tool() + def buggy_tool(x: int) -> float: + """A tool that raises zero division.""" + return 10 / x + + # Model calls buggy_tool with x=0, gets error observation, then returns answer + step1_msg = Message( + role="assistant", + tool_calls=[ + ToolCall( + id="call_err", + name="buggy_tool", + arguments='{"x": 0}', + ) + ], + ) + step2_msg = Message( + role="assistant", + content="Encountered division by zero error.", + tool_calls=[], + ) + + fake_model = ScriptedChatModel([step1_msg, step2_msg]) + agent = ReActAgent(model=fake_model, tools=[buggy_tool], max_steps=5) + + result = agent.run("Divide 10 by 0") + + assert result.finished is True + assert "division by zero" in result.steps[0].observations[0].lower() + + +def test_react_agent_unregistered_tool() -> None: + """Verify calling an unregistered tool returns an error message observation.""" + step1_msg = Message( + role="assistant", + tool_calls=[ + ToolCall( + id="call_unknown", + name="unknown_tool", + arguments="{}", + ) + ], + ) + step2_msg = Message(role="assistant", content="Unknown tool handled.") + + fake_model = ScriptedChatModel([step1_msg, step2_msg]) + agent = ReActAgent(model=fake_model, tools=[], max_steps=5) + + result = agent.run("Call missing tool") + + assert result.finished is True + assert "not registered" in result.steps[0].observations[0].lower() + + +@pytest.mark.asyncio +async def test_react_agent_async_arun() -> None: + """Verify async arun execution works.""" + + @tool() + def greet(name: str) -> str: + """Greet a person.""" + return f"Hello, {name}!" + + step1_msg = Message( + role="assistant", + tool_calls=[ + ToolCall( + id="call_greet", + name="greet", + arguments='{"name": "Dynavec"}', + ) + ], + ) + step2_msg = Message(role="assistant", content="Greeting completed.") + + fake_model = ScriptedChatModel([step1_msg, step2_msg]) + agent = ReActAgent(model=fake_model, tools=[greet], max_steps=5) + + result = await agent.arun("Say hello") + + assert result.finished is True + assert result.total_steps == 2 + assert result.steps[0].observations == ["Hello, Dynavec!"] + + +def test_planner_structured_decomposition() -> None: + """Verify Planner decomposes a goal into ordered PlanSteps from JSON.""" + plan_json = json.dumps( + { + "steps": [ + { + "step_number": 1, + "description": "Fetch documentation from knowledge base", + "tool_hint": "dynavec_search", + }, + { + "step_number": 2, + "description": "Summarize key findings", + "tool_hint": "summarizer", + }, + ] + } + ) + + fake_model = ScriptedChatModel([f"```json\n{plan_json}\n```"]) + planner = Planner(model=fake_model) + + plan = planner.plan("Research vector database scaling") + + assert isinstance(plan, Plan) + assert plan.goal == "Research vector database scaling" + assert len(plan.steps) == 2 + assert plan.steps[0].step_number == 1 + assert "documentation" in plan.steps[0].description + assert plan.steps[0].tool_hint == "dynavec_search" + assert plan.steps[1].step_number == 2 + + +@pytest.mark.asyncio +async def test_planner_async_aplan() -> None: + """Verify async aplan method.""" + plan_json = json.dumps( + { + "steps": [ + {"step_number": 1, "description": "Step A"}, + {"step_number": 2, "description": "Step B"}, + ] + } + ) + fake_model = ScriptedChatModel([plan_json]) + planner = Planner(model=fake_model) + + plan = await planner.aplan("Async task") + assert len(plan.steps) == 2 + assert plan.steps[0].description == "Step A" + assert plan.steps[1].description == "Step B"