lite-agent 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

lite_agent/runner.py

@@ -1,60 +1,45 @@
 import json
+import warnings
 from collections.abc import AsyncGenerator, Sequence
 from datetime import datetime, timedelta, timezone
 from os import PathLike
 from pathlib import Path
-from typing import Any, Literal
+from typing import Any, Literal, cast
 
 from lite_agent.agent import Agent
+from lite_agent.constants import CompletionMode, StreamIncludes, ToolName
 from lite_agent.loggers import logger
 from lite_agent.types import (
     AgentChunk,
     AgentChunkType,
-    AssistantMessageContent,
     AssistantMessageMeta,
     AssistantTextContent,
     AssistantToolCall,
     AssistantToolCallResult,
+    FlexibleInputMessage,
     FlexibleRunnerMessage,
-    MessageDict,
     MessageUsage,
     NewAssistantMessage,
     NewMessage,
     NewSystemMessage,
-    # New structured message types
     NewUserMessage,
     ToolCall,
     ToolCallFunction,
-    UserImageContent,
     UserInput,
-    UserMessageContent,
     UserTextContent,
 )
-from lite_agent.types.events import AssistantMessageEvent
-
-DEFAULT_INCLUDES: tuple[AgentChunkType, ...] = (
-    "completion_raw",
-    "usage",
-    "function_call",
-    "function_call_output",
-    "content_delta",
-    "function_call_delta",
-    "assistant_message",
-)
+from lite_agent.types.events import AssistantMessageEvent, FunctionCallOutputEvent
+from lite_agent.utils.message_builder import MessageBuilder
 
 
 class Runner:
-    def __init__(self, agent: Agent, api: Literal["completion", "responses"] = "responses", streaming: bool = True) -> None:
+    def __init__(self, agent: Agent, api: Literal["completion", "responses"] = "responses", *, streaming: bool = True) -> None:
         self.agent = agent
-        self.messages: list[NewMessage] = []
+        self.messages: list[FlexibleRunnerMessage] = []
         self.api = api
         self.streaming = streaming
         self._current_assistant_message: NewAssistantMessage | None = None
-
-    @property
-    def legacy_messages(self) -> list[NewMessage]:
-        """Return messages in new format (legacy_messages is now an alias)."""
-        return self.messages
+        self.usage = MessageUsage(input_tokens=0, output_tokens=0, total_tokens=0)
 
     def _start_assistant_message(self, content: str = "", meta: AssistantMessageMeta | None = None) -> None:
         """Start a new assistant message."""
@@ -67,7 +52,10 @@ class Runner:
         """Ensure current assistant message exists and return it."""
         if self._current_assistant_message is None:
             self._start_assistant_message()
-        return self._current_assistant_message  # type: ignore[return-value]
+        if self._current_assistant_message is None:
+            msg = "Failed to create current assistant message"
+            raise RuntimeError(msg)
+        return self._current_assistant_message
 
     def _add_to_current_assistant_message(self, content_item: AssistantTextContent | AssistantToolCall | AssistantToolCallResult) -> None:
         """Add content to the current assistant message."""
@@ -100,15 +88,20 @@ class Runner:
 
         if self.messages and isinstance(self.messages[-1], NewAssistantMessage):
             # Add to existing assistant message
-            self.messages[-1].content.append(result)
+            last_message = cast("NewAssistantMessage", self.messages[-1])
+            last_message.content.append(result)
         else:
             # Create new assistant message with just the tool result
             assistant_message = NewAssistantMessage(content=[result])
             self.messages.append(assistant_message)
 
+        # For completion API compatibility, create a separate assistant message
+        # Note: In the new architecture, we store everything as NewMessage format
+        # The conversion to completion format happens when sending to LLM
+
     def _normalize_includes(self, includes: Sequence[AgentChunkType] | None) -> Sequence[AgentChunkType]:
         """Normalize includes parameter to default if None."""
-        return includes if includes is not None else DEFAULT_INCLUDES
+        return includes if includes is not None else StreamIncludes.DEFAULT_INCLUDES
 
     def _normalize_record_path(self, record_to: PathLike | str | None) -> Path | None:
         """Normalize record_to parameter to Path object if provided."""
@@ -120,34 +113,68 @@ class Runner:
             return
 
         # Check for transfer_to_agent calls first
