diff --git a/integrations/autogen/src/agentflow_autogen/adapter.py b/integrations/autogen/src/agentflow_autogen/adapter.py index bd3d0d4..fb0c7c2 100644 --- a/integrations/autogen/src/agentflow_autogen/adapter.py +++ b/integrations/autogen/src/agentflow_autogen/adapter.py @@ -16,7 +16,6 @@ from app.models.run import RunStatus from autogen_agentchat.base import TaskResult from autogen_agentchat.messages import ( - BaseChatMessage, ModelClientStreamingChunkEvent, TextMessage, ToolCallExecutionEvent, @@ -24,6 +23,7 @@ ) from agentflow_autogen.tools import ( + BridgedToolExecution, build_function_tools, inject_tools, ) @@ -68,8 +68,9 @@ async def run(self, ctx: AdapterContext) -> AdapterResult: bridged_tools, bridged_names = build_function_tools( surface, ctx, - step_index=handler.current_step_index, + step_index=lambda: handler.current_step_index, on_tool_error=bridged_failure.append, + on_tool_execution=handler.record_bridged_tool, ) else: bridged_tools = [] @@ -132,6 +133,8 @@ def __init__( self._final_reply = "" self._replies: list[str] = [] self._active_tool_calls: dict[str, tuple[str, str, float]] = {} + self._bridged_tool_executions: list[BridgedToolExecution] = [] + self._pending_bridged_calls: dict[str, tuple[str, dict[str, Any]]] = {} self._pending_user_prompt: str | None = None @property @@ -149,15 +152,30 @@ async def begin_step(self, node: str) -> None: self._started_at = time.monotonic() self._tokens_in = 0 self._tokens_out = 0 + self._final_reply = "" await self.ctx.emit_step_started(index=self._current_step_index, node=self._current_node) + async def ensure_step(self, node: str) -> None: + """Start or rotate the active step when an AutoGen turn changes source.""" + if not self.step_started or self.step_completed: + await self.begin_step(node) + return + if self.per_turn_steps and node != self._current_node: + await self.complete_step(output={"reply": self._final_reply}) + await self.begin_step(node) + async def emit_user_prompt(self, prompt: str) -> None: self._pending_user_prompt = prompt step_index = self._current_step_index if self.step_started else None await self.ctx.emit_message(role="user", content=prompt, step_index=step_index) + def record_bridged_tool(self, execution: BridgedToolExecution) -> None: + """Hold tool events until AutoGen emits an event with the owning source.""" + self._bridged_tool_executions.append(execution) + async def handle(self, event: Any, *, bridged_names: frozenset[str]) -> None: if isinstance(event, ModelClientStreamingChunkEvent): + await self.ensure_step(str(event.source or self.default_node)) if self.stream_tokens and isinstance(event.content, str): await self.ctx.emit_token_delta( step_index=self._current_step_index, @@ -191,13 +209,7 @@ async def _handle_text_message(self, message: TextMessage) -> None: ) return - if self.per_turn_steps: - if self.step_started and not self.step_completed: - await self.complete_step(output={"reply": self._final_reply}) - await self.begin_step(str(message.source or self.default_node)) - - if not self.step_started: - await self.begin_step(str(message.source or self.default_node)) + await self.ensure_step(str(message.source or self.default_node)) self._accumulate_usage(message) content = str(message.content or "") @@ -217,11 +229,19 @@ async def _handle_tool_request( *, bridged_names: frozenset[str], ) -> None: + await self.ensure_step(str(event.source or self.default_node)) self._accumulate_usage(event) for call in event.content: + arguments = _parse_tool_arguments(call.arguments) if call.name in bridged_names: + emitted = await self._emit_bridged_tool( + name=call.name, + arguments=arguments, + call_id=call.id, + ) + if not emitted: + self._pending_bridged_calls[call.id] = (call.name, arguments) continue - arguments = _parse_tool_arguments(call.arguments) call_id = await self.ctx.emit_tool_call_started( step_index=self._current_step_index, name=call.name, @@ -235,8 +255,16 @@ async def _handle_tool_execution( *, bridged_names: frozenset[str], ) -> None: + await self.ensure_step(str(event.source or self.default_node)) for result in event.content: if result.name in bridged_names: + pending = self._pending_bridged_calls.pop(result.call_id, None) + name, arguments = pending or (result.name, {}) + await self._emit_bridged_tool( + name=name, + arguments=arguments, + call_id=result.call_id, + ) continue active = self._active_tool_calls.pop(result.call_id, None) if active is None: @@ -258,6 +286,50 @@ async def _handle_tool_execution( latency_ms=int((time.monotonic() - started) * 1000), ) + async def _emit_bridged_tool( + self, + *, + name: str, + arguments: dict[str, Any], + call_id: str, + ) -> bool: + execution_index = next( + ( + index + for index, execution in enumerate(self._bridged_tool_executions) + if execution.name == name and execution.arguments == arguments + ), + None, + ) + if execution_index is None: + execution_index = next( + ( + index + for index, execution in enumerate(self._bridged_tool_executions) + if execution.name == name + ), + None, + ) + if execution_index is None: + return False + + execution = self._bridged_tool_executions.pop(execution_index) + emitted_call_id = await self.ctx.emit_tool_call_started( + step_index=self._current_step_index, + name=execution.name, + arguments=execution.arguments, + call_id=call_id, + ) + await self.ctx.emit_tool_call_completed( + step_index=self._current_step_index, + name=execution.name, + call_id=emitted_call_id, + result=execution.result, + error=execution.error, + latency_ms=execution.latency_ms, + ) + return True + async def _handle_task_result(self, result: TaskResult) -> None: for message in result.messages: if isinstance(message, TextMessage) and message.source != USER_SOURCE: diff --git a/integrations/autogen/src/agentflow_autogen/tools.py b/integrations/autogen/src/agentflow_autogen/tools.py index 8653338..e113ef6 100644 --- a/integrations/autogen/src/agentflow_autogen/tools.py +++ b/integrations/autogen/src/agentflow_autogen/tools.py @@ -3,7 +3,10 @@ from __future__ import annotations import json -from typing import Annotated, Any, Callable +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import Annotated, Any from app.adapters.adapter_tools import AdapterToolSurface from app.adapters.base import AdapterContext @@ -11,12 +14,24 @@ from autogen_core.tools import FunctionTool +@dataclass(frozen=True, slots=True) +class BridgedToolExecution: + """A tool result waiting for AutoGen to reveal its owning turn.""" + + name: str + arguments: dict[str, Any] + result: dict[str, Any] | None + error: str | None + latency_ms: int + + def build_function_tools( surface: AdapterToolSurface, ctx: AdapterContext, *, - step_index: int, + step_index: int | Callable[[], int], on_tool_error: Callable[[str], None] | None = None, + on_tool_execution: Callable[[BridgedToolExecution], None] | None = None, ) -> tuple[list[FunctionTool], frozenset[str]]: """Create AutoGen tools that delegate execution to ``AdapterToolSurface``.""" tools: list[FunctionTool] = [] @@ -29,6 +44,7 @@ def build_function_tools( step_index=step_index, definition=definition, on_tool_error=on_tool_error, + on_tool_execution=on_tool_execution, ) tools.append( FunctionTool( @@ -46,20 +62,50 @@ def _make_handler( surface: AdapterToolSurface, ctx: AdapterContext, *, - step_index: int, + step_index: int | Callable[[], int], definition: ToolDefinition, on_tool_error: Callable[[str], None] | None = None, + on_tool_execution: Callable[[BridgedToolExecution], None] | None = None, ): properties = (definition.parameters or {}).get("properties") or {} async def _execute(arguments: dict[str, Any]) -> str: try: - result = await surface.execute( - ctx, - step_index=step_index, - name=definition.name, - arguments=arguments, - ) + if on_tool_execution is None: + current_step_index = step_index() if callable(step_index) else step_index + result = await surface.execute( + ctx, + step_index=current_step_index, + name=definition.name, + arguments=arguments, + ) + else: + started = time.monotonic() + try: + result = await surface.lookup(definition.name).handler(arguments) + except Exception as exc: + error = str(exc) or type(exc).__name__ + on_tool_execution( + BridgedToolExecution( + name=definition.name, + arguments=arguments, + result=None, + error=error, + latency_ms=int((time.monotonic() - started) * 1000), + ) + ) + raise + if not isinstance(result, dict): + result = {"result": result} + on_tool_execution( + BridgedToolExecution( + name=definition.name, + arguments=arguments, + result=result, + error=None, + latency_ms=int((time.monotonic() - started) * 1000), + ) + ) return _stringify_result(result) except Exception as exc: if on_tool_error is not None: @@ -67,20 +113,20 @@ async def _execute(arguments: dict[str, Any]) -> str: return str(exc) or type(exc).__name__ if not properties: - async def handler() -> str: + async def no_arguments_handler() -> str: return await _execute({}) - handler.__name__ = definition.name - return handler + no_arguments_handler.__name__ = definition.name + return no_arguments_handler if set(properties) == {"text"}: - async def handler(text: Annotated[str, "Tool input text."] = "") -> str: + async def text_handler(text: Annotated[str, "Tool input text."] = "") -> str: return await _execute({"text": text}) - handler.__name__ = definition.name - return handler + text_handler.__name__ = definition.name + return text_handler - async def handler(input_json: Annotated[str, "JSON-encoded tool arguments."] = "{}") -> str: + async def json_handler(input_json: Annotated[str, "JSON-encoded tool arguments."] = "{}") -> str: try: parsed = json.loads(input_json or "{}") except json.JSONDecodeError: @@ -88,8 +134,8 @@ async def handler(input_json: Annotated[str, "JSON-encoded tool arguments."] = " arguments = parsed if isinstance(parsed, dict) else {"value": parsed} return await _execute(arguments) - handler.__name__ = definition.name - return handler + json_handler.__name__ = definition.name + return json_handler def _stringify_result(result: dict[str, Any]) -> str: diff --git a/integrations/autogen/tests/test_adapter.py b/integrations/autogen/tests/test_adapter.py index c8b2496..89735be 100644 --- a/integrations/autogen/tests/test_adapter.py +++ b/integrations/autogen/tests/test_adapter.py @@ -133,6 +133,64 @@ def create_team() -> RoundRobinGroupChat: return RoundRobinGroupChat([writer_agent(), editor_agent()], max_turns=2) +def bridged_team_agent( + name: str, + reply: str, + *, + agentflow_tools=None, +) -> AssistantAgent: + usage = RequestUsage(prompt_tokens=5, completion_tokens=2) + tools = list(agentflow_tools or []) + tool_name = getattr(tools[0], "name", "echo") if tools else "echo" + call = FunctionCall( + id=f"{name}-bridge-call", + name=tool_name, + arguments=json.dumps({"text": name}), + ) + client = ReplayChatCompletionClient( + [ + CreateResult(content=[call], finish_reason="function_calls", usage=usage, cached=False), + CreateResult(content=reply, finish_reason="stop", usage=usage, cached=False), + ], + model_info=ModelInfo( + vision=False, + function_calling=True, + json_output=False, + family="unknown", + ), + ) + return AssistantAgent( + name=name, + model_client=client, + tools=tools, + reflect_on_tool_use=True, + ) + + +def create_bridged_team(*, agentflow_tools=None) -> RoundRobinGroupChat: + tools = list(agentflow_tools or []) + return RoundRobinGroupChat( + [ + bridged_team_agent("writer", "Writer used the bridge", agentflow_tools=tools), + bridged_team_agent("editor", "Editor used the bridge", agentflow_tools=tools), + ], + max_turns=2, + ) + + +class RepeatedMessageTeam: + async def run_stream(self, *, task: str): + first = TextMessage(content="Draft", source="writer") + final = TextMessage(content="Revised draft", source="writer") + yield first + yield final + yield TaskResult(messages=[first, final]) + + +def repeated_message_team() -> RepeatedMessageTeam: + return RepeatedMessageTeam() + + def failing_bridged_agent(*, agentflow_tools=None) -> AssistantAgent: usage = RequestUsage(prompt_tokens=5, completion_tokens=2) call = FunctionCall(id="bridge-call", name="bridge_failure", arguments="{}") @@ -236,6 +294,61 @@ async def test_team_run_emits_one_step_per_turn() -> None: assert len(ctx.event("step.completed")) == 2 +@pytest.mark.asyncio +async def test_team_bridged_tools_follow_the_live_turn_step() -> None: + ctx = RecordingContext( + { + "team_factory": factory_reference("create_bridged_team"), + "per_turn_steps": True, + "tools": ["echo"], + }, + prompt="collab with tools", + ) + + result = await AutoGenAdapter().run(ctx) + + assert result.status == RunStatus.SUCCEEDED + assert result.output["replies"] == ["Writer used the bridge", "Editor used the bridge"] + assert ctx.event("step.started") == [ + {"index": 0, "node": "writer"}, + {"index": 1, "node": "editor"}, + ] + assert [event["step_index"] for event in ctx.event("tool_call.started")] == [0, 1] + assert [event["step_index"] for event in ctx.event("tool_call.completed")] == [0, 1] + assert [event["step_index"] for event in ctx.event("message.created") if event["role"] == "assistant"] == [0, 1] + + for step_index in (0, 1): + step_position = next( + index + for index, (event_type, data) in enumerate(ctx.events) + if event_type == "step.started" and data["index"] == step_index + ) + tool_position = next( + index + for index, (event_type, data) in enumerate(ctx.events) + if event_type == "tool_call.started" and data["step_index"] == step_index + ) + assert step_position < tool_position + + +@pytest.mark.asyncio +async def test_repeated_messages_from_one_agent_stay_in_the_same_turn_step() -> None: + ctx = RecordingContext( + { + "team_factory": factory_reference("repeated_message_team"), + "per_turn_steps": True, + } + ) + + result = await AutoGenAdapter().run(ctx) + + assert result.status == RunStatus.SUCCEEDED + assert result.output == {"reply": "Revised draft", "replies": ["Draft", "Revised draft"]} + assert ctx.event("step.started") == [{"index": 0, "node": "writer"}] + assert len(ctx.event("step.completed")) == 1 + assert [event["step_index"] for event in ctx.event("message.created") if event["role"] == "assistant"] == [0, 0] + + @pytest.mark.asyncio async def test_invalid_factory_becomes_a_failed_step() -> None: ctx = RecordingContext({"agent_factory": factory_reference("not_a_runnable")})