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

lite_agent/agent.py

@@ -5,12 +5,23 @@ from typing import Any, Optional
 
 from funcall import Funcall
 from jinja2 import Environment, FileSystemLoader
-from litellm import CustomStreamWrapper
 
 from lite_agent.client import BaseLLMClient, LiteLLMClient, ReasoningConfig
+from lite_agent.constants import CompletionMode, ToolName
 from lite_agent.loggers import logger
-from lite_agent.stream_handlers import litellm_completion_stream_handler, litellm_response_stream_handler
-from lite_agent.types import AgentChunk, FunctionCallEvent, FunctionCallOutputEvent, RunnerMessages, ToolCall, message_to_llm_dict, system_message_to_llm_dict
+from lite_agent.response_handlers import CompletionResponseHandler, ResponsesAPIHandler
+from lite_agent.types import (
+    AgentChunk,
+    AssistantTextContent,
+    AssistantToolCall,
+    AssistantToolCallResult,
+    FunctionCallEvent,
+    FunctionCallOutputEvent,
+    RunnerMessages,
+    ToolCall,
+    message_to_llm_dict,
+    system_message_to_llm_dict,
+)
 from lite_agent.types.messages import NewAssistantMessage, NewSystemMessage, NewUserMessage
 
 TEMPLATES_DIR = Path(__file__).parent / "templates"