-        transfer_calls = [tc for tc in tool_calls if tc.function.name == "transfer_to_agent"]
+        transfer_calls = [tc for tc in tool_calls if tc.function.name == ToolName.TRANSFER_TO_AGENT]
         if transfer_calls:
             # Handle all transfer calls but only execute the first one
             for i, tool_call in enumerate(transfer_calls):
                 if i == 0:
                     # Execute the first transfer
-                    await self._handle_agent_transfer(tool_call)
+                    call_id, output = await self._handle_agent_transfer(tool_call)
+                    # Generate function_call_output event if in includes
+                    if "function_call_output" in includes:
+                        yield FunctionCallOutputEvent(
+                            tool_call_id=call_id,
+                            name=tool_call.function.name,
+                            content=output,
+                            execution_time_ms=0,  # Transfer operations are typically fast
+                        )
                 else:
                     # Add response for additional transfer calls without executing them
+                    output = "Transfer already executed by previous call"
                     self._add_tool_call_result(
                         call_id=tool_call.id,
-                        output="Transfer already executed by previous call",
+                        output=output,
                     )
+                    # Generate function_call_output event if in includes
+                    if "function_call_output" in includes:
+                        yield FunctionCallOutputEvent(
+                            tool_call_id=tool_call.id,
+                            name=tool_call.function.name,
+                            content=output,
+                            execution_time_ms=0,
+                        )
             return  # Stop processing other tool calls after transfer
 
-        return_parent_calls = [tc for tc in tool_calls if tc.function.name == "transfer_to_parent"]
+        return_parent_calls = [tc for tc in tool_calls if tc.function.name == ToolName.TRANSFER_TO_PARENT]
         if return_parent_calls:
             # Handle multiple transfer_to_parent calls (only execute the first one)
             for i, tool_call in enumerate(return_parent_calls):
                 if i == 0:
                     # Execute the first transfer
-                    await self._handle_parent_transfer(tool_call)
+                    call_id, output = await self._handle_parent_transfer(tool_call)
+                    # Generate function_call_output event if in includes
+                    if "function_call_output" in includes:
+                        yield FunctionCallOutputEvent(
+                            tool_call_id=call_id,
+                            name=tool_call.function.name,
+                            content=output,
+                            execution_time_ms=0,  # Transfer operations are typically fast
+                        )
                 else:
                     # Add response for additional transfer calls without executing them
+                    output = "Transfer already executed by previous call"
                     self._add_tool_call_result(
                         call_id=tool_call.id,
-                        output="Transfer already executed by previous call",
+                        output=output,
                     )
+                    # Generate function_call_output event if in includes
+                    if "function_call_output" in includes:
+                        yield FunctionCallOutputEvent(
+                            tool_call_id=tool_call.id,
+                            name=tool_call.function.name,
+                            content=output,
+                            execution_time_ms=0,
+                        )
             return  # Stop processing other tool calls after transfer
 
         async for tool_call_chunk in self.agent.handle_tool_calls(tool_calls, context=context):
@@ -163,7 +190,10 @@ class Runner:
                         output=tool_call_chunk.content,
                         execution_time_ms=tool_call_chunk.execution_time_ms,
                     )
-                    self.messages[-1].content.append(tool_result)
+                    last_message = cast("NewAssistantMessage", self.messages[-1])
+                    last_message.content.append(tool_result)
+
+                # Note: For completion API compatibility, the conversion happens when sending to LLM
 
     async def _collect_all_chunks(self, stream: AsyncGenerator[AgentChunk, None]) -> list[AgentChunk]:
         """Collect all chunks from an async generator into a list."""
@@ -171,16 +201,35 @@ class Runner:
 
     def run(
         self,
-        user_input: UserInput,
+        user_input: UserInput | None = None,
         max_steps: int = 20,
         includes: Sequence[AgentChunkType] | None = None,
         context: "Any | None" = None,  # noqa: ANN401
         record_to: PathLike | str | None = None,
         agent_kwargs: dict[str, Any] | None = None,
     ) -> AsyncGenerator[AgentChunk, None]:
-        """Run the agent and return a RunResponse object that can be asynchronously iterated for each chunk."""
+        """Run the agent and return a RunResponse object that can be asynchronously iterated for each chunk.
+
+        If user_input is None, the method will continue execution from the current state,
+        equivalent to calling the continue methods.
+        """
         logger.debug(f"Runner.run called with streaming={self.streaming}, api={self.api}")
         includes = self._normalize_includes(includes)
+
+        # If no user input provided, use continue logic
+        if user_input is None:
+            logger.debug("No user input provided, using continue logic")
+            return self._run_continue_stream(max_steps, includes, self._normalize_record_path(record_to), context)
+
+        # Cancel any pending tool calls before processing new user input
+        # and yield cancellation events if they should be included
+        cancellation_events = self._cancel_pending_tool_calls()
+
+        # We need to handle this differently since run() is not async
+        # Store cancellation events to be yielded by _run
+        self._pending_cancellation_events = cancellation_events
+
+        # Process user input
         match user_input:
             case str():
                 self.messages.append(NewUserMessage(content=[UserTextContent(text=user_input)]))
@@ -204,21 +253,31 @@ class Runner:
     ) -> AsyncGenerator[AgentChunk, None]:
         """Run the agent and return a RunResponse object that can be asynchronously iterated for each chunk."""
         logger.debug(f"Running agent with messages: {self.messages}")
+
+        # First, yield any pending cancellation events
+        if hasattr(self, "_pending_cancellation_events"):
+            for cancellation_event in self._pending_cancellation_events:
+                if "function_call_output" in includes:
+                    yield cancellation_event
+            # Clear the pending events after yielding
+            delattr(self, "_pending_cancellation_events")
+
         steps = 0
         finish_reason = None
 
         # Determine completion condition based on agent configuration
-        completion_condition = getattr(self.agent, "completion_condition", "stop")
+        completion_condition = getattr(self.agent, "completion_condition", CompletionMode.STOP)
 
         def is_finish() -> bool:
-            if completion_condition == "call":
+            if completion_condition == CompletionMode.CALL:
                 # Check if wait_for_user was called in the last assistant message
                 if self.messages and isinstance(self.messages[-1], NewAssistantMessage):
-                    for content_item in self.messages[-1].content:
-                        if content_item.type == "tool_call_result" and self._get_tool_call_name_by_id(content_item.call_id) == "wait_for_user":
+                    last_message = self.messages[-1]
+                    for content_item in last_message.content:
+                        if isinstance(content_item, AssistantToolCallResult) and self._get_tool_call_name_by_id(content_item.call_id) == ToolName.WAIT_FOR_USER:
                             return True
                 return False
-            return finish_reason == "stop"
+            return finish_reason == CompletionMode.STOP
 
         while not is_finish() and steps < max_steps:
             logger.debug(f"Step {steps}: finish_reason={finish_reason}, is_finish()={is_finish()}")
@@ -255,23 +314,26 @@ class Runner:
                 match chunk.type:
                     case "assistant_message":
                         # Start or update assistant message in new format
-                        meta = AssistantMessageMeta(
-                            sent_at=chunk.message.meta.sent_at,
-                            latency_ms=getattr(chunk.message.meta, "latency_ms", None),
-                            total_time_ms=getattr(chunk.message.meta, "output_time_ms", None),
-                        )
                         # If we already have a current assistant message, just update its metadata
                         if self._current_assistant_message is not None:
-                            self._current_assistant_message.meta = meta
+                            # Preserve all existing metadata and only update specific fields
+                            original_meta = self._current_assistant_message.meta
+                            original_meta.sent_at = chunk.message.meta.sent_at
+                            if hasattr(chunk.message.meta, "latency_ms"):
+                                original_meta.latency_ms = chunk.message.meta.latency_ms
+                            if hasattr(chunk.message.meta, "output_time_ms"):
+                                original_meta.total_time_ms = chunk.message.meta.output_time_ms
+                            # Preserve other metadata fields like model, usage, etc.
+                            for attr in ["model", "usage", "input_tokens", "output_tokens"]:
+                                if hasattr(chunk.message.meta, attr):
+                                    setattr(original_meta, attr, getattr(chunk.message.meta, attr))
                         else:
-                            # Extract text content from the new message format
-                            text_content = ""
-                            if chunk.message.content:
-                                for item in chunk.message.content:
-                                    if hasattr(item, "type") and item.type == "text":
-                                        text_content = item.text
-                                        break
-                            self._start_assistant_message(text_content, meta)
+                            # For non-streaming mode, directly use the complete message from the response handler
+                            self._current_assistant_message = chunk.message
+
+                        # If model is None, try to get it from agent client
+                        if self._current_assistant_message is not None and self._current_assistant_message.meta.model is None and hasattr(self.agent.client, "model"):
+                            self._current_assistant_message.meta.model = self.agent.client.model
                         # Only yield assistant_message chunk if it's in includes and has content
                         if chunk.type in includes and self._current_assistant_message is not None:
                             # Create a new chunk with the current assistant message content
@@ -298,28 +360,45 @@ class Runner:
                         if chunk.type in includes:
                             yield chunk
                     case "usage":
-                        # Update the last assistant message with usage data and output_time_ms
+                        # Update the current or last assistant message with usage data and output_time_ms
                         usage_time = datetime.now(timezone.utc)
-                        for i in range(len(self.messages) - 1, -1, -1):
-                            current_message = self.messages[i]
-                            if isinstance(current_message, NewAssistantMessage):
-                                # Update usage information
-                                if current_message.meta.usage is None:
-                                    current_message.meta.usage = MessageUsage()
-                                current_message.meta.usage.input_tokens = chunk.usage.input_tokens
-                                current_message.meta.usage.output_tokens = chunk.usage.output_tokens
-                                current_message.meta.usage.total_tokens = (chunk.usage.input_tokens or 0) + (chunk.usage.output_tokens or 0)
-
-                                # Calculate output_time_ms if latency_ms is available
-                                if current_message.meta.latency_ms is not None:
-                                    # We need to calculate from first output to usage time
-                                    # We'll calculate: usage_time - (sent_at - latency_ms)
-                                    # This gives us the time from first output to usage completion
-                                    # sent_at is when the message was completed, so sent_at - latency_ms approximates first output time
-                                    first_output_time_approx = current_message.meta.sent_at - timedelta(milliseconds=current_message.meta.latency_ms)
-                                    output_time_ms = int((usage_time - first_output_time_approx).total_seconds() * 1000)
-                                    current_message.meta.total_time_ms = max(0, output_time_ms)
-                                break
+
+                        # Always accumulate usage in runner first
+                        self.usage.input_tokens = (self.usage.input_tokens or 0) + (chunk.usage.input_tokens or 0)
+                        self.usage.output_tokens = (self.usage.output_tokens or 0) + (chunk.usage.output_tokens or 0)
+                        self.usage.total_tokens = (self.usage.total_tokens or 0) + (chunk.usage.input_tokens or 0) + (chunk.usage.output_tokens or 0)
+
+                        # Try to find the assistant message to update
+                        target_message = None
+
+                        # First check if we have a current assistant message
+                        if self._current_assistant_message is not None:
+                            target_message = self._current_assistant_message
+                        else:
+                            # Otherwise, look for the last assistant message in the list
+                            for i in range(len(self.messages) - 1, -1, -1):
+                                current_message = self.messages[i]
+                                if isinstance(current_message, NewAssistantMessage):
+                                    target_message = current_message
+                                    break
+
+                        # Update the target message with usage information
+                        if target_message is not None:
+                            if target_message.meta.usage is None:
+                                target_message.meta.usage = MessageUsage()
+                            target_message.meta.usage.input_tokens = chunk.usage.input_tokens
+                            target_message.meta.usage.output_tokens = chunk.usage.output_tokens
+                            target_message.meta.usage.total_tokens = (chunk.usage.input_tokens or 0) + (chunk.usage.output_tokens or 0)
+
+                            # Calculate output_time_ms if latency_ms is available
+                            if target_message.meta.latency_ms is not None:
+                                # We need to calculate from first output to usage time
+                                # We'll calculate: usage_time - (sent_at - latency_ms)
+                                # This gives us the time from first output to usage completion
+                                # sent_at is when the message was completed, so sent_at - latency_ms approximates first output time
+                                first_output_time_approx = target_message.meta.sent_at - timedelta(milliseconds=target_message.meta.latency_ms)
+                                output_time_ms = int((usage_time - first_output_time_approx).total_seconds() * 1000)
+                                target_message.meta.total_time_ms = max(0, output_time_ms)
                         # Always yield usage chunk if it's in includes
                         if chunk.type in includes:
                             yield chunk
@@ -342,7 +421,7 @@ class Runner:
                     yield tool_chunk
                 finish_reason = "tool_calls"
             else:
-                finish_reason = "stop"
+                finish_reason = CompletionMode.STOP
             steps += 1
 
     async def has_require_confirm_tools(self):
@@ -359,6 +438,12 @@ class Runner:
         includes: list[AgentChunkType] | None = None,
         record_to: PathLike | str | None = None,
     ) -> list[AgentChunk]:
+        """Deprecated: Use run_until_complete(None) instead."""
+        warnings.warn(
+            "run_continue_until_complete is deprecated. Use run_until_complete(None) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         resp = self.run_continue_stream(max_steps, includes, record_to=record_to)
         return await self._collect_all_chunks(resp)
 
@@ -369,6 +454,12 @@ class Runner:
         record_to: PathLike | str | None = None,
         context: "Any | None" = None,  # noqa: ANN401
     ) -> AsyncGenerator[AgentChunk, None]:
+        """Deprecated: Use run(None) instead."""
+        warnings.warn(
+            "run_continue_stream is deprecated. Use run(None) instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         return self._run_continue_stream(max_steps, includes, record_to=record_to, context=context)
 
     async def _run_continue_stream(
@@ -403,7 +494,7 @@ class Runner:
 
     async def run_until_complete(
         self,
-        user_input: UserInput,
+        user_input: UserInput | None = None,
         max_steps: int = 20,
         includes: list[AgentChunkType] | None = None,
         record_to: PathLike | str | None = None,
@@ -421,11 +512,12 @@ class Runner:
         tool_results = set()
         tool_call_names = {}
 
-        for content_item in self.messages[-1].content:
-            if content_item.type == "tool_call":
+        last_message = self.messages[-1]
+        for content_item in last_message.content:
+            if isinstance(content_item, AssistantToolCall):
                 tool_calls[content_item.call_id] = content_item
                 tool_call_names[content_item.call_id] = content_item.name
-            elif content_item.type == "tool_call_result":
+            elif isinstance(content_item, AssistantToolCallResult):
                 tool_results.add(content_item.call_id)
 
         # Return pending tool calls and tool call names map
@@ -442,6 +534,38 @@ class Runner:
         _, tool_call_names = self._analyze_last_assistant_message()
         return tool_call_names.get(call_id)
 
+    def _cancel_pending_tool_calls(self) -> list[FunctionCallOutputEvent]:
+        """Cancel all pending tool calls by adding cancellation results.
+
+        Returns:
+            List of FunctionCallOutputEvent for each cancelled tool call
+        """
+        pending_tool_calls = self._find_pending_tool_calls()
+        if not pending_tool_calls:
+            return []
+
+        logger.debug(f"Cancelling {len(pending_tool_calls)} pending tool calls due to new user input")
+
+        cancellation_events = []
+        for tool_call in pending_tool_calls:
+            output = "Operation cancelled by user - new input provided"
+            self._add_tool_call_result(
+                call_id=tool_call.call_id,
+                output=output,
+                execution_time_ms=0,
+            )
+
+            # Create cancellation event
+            cancellation_event = FunctionCallOutputEvent(
+                tool_call_id=tool_call.call_id,
+                name=tool_call.name,
+                content=output,
+                execution_time_ms=0,
+            )
+            cancellation_events.append(cancellation_event)
+
+        return cancellation_events
+
     def _convert_tool_calls_to_tool_calls(self, tool_calls: list[AssistantToolCall]) -> list[ToolCall]:
         """Convert AssistantToolCall objects to ToolCall objects for compatibility."""
         return [
@@ -457,7 +581,7 @@ class Runner:
             for i, tc in enumerate(tool_calls)
         ]
 
-    def set_chat_history(self, messages: Sequence[FlexibleRunnerMessage], root_agent: Agent | None = None) -> None:
+    def set_chat_history(self, messages: Sequence[FlexibleInputMessage], root_agent: Agent | None = None) -> None:
         """Set the entire chat history and track the current agent based on function calls.
 
         This method analyzes the message history to determine which agent should be active
@@ -474,17 +598,54 @@ class Runner:
         current_agent = root_agent if root_agent is not None else self.agent
 
         # Add each message and track agent transfers
-        for message in messages:
-            self.append_message(message)
-            current_agent = self._track_agent_transfer_in_message(message, current_agent)
+        for input_message in messages:
+            # Store length before adding to get the added message
+            prev_length = len(self.messages)
+            self.append_message(input_message)
+
+            # Track transfers using the converted message (now in self.messages)
+            if len(self.messages) > prev_length:
+                converted_message = self.messages[-1]  # Get the last added message
+                current_agent = self._track_agent_transfer_in_message(converted_message, current_agent)
 
         # Set the current agent based on the tracked transfers
         self.agent = current_agent
         logger.info(f"Chat history set with {len(self.messages)} messages. Current agent: {self.agent.name}")
 
-    def get_messages_dict(self) -> list[dict[str, Any]]:
+    def get_messages(self) -> list[NewMessage]:
+        """Get the messages as NewMessage objects.
+
+        Only returns NewMessage objects, filtering out any dict or other legacy formats.
+        """
+        return [msg for msg in self.messages if isinstance(msg, NewMessage)]
+
+    def get_dict_messages(self) -> list[dict[str, Any]]:
         """Get the messages in JSONL format."""
-        return [msg.model_dump(mode="json") for msg in self.messages]
+        result = []
+        for msg in self.messages:
+            if hasattr(msg, "model_dump"):
+                result.append(msg.model_dump(mode="json"))
+            elif isinstance(msg, dict):
+                result.append(msg)
+            else:
+                # Fallback for any other message types
+                result.append(dict(msg))
+        return result
+
+    def add_user_message(self, text: str) -> None:
+        """Convenience method to add a user text message."""
+        message = NewUserMessage(content=[UserTextContent(text=text)])
+        self.append_message(message)
+
+    def add_assistant_message(self, text: str) -> None:
+        """Convenience method to add an assistant text message."""
+        message = NewAssistantMessage(content=[AssistantTextContent(text=text)])
+        self.append_message(message)
+
+    def add_system_message(self, content: str) -> None:
+        """Convenience method to add a system message."""
+        message = NewSystemMessage(content=content)
+        self.append_message(message)
 
     def _track_agent_transfer_in_message(self, message: FlexibleRunnerMessage, current_agent: Agent) -> Agent:
         """Track agent transfers in a single message.
@@ -496,8 +657,6 @@ class Runner:
         Returns:
             The agent that should be active after processing this message
         """
-        if isinstance(message, dict):
-            return self._track_transfer_from_dict_message(message, current_agent)
         if isinstance(message, NewAssistantMessage):
             return self._track_transfer_from_new_assistant_message(message, current_agent)
 
@@ -507,28 +666,13 @@ class Runner:
         """Track transfers from NewAssistantMessage objects."""
         for content_item in message.content:
             if content_item.type == "tool_call":
-                if content_item.name == "transfer_to_agent":
+                if content_item.name == ToolName.TRANSFER_TO_AGENT:
                     arguments = content_item.arguments if isinstance(content_item.arguments, str) else str(content_item.arguments)
                     return self._handle_transfer_to_agent_tracking(arguments, current_agent)
-                if content_item.name == "transfer_to_parent":
+                if content_item.name == ToolName.TRANSFER_TO_PARENT:
                     return self._handle_transfer_to_parent_tracking(current_agent)
         return current_agent
 
-    def _track_transfer_from_dict_message(self, message: dict[str, Any] | MessageDict, current_agent: Agent) -> Agent:
-        """Track transfers from dictionary-format messages."""
-        message_type = message.get("type")
-        if message_type != "function_call":
-            return current_agent
-
-        function_name = message.get("name", "")
-        if function_name == "transfer_to_agent":
-            return self._handle_transfer_to_agent_tracking(message.get("arguments", ""), current_agent)
-
-        if function_name == "transfer_to_parent":
-            return self._handle_transfer_to_parent_tracking(current_agent)
-
-        return current_agent
-
     def _handle_transfer_to_agent_tracking(self, arguments: str | dict, current_agent: Agent) -> Agent:
         """Handle transfer_to_agent function call tracking."""
         try:
@@ -584,145 +728,39 @@ class Runner:
 
         return None
 
-    def append_message(self, message: FlexibleRunnerMessage) -> None:
+    def append_message(self, message: FlexibleInputMessage) -> None:
+        """Append a message to the conversation history.
+
+        Accepts both NewMessage format and dict format (which will be converted internally).
+        """
         if isinstance(message, NewMessage):
-            # Already in new format
             self.messages.append(message)
         elif isinstance(message, dict):
-            # Handle different message types from dict
-            message_type = message.get("type")
-            role = message.get("role")
-
+            # Convert dict to NewMessage using MessageBuilder
+            role = message.get("role", "").lower()
             if role == "user":
-                content = message.get("content", "")
-                if isinstance(content, str):
-                    user_message = NewUserMessage(content=[UserTextContent(text=content)])
-                elif isinstance(content, list):
-                    # Handle complex content array
-                    user_content_items: list[UserMessageContent] = []
-                    for item in content:
-                        if isinstance(item, dict):
-                            item_type = item.get("type")
-                            if item_type in {"input_text", "text"}:
-                                user_content_items.append(UserTextContent(text=item.get("text", "")))
-                            elif item_type in {"input_image", "image_url"}:
-                                if item_type == "image_url":
-                                    # Handle completion API format
-                                    image_url = item.get("image_url", {})
-                                    url = image_url.get("url", "") if isinstance(image_url, dict) else str(image_url)
-                                    user_content_items.append(UserImageContent(image_url=url))
-                                else:
-                                    # Handle response API format
-                                    user_content_items.append(
-                                        UserImageContent(
-                                            image_url=item.get("image_url"),
-                                            file_id=item.get("file_id"),
-                                            detail=item.get("detail", "auto"),
-                                        ),
-                                    )
-                        elif hasattr(item, "type"):
-                            # Handle Pydantic models
-                            if item.type == "input_text":
-                                user_content_items.append(UserTextContent(text=item.text))
-                            elif item.type == "input_image":
-                                user_content_items.append(
-                                    UserImageContent(
-                                        image_url=getattr(item, "image_url", None),
-                                        file_id=getattr(item, "file_id", None),
-                                        detail=getattr(item, "detail", "auto"),
-                                    ),
-                                )
-                        else:
-                            # Fallback: convert to text
-                            user_content_items.append(UserTextContent(text=str(item)))
-
-                    user_message = NewUserMessage(content=user_content_items)
-                else:
-                    # Handle non-string, non-list content
-                    user_message = NewUserMessage(content=[UserTextContent(text=str(content))])
-                self.messages.append(user_message)
-            elif role == "system":
-                content = message.get("content", "")
-                system_message = NewSystemMessage(content=str(content))
-                self.messages.append(system_message)
+                converted_message = MessageBuilder.build_user_message_from_dict(message)
             elif role == "assistant":
-                content = message.get("content", "")
-                assistant_content_items: list[AssistantMessageContent] = [AssistantTextContent(text=str(content))] if content else []
-
-                # Handle tool calls if present
-                if "tool_calls" in message:
-                    for tool_call in message.get("tool_calls", []):
-                        try:
-                            arguments = json.loads(tool_call["function"]["arguments"]) if isinstance(tool_call["function"]["arguments"], str) else tool_call["function"]["arguments"]
-                        except (json.JSONDecodeError, TypeError):
-                            arguments = tool_call["function"]["arguments"]
-
-                        assistant_content_items.append(
-                            AssistantToolCall(
-                                call_id=tool_call["id"],
-                                name=tool_call["function"]["name"],
-                                arguments=arguments,
-                            ),
-                        )
-
-                assistant_message = NewAssistantMessage(content=assistant_content_items)
-                self.messages.append(assistant_message)
-            elif message_type == "function_call":
-                # Handle function_call directly like AgentFunctionToolCallMessage
-                # Type guard: ensure we have the right message type
-                if "call_id" in message and "name" in message and "arguments" in message:
-                    function_call_msg = message  # Type should be FunctionCallDict now
-                    if self.messages and isinstance(self.messages[-1], NewAssistantMessage):
-                        tool_call = AssistantToolCall(
-                            call_id=function_call_msg["call_id"],  # type: ignore
-                            name=function_call_msg["name"],  # type: ignore
-                            arguments=function_call_msg["arguments"],  # type: ignore
-                        )
-                        self.messages[-1].content.append(tool_call)
-                    else:
-                        assistant_message = NewAssistantMessage(
-                            content=[
-                                AssistantToolCall(
-                                    call_id=function_call_msg["call_id"],  # type: ignore
-                                    name=function_call_msg["name"],  # type: ignore
-                                    arguments=function_call_msg["arguments"],  # type: ignore
-                                ),
-                            ],
-                        )
-                        self.messages.append(assistant_message)
-            elif message_type == "function_call_output":
-                # Handle function_call_output directly like AgentFunctionCallOutput
-                # Type guard: ensure we have the right message type
-                if "call_id" in message and "output" in message:
-                    function_output_msg = message  # Type should be FunctionCallOutputDict now
-                    if self.messages and isinstance(self.messages[-1], NewAssistantMessage):
-                        tool_result = AssistantToolCallResult(
-                            call_id=function_output_msg["call_id"],  # type: ignore
-                            output=function_output_msg["output"],  # type: ignore
-                        )
-                        self.messages[-1].content.append(tool_result)
-                    else:
-                        assistant_message = NewAssistantMessage(
-                            content=[
-                                AssistantToolCallResult(
-                                    call_id=function_output_msg["call_id"],  # type: ignore
-                                    output=function_output_msg["output"],  # type: ignore
-                                ),
-                            ],
-                        )
-                        self.messages.append(assistant_message)
+                converted_message = MessageBuilder.build_assistant_message_from_dict(message)
+            elif role == "system":
+                converted_message = MessageBuilder.build_system_message_from_dict(message)
             else:
-                msg = "Message must have a 'role' or 'type' field."
+                msg = f"Unsupported message role: {role}. Must be 'user', 'assistant', or 'system'."
                 raise ValueError(msg)
+
+            self.messages.append(converted_message)
         else:
-            msg = f"Unsupported message type: {type(message)}"
+            msg = f"Unsupported message type: {type(message)}. Supports NewMessage types and dict."
             raise TypeError(msg)
 
-    async def _handle_agent_transfer(self, tool_call: ToolCall) -> None:
+    async def _handle_agent_transfer(self, tool_call: ToolCall) -> tuple[str, str]:
         """Handle agent transfer when transfer_to_agent tool is called.
 
         Args:
             tool_call: The transfer_to_agent tool call
+
+        Returns:
+            Tuple of (call_id, output) for the tool call result
         """
 
         # Parse the arguments to get the target agent name
@@ -731,31 +769,34 @@ class Runner:
             target_agent_name = arguments.get("name")
         except (json.JSONDecodeError, KeyError):
             logger.error("Failed to parse transfer_to_agent arguments: %s", tool_call.function.arguments)
+            output = "Failed to parse transfer arguments"
             # Add error result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output="Failed to parse transfer arguments",
+                output=output,
             )
-            return
+            return tool_call.id, output
 
         if not target_agent_name:
             logger.error("No target agent name provided in transfer_to_agent call")
+            output = "No target agent name provided"
             # Add error result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output="No target agent name provided",
+                output=output,
             )
-            return
+            return tool_call.id, output
 
         # Find the target agent in handoffs
         if not self.agent.handoffs:
             logger.error("Current agent has no handoffs configured")
+            output = "Current agent has no handoffs configured"
             # Add error result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output="Current agent has no handoffs configured",
+                output=output,
             )
-            return
+            return tool_call.id, output
 
         target_agent = None
         for agent in self.agent.handoffs:
@@ -765,12 +806,13 @@ class Runner:
 
         if not target_agent:
             logger.error("Target agent '%s' not found in handoffs", target_agent_name)
+            output = f"Target agent '{target_agent_name}' not found in handoffs"
             # Add error result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output=f"Target agent '{target_agent_name}' not found in handoffs",
+                output=output,
             )
-            return
+            return tool_call.id, output
 
         # Execute the transfer tool call to get the result
         try:
@@ -779,40 +821,49 @@ class Runner:
                 tool_call.function.arguments or "",
             )
 
+            output = str(result)
             # Add the tool call result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output=str(result),
+                output=output,
             )
 
             # Switch to the target agent
             logger.info("Transferring conversation from %s to %s", self.agent.name, target_agent_name)
             self.agent = target_agent
 
+            return tool_call.id, output
+
         except Exception as e:
             logger.exception("Failed to execute transfer_to_agent tool call")
+            output = f"Transfer failed: {e!s}"
             # Add error result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output=f"Transfer failed: {e!s}",
+                output=output,
             )
+            return tool_call.id, output
 
-    async def _handle_parent_transfer(self, tool_call: ToolCall) -> None:
+    async def _handle_parent_transfer(self, tool_call: ToolCall) -> tuple[str, str]:
         """Handle parent transfer when transfer_to_parent tool is called.
 
         Args:
             tool_call: The transfer_to_parent tool call
+
+        Returns:
+            Tuple of (call_id, output) for the tool call result
         """
 
         # Check if current agent has a parent
         if not self.agent.parent:
             logger.error("Current agent has no parent to transfer back to.")
+            output = "Current agent has no parent to transfer back to"
             # Add error result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output="Current agent has no parent to transfer back to",
+                output=output,
             )
-            return
+            return tool_call.id, output
 
         # Execute the transfer tool call to get the result
         try:
@@ -821,20 +872,25 @@ class Runner:
                 tool_call.function.arguments or "",
             )
 
+            output = str(result)
             # Add the tool call result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output=str(result),
+                output=output,
             )
 
             # Switch to the parent agent
             logger.info("Transferring conversation from %s back to parent %s", self.agent.name, self.agent.parent.name)
             self.agent = self.agent.parent
 
+            return tool_call.id, output
+
         except Exception as e:
             logger.exception("Failed to execute transfer_to_parent tool call")
+            output = f"Transfer to parent failed: {e!s}"
             # Add error result to messages
             self._add_tool_call_result(
                 call_id=tool_call.id,
-                output=f"Transfer to parent failed: {e!s}",
+                output=output,
             )
+            return tool_call.id, output