@@ -22,7 +33,7 @@ WAIT_FOR_USER_INSTRUCTIONS_TEMPLATE = jinja_env.get_template("wait_for_user_inst
 
 
 class Agent:
-    def __init__(  # noqa: PLR0913
+    def __init__(
         self,
         *,
         model: str | BaseLLMClient,
@@ -33,10 +44,24 @@ class Agent:
         message_transfer: Callable[[RunnerMessages], RunnerMessages] | None = None,
         completion_condition: str = "stop",
         reasoning: ReasoningConfig = None,
+        stop_before_tools: list[str] | list[Callable] | None = None,
     ) -> None:
         self.name = name
         self.instructions = instructions
         self.reasoning = reasoning
+        # Convert stop_before_functions to function names
+        if stop_before_tools:
+            self.stop_before_functions = set()
+            for func in stop_before_tools:
+                if isinstance(func, str):
+                    self.stop_before_functions.add(func)
+                elif callable(func):
+                    self.stop_before_functions.add(func.__name__)
+                else:
+                    msg = f"stop_before_functions must contain strings or callables, got {type(func)}"
+                    raise TypeError(msg)
+        else:
+            self.stop_before_functions = set()
 
         if isinstance(model, BaseLLMClient):
             # If model is a BaseLLMClient instance, use it directly
@@ -55,7 +80,7 @@ class Agent:
         self.fc = Funcall(tools)
 
         # Add wait_for_user tool if completion condition is "call"
-        if completion_condition == "call":
+        if completion_condition == CompletionMode.CALL:
             self._add_wait_for_user_tool()
 
         # Set parent for handoff agents
@@ -100,7 +125,7 @@ class Agent:
 
         # Add single dynamic tool for all transfers
         self.fc.add_dynamic_tool(
-            name="transfer_to_agent",
+            name=ToolName.TRANSFER_TO_AGENT,
             description="Transfer conversation to another agent.",
             parameters={
                 "name": {
@@ -130,7 +155,7 @@ class Agent:
 
         # Add dynamic tool for parent transfer
         self.fc.add_dynamic_tool(
-            name="transfer_to_parent",
+            name=ToolName.TRANSFER_TO_PARENT,
             description="Transfer conversation back to parent agent when current task is completed or cannot be solved by current agent",
             parameters={},
             required=[],
@@ -161,7 +186,7 @@ class Agent:
             try:
                 # Try to remove the existing transfer tool
                 if hasattr(self.fc, "remove_dynamic_tool"):
-                    self.fc.remove_dynamic_tool("transfer_to_agent")
+                    self.fc.remove_dynamic_tool(ToolName.TRANSFER_TO_AGENT)
             except Exception as e:
                 # If removal fails, log and continue anyway
                 logger.debug(f"Failed to remove existing transfer tool: {e}")
@@ -206,31 +231,30 @@ class Agent:
         for message in messages:
             if isinstance(message, NewAssistantMessage):
                 for item in message.content:
-                    match item.type:
-                        case "text":
-                            res.append(
-                                {
-                                    "role": "assistant",
-                                    "content": item.text,
-                                },
-                            )
-                        case "tool_call":
-                            res.append(
-                                {
-                                    "type": "function_call",
-                                    "call_id": item.call_id,
-                                    "name": item.name,
-                                    "arguments": item.arguments,
-                                },
-                            )
-                        case "tool_call_result":
-                            res.append(
-                                {
-                                    "type": "function_call_output",
-                                    "call_id": item.call_id,
-                                    "output": item.output,
-                                },
-                            )
+                    if isinstance(item, AssistantTextContent):
+                        res.append(
+                            {
+                                "role": "assistant",
+                                "content": item.text,
+                            },
+                        )
+                    elif isinstance(item, AssistantToolCall):
+                        res.append(
+                            {
+                                "type": "function_call",
+                                "call_id": item.call_id,
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        )
+                    elif isinstance(item, AssistantToolCallResult):
+                        res.append(
+                            {
+                                "type": "function_call_output",
+                                "call_id": item.call_id,
+                                "output": item.output,
+                            },
+                        )
             elif isinstance(message, NewSystemMessage):
                 res.append(
                     {
@@ -270,9 +294,6 @@ class Agent:
                         "content": contents,
                     },
                 )
-            # Handle dict messages (legacy format)
-            elif isinstance(message, dict):
-                res.append(message)
         return res
 
     async def completion(
@@ -280,6 +301,8 @@ class Agent:
         messages: RunnerMessages,
         record_to_file: Path | None = None,
         reasoning: ReasoningConfig = None,
+        *,
+        streaming: bool = True,
     ) -> AsyncGenerator[AgentChunk, None]:
         # Apply message transfer callback if provided - always use legacy format for LLM compatibility
         processed_messages = messages
@@ -296,19 +319,20 @@ class Agent:
             tools=tools,
             tool_choice="auto",  # TODO: make this configurable
             reasoning=reasoning,
+            streaming=streaming,
         )
 
-        # Ensure resp is a CustomStreamWrapper
-        if isinstance(resp, CustomStreamWrapper):
-            return litellm_completion_stream_handler(resp, record_to=record_to_file)
-        msg = "Response is not a CustomStreamWrapper, cannot stream chunks."
-        raise TypeError(msg)
+        # Use response handler for unified processing
+        handler = CompletionResponseHandler()
+        return handler.handle(resp, streaming=streaming, record_to=record_to_file)
 
     async def responses(
         self,
         messages: RunnerMessages,
         record_to_file: Path | None = None,
         reasoning: ReasoningConfig = None,
+        *,
+        streaming: bool = True,
     ) -> AsyncGenerator[AgentChunk, None]:
         # Apply message transfer callback if provided - always use legacy format for LLM compatibility
         processed_messages = messages
@@ -324,21 +348,33 @@ class Agent:
             tools=tools,
             tool_choice="auto",  # TODO: make this configurable
             reasoning=reasoning,
+            streaming=streaming,
         )
-        return litellm_response_stream_handler(resp, record_to=record_to_file)
+        # Use response handler for unified processing
+        handler = ResponsesAPIHandler()
+        return handler.handle(resp, streaming=streaming, record_to=record_to_file)
 
     async def list_require_confirm_tools(self, tool_calls: Sequence[ToolCall] | None) -> Sequence[ToolCall]:
         if not tool_calls:
             return []
         results = []
         for tool_call in tool_calls:
-            tool_func = self.fc.function_registry.get(tool_call.function.name)
+            function_name = tool_call.function.name
+
+            # Check if function is in dynamic stop_before_functions list
+            if function_name in self.stop_before_functions:
+                logger.debug('Tool call "%s" requires confirmation (stop_before_functions)', tool_call.id)
+                results.append(tool_call)
+                continue
+
+            # Check decorator-based require_confirmation
+            tool_func = self.fc.function_registry.get(function_name)
             if not tool_func:
-                logger.warning("Tool function %s not found in registry", tool_call.function.name)
+                logger.warning("Tool function %s not found in registry", function_name)
                 continue
-            tool_meta = self.fc.get_tool_meta(tool_call.function.name)
+            tool_meta = self.fc.get_tool_meta(function_name)
             if tool_meta["require_confirm"]:
-                logger.debug('Tool call "%s" requires confirmation', tool_call.id)
+                logger.debug('Tool call "%s" requires confirmation (decorator)', tool_call.id)
                 results.append(tool_call)
         return results
 
@@ -393,10 +429,42 @@ class Agent:
             role = message_dict.get("role")
 
             if role == "assistant":
-                # Look ahead for function_call messages
+                # Extract tool_calls from content if present
                 tool_calls = []
+                content = message_dict.get("content", [])
+
+                # Handle both string and array content
+                if isinstance(content, list):
+                    # Extract tool_calls from content array and filter out non-text content
+                    filtered_content = []
+                    for item in content:
+                        if isinstance(item, dict):
+                            if item.get("type") == "tool_call":
+                                tool_call = {
+                                    "id": item.get("call_id", ""),
+                                    "type": "function",
+                                    "function": {
+                                        "name": item.get("name", ""),
+                                        "arguments": item.get("arguments", "{}"),
+                                    },
+                                    "index": len(tool_calls),
+                                }
+                                tool_calls.append(tool_call)
+                            elif item.get("type") == "text":
+                                filtered_content.append(item)
+                            # Skip tool_call_result - they should be handled by separate function_call_output messages
+
+                    # Update content to only include text items
+                    if filtered_content:
+                        message_dict = message_dict.copy()
+                        message_dict["content"] = filtered_content
+                    elif tool_calls:
+                        # If we have tool_calls but no text content, set content to None per OpenAI API spec
+                        message_dict = message_dict.copy()
+                        message_dict["content"] = None
+
+                # Look ahead for function_call messages (legacy support)
                 j = i + 1
-
                 while j < len(messages):
                     next_message = messages[j]
                     next_dict = message_to_llm_dict(next_message) if isinstance(next_message, (NewUserMessage, NewSystemMessage, NewAssistantMessage)) else next_message
@@ -421,6 +489,13 @@ class Agent:
                 if tool_calls:
                     assistant_msg["tool_calls"] = tool_calls  # type: ignore
 
+                # Convert content format for OpenAI API compatibility
+                content = assistant_msg.get("content", [])
+                if isinstance(content, list):
+                    # Extract text content and convert to string using list comprehension
+                    text_parts = [item.get("text", "") for item in content if isinstance(item, dict) and item.get("type") == "text"]
+                    assistant_msg["content"] = " ".join(text_parts) if text_parts else None
+
                 converted_messages.append(assistant_msg)
                 i = j  # Skip the function_call messages we've processed
 
@@ -533,9 +608,73 @@ class Agent:
 
         # Add dynamic tool for task completion
         self.fc.add_dynamic_tool(
-            name="wait_for_user",
+            name=ToolName.WAIT_FOR_USER,
             description="Call this function when you have completed your assigned task or need more information from the user.",
             parameters={},
             required=[],
             handler=wait_for_user_handler,
         )
+
+    def set_stop_before_functions(self, functions: list[str] | list[Callable]) -> None:
+        """Set the list of functions that require confirmation before execution.
+
+        Args:
+            functions: List of function names (str) or callable objects
+        """
+        self.stop_before_functions = set()
+        for func in functions:
+            if isinstance(func, str):
+                self.stop_before_functions.add(func)
+            elif callable(func):
+                self.stop_before_functions.add(func.__name__)
+            else:
+                msg = f"stop_before_functions must contain strings or callables, got {type(func)}"
+                raise TypeError(msg)
+        logger.debug(f"Set stop_before_functions to: {self.stop_before_functions}")
+
+    def add_stop_before_function(self, function: str | Callable) -> None:
+        """Add a function to the stop_before_functions list.
+
+        Args:
+            function: Function name (str) or callable object to add
+        """
+        if isinstance(function, str):
+            function_name = function
+        elif callable(function):
+            function_name = function.__name__
+        else:
+            msg = f"function must be a string or callable, got {type(function)}"
+            raise TypeError(msg)
+
+        self.stop_before_functions.add(function_name)
+        logger.debug(f"Added '{function_name}' to stop_before_functions")
+
+    def remove_stop_before_function(self, function: str | Callable) -> None:
+        """Remove a function from the stop_before_functions list.
+
+        Args:
+            function: Function name (str) or callable object to remove
+        """
+        if isinstance(function, str):
+            function_name = function
+        elif callable(function):
+            function_name = function.__name__
+        else:
+            msg = f"function must be a string or callable, got {type(function)}"
+            raise TypeError(msg)
+
+        self.stop_before_functions.discard(function_name)
+        logger.debug(f"Removed '{function_name}' from stop_before_functions")
+
+    def clear_stop_before_functions(self) -> None:
+        """Clear all function names from the stop_before_functions list."""
+        self.stop_before_functions.clear()
+        logger.debug("Cleared all stop_before_functions")
+
+    def get_stop_before_functions(self) -> set[str]:
+        """Get the current set of function names that require confirmation.
+
+        Returns:
+            Set of function names
+        """
+        return self.stop_before_functions.copy()

lite_agent/chat_display.py

@@ -26,6 +26,8 @@ from lite_agent.types import (
     AgentSystemMessage,
     AgentUserMessage,
     AssistantMessageMeta,
+    AssistantToolCall,
+    AssistantToolCallResult,
     BasicMessageMeta,
     FlexibleRunnerMessage,
     LLMResponseMeta,
@@ -228,9 +230,9 @@ def _update_message_counts(message: FlexibleRunnerMessage, counts: dict[str, int
         counts["Assistant"] += 1
         # Count tool calls and outputs within the assistant message
         for content_item in message.content:
-            if content_item.type == "tool_call":
+            if isinstance(content_item, AssistantToolCall):
                 counts["Function Call"] += 1
-            elif content_item.type == "tool_call_result":
+            elif isinstance(content_item, AssistantToolCallResult):
                 counts["Function Output"] += 1
     elif isinstance(message, NewSystemMessage):
         counts["System"] += 1
@@ -295,10 +297,18 @@ def _process_object_meta(meta: BasicMessageMeta | LLMResponseMeta | AssistantMes
     """处理对象类型的 meta 数据。"""
     # LLMResponseMeta 和 AssistantMessageMeta 都有这些字段
     if isinstance(meta, (LLMResponseMeta, AssistantMessageMeta)):
-        if hasattr(meta, "input_tokens") and meta.input_tokens is not None:
-            total_input += int(meta.input_tokens)
-        if hasattr(meta, "output_tokens") and meta.output_tokens is not None:
-            total_output += int(meta.output_tokens)
+        # For AssistantMessageMeta, use the structured usage field
+        if isinstance(meta, AssistantMessageMeta) and meta.usage is not None:
+            if meta.usage.input_tokens is not None:
+                total_input += int(meta.usage.input_tokens)
+            if meta.usage.output_tokens is not None:
+                total_output += int(meta.usage.output_tokens)
+        # For LLMResponseMeta, use the flat fields
+        elif isinstance(meta, LLMResponseMeta):
+            if hasattr(meta, "input_tokens") and meta.input_tokens is not None:
+                total_input += int(meta.input_tokens)
+            if hasattr(meta, "output_tokens") and meta.output_tokens is not None:
+                total_output += int(meta.output_tokens)
         if hasattr(meta, "latency_ms") and meta.latency_ms is not None:
             total_latency += int(meta.latency_ms)
         if hasattr(meta, "output_time_ms") and meta.output_time_ms is not None:
@@ -363,11 +373,9 @@ def display_chat_summary(messages: RunnerMessages, *, console: Console | None =
         messages: 要汇总的消息列表
         console: Rich Console 实例，如果为 None 则创建新的
     """
-    if console is None:
-        console = Console()
-
+    active_console = console or Console()
     summary_table = build_chat_summary_table(messages)
-    console.print(summary_table)
+    active_console.print(summary_table)
 
 
 def display_messages(
@@ -437,7 +445,7 @@ def display_messages(
         )
 
 
-def _display_single_message_compact(  # noqa: PLR0913
+def _display_single_message_compact(
     message: FlexibleRunnerMessage,
     *,
     index: int | None = None,
@@ -577,9 +585,9 @@ def _display_assistant_message_compact_v2(message: AgentAssistantMessage, contex
             meta_parts.append(f"Latency:{message.meta.latency_ms}ms")
         if message.meta.output_time_ms is not None:
             meta_parts.append(f"Output:{message.meta.output_time_ms}ms")
-        if message.meta.input_tokens is not None and message.meta.output_tokens is not None:
-            total_tokens = message.meta.input_tokens + message.meta.output_tokens
-            meta_parts.append(f"Tokens:↑{message.meta.input_tokens}↓{message.meta.output_tokens}={total_tokens}")
+        if message.meta.usage and message.meta.usage.input_tokens is not None and message.meta.usage.output_tokens is not None:
+            total_tokens = message.meta.usage.input_tokens + message.meta.usage.output_tokens
+            meta_parts.append(f"Tokens:↑{message.meta.usage.input_tokens}↓{message.meta.usage.output_tokens}={total_tokens}")
 
         if meta_parts:
             meta_info = f" [dim]({' | '.join(meta_parts)})[/dim]"

lite_agent/client.py

@@ -5,6 +5,7 @@ from typing import Any, Literal
 import litellm
 from openai.types.chat import ChatCompletionToolParam
 from openai.types.responses import FunctionToolParam
+from pydantic import BaseModel
 
 ReasoningEffort = Literal["minimal", "low", "medium", "high"]
 ThinkingConfig = dict[str, Any] | None
@@ -18,6 +19,17 @@ ReasoningConfig = (
 )
 
 
+class LLMConfig(BaseModel):
+    """LLM generation parameters configuration."""
+
+    temperature: float | None = None
+    max_tokens: int | None = None
+    top_p: float | None = None
+    frequency_penalty: float | None = None
+    presence_penalty: float | None = None
+    stop: list[str] | str | None = None
+
+
 def parse_reasoning_config(reasoning: ReasoningConfig) -> tuple[ReasoningEffort | None, ThinkingConfig]:
     """
     解析统一的推理配置，返回 reasoning_effort 和 thinking_config。
@@ -36,7 +48,10 @@ def parse_reasoning_config(reasoning: ReasoningConfig) -> tuple[ReasoningEffort
         return None, None
     if isinstance(reasoning, str):
         # 字符串类型，使用 reasoning_effort
-        return reasoning, None
+        # 确保字符串是有效的 ReasoningEffort 值
+        if reasoning in ("minimal", "low", "medium", "high"):
+            return reasoning, None  # type: ignore[return-value]
+        return None, None
     if isinstance(reasoning, dict):
         # 字典类型，使用 thinking_config
         return None, reasoning
@@ -58,13 +73,24 @@ class BaseLLMClient(abc.ABC):
         api_base: str | None = None,
         api_version: str | None = None,
         reasoning: ReasoningConfig = None,
+        llm_config: LLMConfig | None = None,
+        **llm_params: Any,  # noqa: ANN401
     ):
         self.model = model
         self.api_key = api_key
         self.api_base = api_base
         self.api_version = api_version
 
+        # 处理 LLM 生成参数
+        if llm_config is not None:
+            self.llm_config = llm_config
+        else:
+            # 从 **llm_params 创建配置
+            self.llm_config = LLMConfig(**llm_params)
+
         # 处理推理配置
+        self.reasoning_effort: ReasoningEffort | None
+        self.thinking_config: ThinkingConfig
         self.reasoning_effort, self.thinking_config = parse_reasoning_config(reasoning)
 
     @abc.abstractmethod
@@ -74,6 +100,8 @@ class BaseLLMClient(abc.ABC):
         tools: list[ChatCompletionToolParam] | None = None,
         tool_choice: str = "auto",
         reasoning: ReasoningConfig = None,
+        *,
+        streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # noqa: ANN401
         """Perform a completion request to the LLM."""
@@ -85,6 +113,8 @@ class BaseLLMClient(abc.ABC):
         tools: list[FunctionToolParam] | None = None,
         tool_choice: Literal["none", "auto", "required"] = "auto",
         reasoning: ReasoningConfig = None,
+        *,
+        streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # noqa: ANN401
         """Perform a response request to the LLM."""
@@ -108,6 +138,8 @@ class LiteLLMClient(BaseLLMClient):
         tools: list[ChatCompletionToolParam] | None = None,
         tool_choice: str = "auto",
         reasoning: ReasoningConfig = None,
+        *,
+        streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # noqa: ANN401
         """Perform a completion request to the Litellm API."""
@@ -126,10 +158,24 @@ class LiteLLMClient(BaseLLMClient):
             "api_version": self.api_version,
             "api_key": self.api_key,
             "api_base": self.api_base,
-            "stream": True,
+            "stream": streaming,
             **kwargs,
         }
 
+        # Add LLM generation parameters if specified
+        if self.llm_config.temperature is not None:
+            completion_params["temperature"] = self.llm_config.temperature
+        if self.llm_config.max_tokens is not None:
+            completion_params["max_tokens"] = self.llm_config.max_tokens
+        if self.llm_config.top_p is not None:
+            completion_params["top_p"] = self.llm_config.top_p
+        if self.llm_config.frequency_penalty is not None:
+            completion_params["frequency_penalty"] = self.llm_config.frequency_penalty
+        if self.llm_config.presence_penalty is not None:
+            completion_params["presence_penalty"] = self.llm_config.presence_penalty
+        if self.llm_config.stop is not None:
+            completion_params["stop"] = self.llm_config.stop
+
         # Add reasoning parameters if specified
         if final_reasoning_effort is not None:
             completion_params["reasoning_effort"] = final_reasoning_effort
@@ -144,6 +190,8 @@ class LiteLLMClient(BaseLLMClient):
         tools: list[FunctionToolParam] | None = None,
         tool_choice: Literal["none", "auto", "required"] = "auto",
         reasoning: ReasoningConfig = None,
+        *,
+        streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # type: ignore[return]  # noqa: ANN401
         """Perform a response request to the Litellm API."""
@@ -164,11 +212,25 @@ class LiteLLMClient(BaseLLMClient):
             "api_version": self.api_version,
             "api_key": self.api_key,
             "api_base": self.api_base,
-            "stream": True,
+            "stream": streaming,
             "store": False,
             **kwargs,
         }
 
+        # Add LLM generation parameters if specified
+        if self.llm_config.temperature is not None:
+            response_params["temperature"] = self.llm_config.temperature
+        if self.llm_config.max_tokens is not None:
+            response_params["max_tokens"] = self.llm_config.max_tokens
+        if self.llm_config.top_p is not None:
+            response_params["top_p"] = self.llm_config.top_p
+        if self.llm_config.frequency_penalty is not None:
+            response_params["frequency_penalty"] = self.llm_config.frequency_penalty
+        if self.llm_config.presence_penalty is not None:
+            response_params["presence_penalty"] = self.llm_config.presence_penalty
+        if self.llm_config.stop is not None:
+            response_params["stop"] = self.llm_config.stop
+
         # Add reasoning parameters if specified
         if final_reasoning_effort is not None:
             response_params["reasoning_effort"] = final_reasoning_effort

lite_agent/constants.py

@@ -0,0 +1,30 @@
+from typing import Literal
+
+
+class CompletionMode:
+    """Agent completion modes."""
+
+    STOP: Literal["stop"] = "stop"  # Traditional completion until model decides to stop
+    CALL: Literal["call"] = "call"  # Completion until specific tool is called
+
+
+class ToolName:
+    """System tool names."""
+
+    TRANSFER_TO_AGENT = "transfer_to_agent"
+    TRANSFER_TO_PARENT = "transfer_to_parent"
+    WAIT_FOR_USER = "wait_for_user"
+
+
+class StreamIncludes:
+    """Default stream includes configuration."""
+
+    DEFAULT_INCLUDES = (
+        "completion_raw",
+        "usage",
+        "function_call",
+        "function_call_output",
+        "content_delta",
+        "function_call_delta",
+        "assistant_message",
+    )

lite_agent/message_transfers.py

@@ -5,7 +5,7 @@ This module provides common message transfer functions that can be used
 with agents to preprocess messages before sending them to the API.
 """
 
-from lite_agent.types import RunnerMessages
+from lite_agent.types import NewUserMessage, RunnerMessages, UserTextContent
 
 
 def consolidate_history_transfer(messages: RunnerMessages) -> RunnerMessages:
@@ -43,8 +43,8 @@ def consolidate_history_transfer(messages: RunnerMessages) -> RunnerMessages:
     # Create the consolidated message
     consolidated_content = "以下是目前发生的所有交互:\n\n" + "\n".join(xml_content) + "\n\n接下来该做什么?"
 
-    # Return a single user message
-    return [{"role": "user", "content": consolidated_content}]
+    # Return a single user message using NewMessage format
+    return [NewUserMessage(content=[UserTextContent(text=consolidated_content)])]
 
 
 def _process_message_to_xml(message: dict | object) -> list[str]: