lite-agent 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

lite_agent/__init__.py

@@ -2,6 +2,7 @@
 
 from .agent import Agent
 from .message_transfers import consolidate_history_transfer
+from .rich_helpers import print_chat_history, print_chat_summary
 from .runner import Runner
 
-__all__ = ["Agent", "Runner", "consolidate_history_transfer"]
+__all__ = ["Agent", "Runner", "consolidate_history_transfer", "print_chat_history", "print_chat_summary"]

lite_agent/agent.py

@@ -2,51 +2,55 @@ from collections.abc import AsyncGenerator, Callable, Sequence
 from pathlib import Path
 from typing import Any, Optional
 
-import litellm
 from funcall import Funcall
+from jinja2 import Environment, FileSystemLoader
 from litellm import CustomStreamWrapper
 from pydantic import BaseModel
 
+from lite_agent.client import BaseLLMClient, LiteLLMClient
 from lite_agent.loggers import logger
 from lite_agent.stream_handlers import litellm_stream_handler
 from lite_agent.types import AgentChunk, AgentSystemMessage, RunnerMessages, ToolCall, ToolCallChunk, ToolCallResultChunk
 
-HANDOFFS_SOURCE_INSTRUCTIONS = """<ExtraGuide>
-You are a parent agent that can assign tasks to sub-agents.
+TEMPLATES_DIR = Path(__file__).parent / "templates"
+jinja_env = Environment(loader=FileSystemLoader(str(TEMPLATES_DIR)), autoescape=True)
 
-You can transfer conversations to other agents for specific tasks.
-If you need to assign tasks to multiple agents, you should break down the tasks and assign them one by one.
-You need to wait for one sub-agent to finish before assigning the task to the next sub-agent.
-</ExtraGuide>"""
-
-HANDOFFS_TARGET_INSTRUCTIONS = """<ExtraGuide>
-You are a sub-agent that is assigned to a specific task by your parent agent.
-
-Everything you output is intended for your parent agent to read.
-When you finish your task, you should call `transfer_to_parent` to transfer back to parent agent.
-</ExtraGuide>"""
+HANDOFFS_SOURCE_INSTRUCTIONS_TEMPLATE = jinja_env.get_template("handoffs_source_instructions.xml.j2")
+HANDOFFS_TARGET_INSTRUCTIONS_TEMPLATE = jinja_env.get_template("handoffs_target_instructions.xml.j2")
+WAIT_FOR_USER_INSTRUCTIONS_TEMPLATE = jinja_env.get_template("wait_for_user_instructions.xml.j2")
 
 
 class Agent:
     def __init__(  # noqa: PLR0913
         self,
         *,
-        model: str,
+        model: str | BaseLLMClient,
         name: str,
         instructions: str,
         tools: list[Callable] | None = None,
         handoffs: list["Agent"] | None = None,
         message_transfer: Callable[[RunnerMessages], RunnerMessages] | None = None,
+        completion_condition: str = "stop",
     ) -> None:
         self.name = name
         self.instructions = instructions
-        self.model = model
+        if isinstance(model, BaseLLMClient):
+            # If model is a BaseLLMClient instance, use it directly
+            self.client = model
+        else:
+            # Otherwise, create a LitellmClient instance
+            self.client = LiteLLMClient(model=model)
+        self.completion_condition = completion_condition
         self.handoffs = handoffs if handoffs else []
         self._parent: Agent | None = None
         self.message_transfer = message_transfer
         # Initialize Funcall with regular tools
         self.fc = Funcall(tools)
 
+        # Add wait_for_user tool if completion condition is "call"
+        if completion_condition == "call":
+            self._add_wait_for_user_tool()
+
         # Set parent for handoff agents
         if handoffs:
             for handoff_agent in handoffs:
@@ -167,11 +171,15 @@ class Agent:
 
         # Add source instructions if this agent can handoff to others
         if self.handoffs:
-            instructions = HANDOFFS_SOURCE_INSTRUCTIONS + "\n\n" + instructions
+            instructions = HANDOFFS_SOURCE_INSTRUCTIONS_TEMPLATE.render(extra_instructions=None) + "\n\n" + instructions
 
         # Add target instructions if this agent can be handed off to (has a parent)
         if self.parent:
-            instructions = HANDOFFS_TARGET_INSTRUCTIONS + "\n\n" + instructions
+            instructions = HANDOFFS_TARGET_INSTRUCTIONS_TEMPLATE.render(extra_instructions=None) + "\n\n" + instructions
+
+        # Add wait_for_user instructions if completion condition is "call"
+        if self.completion_condition == "call":
+            instructions = WAIT_FOR_USER_INSTRUCTIONS_TEMPLATE.render(extra_instructions=None) + "\n\n" + instructions
 
         return [
             AgentSystemMessage(
@@ -190,12 +198,10 @@ class Agent:
 
         self.message_histories = self.prepare_completion_messages(processed_messages)
         tools = self.fc.get_tools(target="completion")
-        resp = await litellm.acompletion(
-            model=self.model,
+        resp = await self.client.completion(
             messages=self.message_histories,
             tools=tools,
             tool_choice="auto",  # TODO: make this configurable
-            stream=True,
         )
 
         # Ensure resp is a CustomStreamWrapper
@@ -275,11 +281,11 @@ class Agent:
 
                     if next_dict.get("type") == "function_call":
                         tool_call = {
-                            "id": next_dict["function_call_id"],
+                            "id": next_dict["function_call_id"],  # type: ignore
                             "type": "function",
                             "function": {
-                                "name": next_dict["name"],
-                                "arguments": next_dict["arguments"],
+                                "name": next_dict["name"],  # type: ignore
+                                "arguments": next_dict["arguments"],  # type: ignore
                             },
                             "index": len(tool_calls),
                         }
@@ -291,7 +297,7 @@ class Agent:
                 # Create assistant message with tool_calls if any
                 assistant_msg = message_dict.copy()
                 if tool_calls:
-                    assistant_msg["tool_calls"] = tool_calls
+                    assistant_msg["tool_calls"] = tool_calls  # type: ignore
 
                 converted_messages.append(assistant_msg)
                 i = j  # Skip the function_call messages we've processed
@@ -301,8 +307,8 @@ class Agent:
                 converted_messages.append(
                     {
                         "role": "tool",
-                        "tool_call_id": message_dict["call_id"],
-                        "content": message_dict["output"],
+                        "tool_call_id": message_dict["call_id"],  # type: ignore
+                        "content": message_dict["output"],  # type: ignore
                     },
                 )
                 i += 1
@@ -314,11 +320,67 @@ class Agent:
 
             else:
                 # Regular message (user, system)
-                converted_messages.append(message_dict)
+                converted_msg = message_dict.copy()
+
+                # Handle new Response API format for user messages
+                content = message_dict.get("content")
+                if role == "user" and isinstance(content, list):
+                    converted_msg["content"] = self._convert_user_content_to_completions_format(content)  # type: ignore
+
+                converted_messages.append(converted_msg)
                 i += 1
 
         return converted_messages
 
+    def _convert_user_content_to_completions_format(self, content: list) -> list:
+        """Convert user message content from Response API format to Completion API format."""
+        # Handle the case where content might not actually be a list due to test mocking
+        if type(content) is not list:  # Use type() instead of isinstance() to avoid test mocking issues
+            return content
+
+        converted_content = []
+        for item in content:
+            if isinstance(item, dict):
+                item_type = item.get("type")
+                if item_type == "input_text":
+                    # Convert ResponseInputText to completion API format
+                    converted_content.append(
+                        {
+                            "type": "text",
+                            "text": item["text"],
+                        },
+                    )
+                elif item_type == "input_image":
+                    # Convert ResponseInputImage to completion API format
+                    if item.get("file_id"):
+                        msg = "File ID input is not supported for Completion API. Please use image_url instead of file_id for image input."
+                        raise ValueError(msg)
+
+                    if not item.get("image_url"):
+                        msg = "ResponseInputImage must have either file_id or image_url, but image_url is required for Completion API."
+                        raise ValueError(msg)
+
+                    # Build image_url object with detail inside
+                    image_data = {"url": item["image_url"]}
+                    detail = item.get("detail", "auto")
+                    if detail:  # Include detail if provided
+                        image_data["detail"] = detail
+
+                    converted_content.append(
+                        {
+                            "type": "image_url",
+                            "image_url": image_data,
+                        },
+                    )
+                else:
+                    # Keep existing format (text, image_url)
+                    converted_content.append(item)
+            else:
+                # Handle non-dict items (shouldn't happen, but just in case)
+                converted_content.append(item)
+
+        return converted_content
+
     def set_message_transfer(self, message_transfer: Callable[[RunnerMessages], RunnerMessages] | None) -> None:
         """Set or update the message transfer callback function.
 
@@ -328,3 +390,22 @@ class Agent:
                              called before making API calls to allow preprocessing of messages.
         """
         self.message_transfer = message_transfer
+
+    def _add_wait_for_user_tool(self) -> None:
+        """Add wait_for_user tool for agents with completion_condition='call'.
+
+        This tool allows the agent to signal when it has completed its task.
+        """
+
+        def wait_for_user_handler() -> str:
+            """Handler for wait_for_user function."""
+            return "Waiting for user input."
+
+        # Add dynamic tool for task completion
+        self.fc.add_dynamic_tool(
+            name="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,
+        )

lite_agent/client.py

@@ -0,0 +1,34 @@
+import abc
+from typing import Any
+
+import litellm
+from openai.types.chat import ChatCompletionToolParam
+
+
+class BaseLLMClient(abc.ABC):
+    """Base class for LLM clients."""
+
+    def __init__(self, *, model: str, api_key: str | None = None, api_base: str | None = None, api_version: str | None = None):
+        self.model = model
+        self.api_key = api_key
+        self.api_base = api_base
+        self.api_version = api_version
+
+    @abc.abstractmethod
+    async def completion(self, messages: list[Any], tools: list[ChatCompletionToolParam] | None = None, tool_choice: str = "auto") -> Any:  # noqa: ANN401
+        """Perform a completion request to the LLM."""
+
+
+class LiteLLMClient(BaseLLMClient):
+    async def completion(self, messages: list[Any], tools: list[ChatCompletionToolParam] | None = None, tool_choice: str = "auto") -> Any:  # noqa: ANN401
+        """Perform a completion request to the Litellm API."""
+        return await litellm.acompletion(
+            model=self.model,
+            messages=messages,
+            tools=tools,
+            tool_choice=tool_choice,
+            api_version=self.api_version,
+            api_key=self.api_key,
+            api_base=self.api_base,
+            stream=True,
+        )

lite_agent/rich_helpers.py

@@ -0,0 +1,503 @@
+"""
+Rich chat history renderer for lite-agent.
+
+This module provides utilities to beautifully render chat history using the rich library.
+It supports all message types including user messages, assistant messages, function calls,
+and function call outputs.
+"""
+
+import json
+from datetime import datetime, timezone
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.syntax import Syntax
+from rich.table import Table
+
+from lite_agent.types import (
+    AgentAssistantMessage,
+    AgentFunctionCallOutput,
+    AgentFunctionToolCallMessage,
+    AgentSystemMessage,
+    AgentUserMessage,
+    RunnerMessages,
+)
+
+
+def print_chat_history(
+    messages: RunnerMessages,
+    *,
+    console: Console | None = None,
+    show_timestamps: bool = True,
+    show_indices: bool = True,
+    chat_width: int = 80,
+) -> None:
+    """
+    使用 rich 库美观地渲染聊天记录。
+
+    Args:
+        messages: 要渲染的消息列表
+        console: Rich Console 实例，如果为 None 则创建新的
+        show_timestamps: 是否显示时间戳
+        show_indices: 是否显示消息索引
+        chat_width: 聊天气泡的最大宽度
+
+    Example:
+        >>> from lite_agent.runner import Runner
+        >>> from lite_agent.rich_helpers import render_chat_history
+        >>>
+        >>> runner = Runner(agent=my_agent)
+        >>> # ... add some messages ...
+        >>> render_chat_history(runner.messages)
+    """
+    if console is None:
+        console = Console()
+
+    if not messages:
+        console.print("[dim]No messages to display[/dim]")
+        return
+
+    console.print(f"\n[bold blue]Chat History[/bold blue] ([dim]{len(messages)} messages[/dim])\n")
+
+    for i, message in enumerate(messages):
+        _render_single_message(
+            message,
+            index=i if show_indices else None,
+            console=console,
+            show_timestamp=show_timestamps,
+            chat_width=chat_width,
+        )
+
+
+def _render_single_message(
+    message: object,
+    *,
+    index: int | None = None,
+    console: Console,
+    show_timestamp: bool = True,
+    chat_width: int = 80,
+) -> None:
+    """渲染单个消息。"""
+    timestamp = datetime.now(timezone.utc).strftime("%H:%M:%S") if show_timestamp else None
+
+    # 处理不同类型的消息
+    if isinstance(message, AgentUserMessage):
+        _render_user_message(message, index, console, timestamp, chat_width)
+    elif isinstance(message, AgentAssistantMessage):
+        _render_assistant_message(message, index, console, timestamp, chat_width)
+    elif isinstance(message, AgentSystemMessage):
+        _render_system_message(message, index, console, timestamp, chat_width)
+    elif isinstance(message, AgentFunctionToolCallMessage):
+        _render_function_call_message(message, index, console, timestamp, chat_width)
+    elif isinstance(message, AgentFunctionCallOutput):
+        _render_function_output_message(message, index, console, timestamp, chat_width)
+    elif isinstance(message, dict):
+        _render_dict_message(message, index, console, timestamp, chat_width)
+    else:
+        _render_unknown_message(message, index, console, timestamp, chat_width)
+
+
+def _render_user_message(
+    message: AgentUserMessage,
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染用户消息 - 靠右显示的蓝色气泡。"""
+    content = str(message.content)  # 显示完整内容，不截断
+
+    title_parts = ["👤 User"]
+    if index is not None:
+        title_parts.append(f"#{index}")
+    if timestamp:
+        title_parts.append(f"[dim]{timestamp}[/dim]")
+
+    title = " ".join(title_parts)
+
+    # 计算内容的实际宽度，用于气泡大小
+    content_width = min(len(content) + 4, chat_width)  # +4 for padding
+    bubble_width = max(content_width, 20)  # 最小宽度
+
+    # 创建用户消息气泡 - 靠右
+    panel = Panel(
+        content,
+        title=title,
+        title_align="left",
+        border_style="blue",
+        padding=(0, 1),
+        width=bubble_width,
+    )
+
+    # 用户消息靠右
+    console.print(panel, justify="right")
+
+
+def _render_assistant_message(
+    message: AgentAssistantMessage,
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染助手消息 - 靠左显示的绿色气泡。"""
+    content = message.content  # 显示完整内容，不截断
+
+    title_parts = ["🤖 Assistant"]
+    if index is not None:
+        title_parts.append(f"#{index}")
+    if timestamp:
+        title_parts.append(f"[dim]{timestamp}[/dim]")
+
+    title = " ".join(title_parts)
+
+    # 计算内容的实际宽度，用于气泡大小
+    content_width = min(len(content) + 4, chat_width)  # +4 for padding
+    bubble_width = max(content_width, 20)  # 最小宽度
+
+    # 创建助手消息气泡 - 靠左
+    panel = Panel(
+        content,
+        title=title,
+        title_align="left",
+        border_style="green",
+        padding=(0, 1),
+        width=bubble_width,
+    )
+
+    # 助手消息靠左
+    console.print(panel)
+
+
+def _render_system_message(
+    message: AgentSystemMessage,
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染系统消息 - 居中显示的黄色气泡。"""
+    content = message.content  # 显示完整内容，不截断
+
+    title_parts = ["⚙️ System"]
+    if index is not None:
+        title_parts.append(f"#{index}")
+    if timestamp:
+        title_parts.append(f"[dim]{timestamp}[/dim]")
+
+    title = " ".join(title_parts)
+
+    # 系统消息居中显示，使用较小的宽度
+    console.print(
+        Panel(
+            content,
+            title=title,
+            title_align="center",
+            border_style="yellow",
+            padding=(0, 1),
+            width=min(len(content) + 10, chat_width),
+        ),
+        justify="center",
+    )
+
+
+def _render_function_call_message(
+    message: AgentFunctionToolCallMessage,
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染函数调用消息 - 靠左显示的紫色气泡。"""
+    title_parts = ["🛠️ Function Call"]
+    if index is not None:
+        title_parts.append(f"#{index}")
+    if timestamp:
+        title_parts.append(f"[dim]{timestamp}[/dim]")
+
+    title = " ".join(title_parts)
+
+    # 创建表格显示函数调用详情
+    table = Table(show_header=False, box=None, padding=0)
+    table.add_column("Field", style="cyan", width=12)
+    table.add_column("Value", style="white")
+
+    table.add_row("Name:", f"[bold]{message.name}[/bold]")
+    table.add_row("Call ID:", f"[dim]{message.function_call_id}[/dim]")
+
+    if message.arguments:
+        # 尝试格式化 JSON 参数 - 显示完整内容
+        try:
+            parsed_args = json.loads(message.arguments)
+            formatted_args = json.dumps(parsed_args, indent=2, ensure_ascii=False)
+            syntax = Syntax(formatted_args, "json", theme="monokai", line_numbers=False)
+            table.add_row("Arguments:", syntax)
+        except (json.JSONDecodeError, TypeError):
+            table.add_row("Arguments:", message.arguments)
+
+    # 函数调用消息靠左
+    console.print(
+        Panel(
+            table,
+            title=title,
+            title_align="left",
+            border_style="magenta",
+            padding=(0, 1),
+            width=min(chat_width, 100),
+        ),
+    )
+
+
+def _render_function_output_message(
+    message: AgentFunctionCallOutput,
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染函数输出消息 - 靠左显示的青色气泡。"""
+    title_parts = ["📤 Function Output"]
+    if index is not None:
+        title_parts.append(f"#{index}")
+    if timestamp:
+        title_parts.append(f"[dim]{timestamp}[/dim]")
+
+    title = " ".join(title_parts)
+
+    output_content = message.output  # 显示完整内容，不截断
+
+    # 创建表格显示函数输出详情
+    table = Table(show_header=False, box=None, padding=0)
+    table.add_column("Field", style="cyan", width=12)
+    table.add_column("Value", style="white")
+
+    table.add_row("Call ID:", f"[dim]{message.call_id}[/dim]")
+    table.add_row("Output:", output_content)
+
+    # 函数输出消息靠左
+    console.print(
+        Panel(
+            table,
+            title=title,
+            title_align="left",
+            border_style="cyan",
+            padding=(0, 1),
+            width=min(chat_width, 100),
+        ),
+    )
+
+
+def _render_role_based_dict_message(  # noqa: PLR0913
+    *,
+    message: dict[str, object],
+    role: str,
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染基于角色的字典消息。"""
+    content = str(message.get("content", ""))  # 显示完整内容，不截断
+
+    title_parts = []
+    if role == "user":
+        title_parts = ["👤 User"]
+        border_style = "blue"
+        # 用户消息靠右
+        content_width = min(len(content) + 4, chat_width)
+        bubble_width = max(content_width, 20)
+        if index is not None:
+            title_parts.append(f"#{index}")
+        if timestamp:
+            title_parts.append(f"[dim]{timestamp}[/dim]")
+
+        panel = Panel(
+            content,
+            title=" ".join(title_parts),
+            title_align="left",
+            border_style=border_style,
+            padding=(0, 1),
+            width=bubble_width,
+        )
+        console.print(panel, justify="right")
+    elif role == "assistant":
+        title_parts = ["🤖 Assistant"]
+        border_style = "green"
+        # 助手消息靠左
+        content_width = min(len(content) + 4, chat_width)
+        bubble_width = max(content_width, 20)
+        if index is not None:
+            title_parts.append(f"#{index}")
+        if timestamp:
+            title_parts.append(f"[dim]{timestamp}[/dim]")
+
+        panel = Panel(
+            content,
+            title=" ".join(title_parts),
+            title_align="left",
+            border_style=border_style,
+            padding=(0, 1),
+            width=bubble_width,
+        )
+        # 助手消息靠左
+        console.print(panel)
+    else:  # system
+        title_parts = ["⚙️ System"]
+        border_style = "yellow"
+        if index is not None:
+            title_parts.append(f"#{index}")
+        if timestamp:
+            title_parts.append(f"[dim]{timestamp}[/dim]")
+
+        # 系统消息居中
+        console.print(
+            Panel(
+                content,
+                title=" ".join(title_parts),
+                title_align="center",
+                border_style=border_style,
+                padding=(0, 1),
+                width=min(len(content) + 10, chat_width),
+            ),
+            justify="center",
+        )
+
+
+def _render_dict_message(
+    message: dict[str, object],
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染字典格式的消息。"""
+    message_type = message.get("type")
+    role = message.get("role")
+
+    if message_type == "function_call":
+        # 创建临时 AgentFunctionToolCallMessage 对象进行渲染
+        temp_message = AgentFunctionToolCallMessage(
+            type="function_call",
+            function_call_id=str(message.get("function_call_id", "")),
+            name=str(message.get("name", "unknown")),
+            arguments=str(message.get("arguments", "")),
+            content=str(message.get("content", "")),
+        )
+        _render_function_call_message(temp_message, index, console, timestamp, chat_width)
+    elif message_type == "function_call_output":
+        # 创建临时 AgentFunctionCallOutput 对象进行渲染
+        temp_message = AgentFunctionCallOutput(
+            type="function_call_output",
+            call_id=str(message.get("call_id", "")),
+            output=str(message.get("output", "")),
+        )
+        _render_function_output_message(temp_message, index, console, timestamp, chat_width)
+    elif role in ["user", "assistant", "system"]:
+        _render_role_based_dict_message(
+            message=message,
+            role=str(role),
+            index=index,
+            console=console,
+            timestamp=timestamp,
+            chat_width=chat_width,
+        )
+    else:
+        _render_unknown_message(message, index, console, timestamp, chat_width)
+
+
+def _render_unknown_message(
+    message: object,
+    index: int | None,
+    console: Console,
+    timestamp: str | None,
+    chat_width: int,
+) -> None:
+    """渲染未知类型的消息 - 居中显示的红色气泡。"""
+    title_parts = ["❓ Unknown"]
+    if index is not None:
+        title_parts.append(f"#{index}")
+    if timestamp:
+        title_parts.append(f"[dim]{timestamp}[/dim]")
+
+    title = " ".join(title_parts)
+
+    # 尝试将消息转换为可读格式 - 显示完整内容
+    try:
+        content = str(message.model_dump()) if hasattr(message, "model_dump") else str(message)  # type: ignore[attr-defined]
+    except Exception:
+        content = str(message)
+
+    console.print(
+        Panel(
+            content,
+            title=title,
+            title_align="center",
+            border_style="red",
+            padding=(0, 1),
+            width=min(len(content) + 10, chat_width),
+        ),
+        justify="center",
+    )
+
+
+def create_chat_summary_table(messages: RunnerMessages) -> Table:
+    """
+    创建聊天记录摘要表格。
+
+    Args:
+        messages: 要汇总的消息列表
+
+    Returns:
+        Rich Table 对象，包含消息统计信息
+    """
+    table = Table(title="Chat Summary")
+    table.add_column("Message Type", style="cyan")
+    table.add_column("Count", justify="right", style="green")
+
+    # 统计各种消息类型
+    counts = {
+        "User": 0,
+        "Assistant": 0,
+        "System": 0,
+        "Function Call": 0,
+        "Function Output": 0,
+        "Unknown": 0,
+    }
+
+    for message in messages:
+        if isinstance(message, AgentUserMessage) or (isinstance(message, dict) and message.get("role") == "user"):
+            counts["User"] += 1
+        elif isinstance(message, AgentAssistantMessage) or (isinstance(message, dict) and message.get("role") == "assistant"):
+            counts["Assistant"] += 1
+        elif isinstance(message, AgentSystemMessage) or (isinstance(message, dict) and message.get("role") == "system"):
+            counts["System"] += 1
+        elif isinstance(message, AgentFunctionToolCallMessage) or (isinstance(message, dict) and message.get("type") == "function_call"):
+            counts["Function Call"] += 1
+        elif isinstance(message, AgentFunctionCallOutput) or (isinstance(message, dict) and message.get("type") == "function_call_output"):
+            counts["Function Output"] += 1
+        else:
+            counts["Unknown"] += 1
+
+    # 只显示计数大于0的类型
+    for msg_type, count in counts.items():
+        if count > 0:
+            table.add_row(msg_type, str(count))
+
+    table.add_row("[bold]Total[/bold]", f"[bold]{len(messages)}[/bold]")
+
+    return table
+
+
+def print_chat_summary(messages: RunnerMessages, *, console: Console | None = None) -> None:
+    """
+    打印聊天记录摘要。
+
+    Args:
+        messages: 要汇总的消息列表
+        console: Rich Console 实例，如果为 None 则创建新的
+    """
+    if console is None:
+        console = Console()
+
+    summary_table = create_chat_summary_table(messages)
+    console.print(summary_table)

lite_agent/runner.py

@@ -14,10 +14,12 @@ from lite_agent.types import (
     AgentFunctionToolCallMessage,
     AgentSystemMessage,
     AgentUserMessage,
+    FlexibleRunnerMessage,
+    MessageDict,
     RunnerMessage,
-    RunnerMessages,
     ToolCall,
     ToolCallFunction,
+    UserInput,
 )
 
 if TYPE_CHECKING:
@@ -70,6 +72,7 @@ class Runner:
                         ),
                     )
             return  # Stop processing other tool calls after transfer
+
         return_parent_calls = [tc for tc in tool_calls if tc.function.name == "transfer_to_parent"]
         if return_parent_calls:
             # Handle multiple transfer_to_parent calls (only execute the first one)
@@ -87,6 +90,7 @@ class Runner:
                         ),
                     )
             return  # Stop processing other tool calls after transfer
+
         async for tool_call_chunk in self.agent.handle_tool_calls(tool_calls, context=context):
             if tool_call_chunk.type == "tool_call" and tool_call_chunk.type in includes:
                 yield tool_call_chunk
@@ -108,7 +112,7 @@ class Runner:
 
     def run(
         self,
-        user_input: RunnerMessages | str,
+        user_input: UserInput,
         max_steps: int = 20,
         includes: Sequence[AgentChunkType] | None = None,
         context: "Any | None" = None,  # noqa: ANN401
@@ -118,18 +122,32 @@ class Runner:
         includes = self._normalize_includes(includes)
         if isinstance(user_input, str):
             self.messages.append(AgentUserMessage(role="user", content=user_input))
-        else:
+        elif isinstance(user_input, (list, tuple)):
+            # Handle sequence of messages
             for message in user_input:
                 self.append_message(message)
+        else:
+            # Handle single message (BaseModel, TypedDict, or dict)
+            # Type assertion needed due to the complex union type
+            self.append_message(user_input)  # type: ignore[arg-type]
         return self._run(max_steps, includes, self._normalize_record_path(record_to), context=context)
 
-    async def _run(self, max_steps: int, includes: Sequence[AgentChunkType], record_to: Path | None = None, context: "Any | None" = None) -> AsyncGenerator[AgentChunk, None]:  # noqa: ANN401
+    async def _run(self, max_steps: int, includes: Sequence[AgentChunkType], record_to: Path | None = None, context: "Any | None" = None) -> AsyncGenerator[AgentChunk, None]:  # noqa: ANN401, C901
         """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}")
         steps = 0
         finish_reason = None
 
-        while finish_reason != "stop" and steps < max_steps:
+        # Determine completion condition based on agent configuration
+        completion_condition = getattr(self.agent, "completion_condition", "stop")
+
+        def is_finish() -> bool:
+            if completion_condition == "call":
+                function_calls = self._find_pending_function_calls()
+                return any(getattr(fc, "name", None) == "wait_for_user" for fc in function_calls)
+            return finish_reason == "stop"
+
+        while not is_finish() and steps < max_steps:
             resp = await self.agent.completion(self.messages, record_to_file=record_to)
             async for chunk in resp:
                 if chunk.type in includes:
@@ -202,14 +220,13 @@ class Runner:
                 msg = "Cannot continue running without a valid last message from the assistant."
                 raise ValueError(msg)
 
-            # If we have an assistant message but no pending function calls,
-            # that means there's nothing to continue
-            msg = "Cannot continue running without pending function calls."
-            raise ValueError(msg)
+            resp = self._run(max_steps=max_steps, includes=includes, record_to=self._normalize_record_path(record_to), context=context)
+            async for chunk in resp:
+                yield chunk
 
     async def run_until_complete(
         self,
-        user_input: RunnerMessages | str,
+        user_input: UserInput,
         max_steps: int = 20,
         includes: list[AgentChunkType] | None = None,
         record_to: PathLike | str | None = None,
@@ -223,12 +240,13 @@ class Runner:
         # The final message from the stream handler might still contain tool_calls
         # We need to convert it to responses format
         if hasattr(message, "tool_calls") and message.tool_calls:
-            # Add the assistant message without tool_calls
-            assistant_msg = AgentAssistantMessage(
-                role="assistant",
-                content=message.content,
-            )
-            self.messages.append(assistant_msg)
+            if message.content:
+                # Add the assistant message without tool_calls
+                assistant_msg = AgentAssistantMessage(
+                    role="assistant",
+                    content=message.content,
+                )
+                self.messages.append(assistant_msg)
 
             # Add function call messages
             for tool_call in message.tool_calls:
@@ -285,7 +303,134 @@ class Runner:
             tool_calls.append(tool_call)
         return tool_calls
 
-    def append_message(self, message: RunnerMessage | dict) -> None:
+    def set_chat_history(self, messages: Sequence[FlexibleRunnerMessage], 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
+        based on transfer_to_agent and transfer_to_parent function calls.
+
+        Args:
+            messages: List of messages to set as the chat history
+            root_agent: The root agent to use if no transfers are found. If None, uses self.agent
+        """
+        # Clear current messages
+        self.messages.clear()
+
+        # Set initial agent
+        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)
+
+        # 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]]:
+        """Get the messages in JSONL format."""
+        return [msg.model_dump(mode="json") for msg in self.messages]
+
+    def _track_agent_transfer_in_message(self, message: FlexibleRunnerMessage, current_agent: Agent) -> Agent:
+        """Track agent transfers in a single message.
+
+        Args:
+            message: The message to analyze for transfers
+            current_agent: The currently active agent
+
+        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, AgentFunctionToolCallMessage):
+            return self._track_transfer_from_function_call_message(message, 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 _track_transfer_from_function_call_message(self, message: AgentFunctionToolCallMessage, current_agent: Agent) -> Agent:
+        """Track transfers from AgentFunctionToolCallMessage objects."""
+        if message.name == "transfer_to_agent":
+            return self._handle_transfer_to_agent_tracking(message.arguments, current_agent)
+
+        if message.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:
+            args_dict = json.loads(arguments) if isinstance(arguments, str) else arguments
+
+            target_agent_name = args_dict.get("name")
+            if target_agent_name:
+                target_agent = self._find_agent_by_name(current_agent, target_agent_name)
+                if target_agent:
+                    logger.debug(f"History tracking: Transferring from {current_agent.name} to {target_agent_name}")
+                    return target_agent
+
+                logger.warning(f"Target agent '{target_agent_name}' not found in handoffs during history setup")
+        except (json.JSONDecodeError, KeyError, TypeError) as e:
+            logger.warning(f"Failed to parse transfer_to_agent arguments during history setup: {e}")
+
+        return current_agent
+
+    def _handle_transfer_to_parent_tracking(self, current_agent: Agent) -> Agent:
+        """Handle transfer_to_parent function call tracking."""
+        if current_agent.parent:
+            logger.debug(f"History tracking: Transferring from {current_agent.name} back to parent {current_agent.parent.name}")
+            return current_agent.parent
+
+        logger.warning(f"Agent {current_agent.name} has no parent to transfer back to during history setup")
+        return current_agent
+
+    def _find_agent_by_name(self, root_agent: Agent, target_name: str) -> Agent | None:
+        """Find an agent by name in the handoffs tree starting from root_agent.
+
+        Args:
+            root_agent: The root agent to start searching from
+            target_name: The name of the agent to find
+
+        Returns:
+            The agent if found, None otherwise
+        """
+        # Check direct handoffs from current agent
+        if root_agent.handoffs:
+            for agent in root_agent.handoffs:
+                if agent.name == target_name:
+                    return agent
+
+        # If not found in direct handoffs, check if we need to look in parent's handoffs
+        # This handles cases where agents can transfer to siblings
+        current = root_agent
+        while current.parent is not None:
+            current = current.parent
+            if current.handoffs:
+                for agent in current.handoffs:
+                    if agent.name == target_name:
+                        return agent
+
+        return None
+
+    def append_message(self, message: FlexibleRunnerMessage) -> None:
         if isinstance(message, RunnerMessage):
             self.messages.append(message)
         elif isinstance(message, dict):

lite_agent/templates/handoffs_source_instructions.xml.j2

@@ -0,0 +1,10 @@
+<HandoffsGuide>
+You are a parent agent that can assign tasks to sub-agents.
+
+You can transfer conversations to other agents for specific tasks.
+If you need to assign tasks to multiple agents, you should break down the tasks and assign them one by one.
+You need to wait for one sub-agent to finish before assigning the task to the next sub-agent.
+{% if extra_instructions %}
+{{ extra_instructions }}
+{% endif %}
+</HandoffsGuide>

lite_agent/templates/handoffs_target_instructions.xml.j2

@@ -0,0 +1,9 @@
+<TransferToParentGuide>
+You are a sub-agent that is assigned to a specific task by your parent agent.
+
+Everything you output is intended for your parent agent to read.
+When you finish your task, you should call `transfer_to_parent` to transfer back to parent agent.
+{% if extra_instructions %}
+{{ extra_instructions }}
+{% endif %}
+</TransferToParentGuide>

lite_agent/templates/wait_for_user_instructions.xml.j2

@@ -0,0 +1,6 @@
+<WaitForUserGuide>
+When you have completed your assigned task or need more information from the user, you must call the `wait_for_user` function.
+{% if extra_instructions %}
+{{ extra_instructions }}
+{% endif %}
+</WaitForUserGuide>

lite_agent/types/__init__.py

@@ -18,12 +18,22 @@ from .messages import (
     AgentSystemMessage,
     AgentUserMessage,
     AssistantMessage,
+    AssistantMessageDict,
+    FlexibleRunnerMessage,
+    FunctionCallDict,
+    FunctionCallOutputDict,
     Message,
+    MessageDict,
+    ResponseInputImage,
+    ResponseInputText,
     RunnerMessage,
     RunnerMessages,
+    SystemMessageDict,
+    UserInput,
     UserMessageContentItemImageURL,
     UserMessageContentItemImageURLImageURL,
     UserMessageContentItemText,
+    UserMessageDict,
 )
 from .tool_calls import ToolCall, ToolCallFunction
 
@@ -37,19 +47,29 @@ __all__ = [
     "AgentSystemMessage",
     "AgentUserMessage",
     "AssistantMessage",
+    "AssistantMessageDict",
     "CompletionRawChunk",
     "ContentDeltaChunk",
     "FinalMessageChunk",
+    "FlexibleRunnerMessage",
+    "FunctionCallDict",
+    "FunctionCallOutputDict",
     "Message",
+    "MessageDict",
+    "ResponseInputImage",
+    "ResponseInputText",
     "RunnerMessage",
     "RunnerMessages",
+    "SystemMessageDict",
     "ToolCall",
     "ToolCallChunk",
     "ToolCallDeltaChunk",
     "ToolCallFunction",
     "ToolCallResultChunk",
     "UsageChunk",
+    "UserInput",
     "UserMessageContentItemImageURL",
     "UserMessageContentItemImageURLImageURL",
     "UserMessageContentItemText",
+    "UserMessageDict",
 ]

lite_agent/types/messages.py

@@ -1,25 +1,71 @@
 from collections.abc import Sequence
-from typing import Literal
+from typing import Any, Literal, NotRequired, TypedDict
 
 from pydantic import BaseModel
-from rich import Any
 
 from .tool_calls import ToolCall
 
 
-class AssistantMessage(BaseModel):
-    id: str
-    index: int
-    role: Literal["assistant"] = "assistant"
-    content: str = ""
-    tool_calls: list[ToolCall] | None = None
+class ResponseInputImageDict(TypedDict):
+    detail: NotRequired[Literal["low", "high", "auto"]]
+    type: Literal["input_image"]
+    file_id: str | None
+    image_url: str | None
 
 
-class Message(BaseModel):
-    role: str
+class ResponseInputTextDict(TypedDict):
+    text: str
+    type: Literal["input_text"]
+
+
+# TypedDict definitions for better type hints
+class UserMessageDict(TypedDict):
+    role: Literal["user"]
+    content: str | Sequence[ResponseInputTextDict | ResponseInputImageDict]
+
+
+class AssistantMessageDict(TypedDict):
+    role: Literal["assistant"]
     content: str
 
 
+class SystemMessageDict(TypedDict):
+    role: Literal["system"]
+    content: str
+
+
+class FunctionCallDict(TypedDict):
+    type: Literal["function_call"]
+    function_call_id: str
+    name: str
+    arguments: str
+    content: str
+
+
+class FunctionCallOutputDict(TypedDict):
+    type: Literal["function_call_output"]
+    call_id: str
+    output: str
+
+
+# Union type for all supported message dictionary formats
+MessageDict = UserMessageDict | AssistantMessageDict | SystemMessageDict | FunctionCallDict | FunctionCallOutputDict
+
+
+# Response API format input types
+class ResponseInputText(BaseModel):
+    text: str
+    type: Literal["input_text"]
+
+
+class ResponseInputImage(BaseModel):
+    detail: Literal["low", "high", "auto"] = "auto"
+    type: Literal["input_image"]
+    file_id: str | None = None
+    image_url: str | None = None
+
+
+# Compatibility types for old completion API format
 class UserMessageContentItemText(BaseModel):
     type: Literal["text"]
     text: str
@@ -34,9 +80,23 @@ class UserMessageContentItemImageURL(BaseModel):
     image_url: UserMessageContentItemImageURLImageURL
 
 
+# Legacy types - keeping for compatibility
+class AssistantMessage(BaseModel):
+    id: str
+    index: int
+    role: Literal["assistant"] = "assistant"
+    content: str = ""
+    tool_calls: list[ToolCall] | None = None
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+
 class AgentUserMessage(BaseModel):
     role: Literal["user"]
-    content: str | Sequence[UserMessageContentItemText | UserMessageContentItemImageURL]
+    content: str | Sequence[ResponseInputText | ResponseInputImage | UserMessageContentItemText | UserMessageContentItemImageURL]
 
 
 class AgentAssistantMessage(BaseModel):
@@ -65,4 +125,11 @@ class AgentFunctionCallOutput(BaseModel):
 
 RunnerMessage = AgentUserMessage | AgentAssistantMessage | AgentSystemMessage | AgentFunctionToolCallMessage | AgentFunctionCallOutput
 AgentMessage = RunnerMessage | AgentSystemMessage
-RunnerMessages = Sequence[RunnerMessage | dict[str, Any]]
+
+# Enhanced type definitions for better type hints
+# Supports BaseModel instances, TypedDict, and plain dict
+FlexibleRunnerMessage = RunnerMessage | MessageDict | dict[str, Any]
+RunnerMessages = Sequence[FlexibleRunnerMessage]
+
+# Type alias for user input - supports string, single message, or sequence of messages
+UserInput = str | FlexibleRunnerMessage | RunnerMessages

{lite_agent-0.2.0.dist-info → lite_agent-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lite-agent
-Version: 0.2.0
+Version: 0.3.0
 Summary: A lightweight, extensible framework for building AI agent.
 Author-email: Jianqi Pan <jannchie@gmail.com>
 License: MIT

lite_agent-0.3.0.dist-info/RECORD

@@ -0,0 +1,22 @@
+lite_agent/__init__.py,sha256=Xaex4kVGxZzg_hhO17b8_tVXf63xFbSzBOlVdhRf-Ng,338
+lite_agent/agent.py,sha256=MsdAnM2pqh9RLTGOFuynowkS52QXUb7vHaMHQyiOpoQ,17627
+lite_agent/client.py,sha256=e_BsXo6KUgleRFkSPSESUoIPvyLXWyJ9E1AzExYXXsk,1236
+lite_agent/loggers.py,sha256=XkNkdqwD_nQGfhQJ-bBWT7koci_mMkNw3aBpyMhOICw,57
+lite_agent/message_transfers.py,sha256=nT7-tID20RK2yoN-rDiEE6sSclluSlhYSkayCzmPwk8,3984
+lite_agent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+lite_agent/rich_helpers.py,sha256=6HdjMwR2U7P5ieZ5KHQNUsQbW-apG67GzQ_nvJm774E,15583
+lite_agent/runner.py,sha256=LWhtmxrYxL4sFi3ZKeWWRymK3tGCvWfIh1-EcWf8P0g,27822
+lite_agent/processors/__init__.py,sha256=X78GKL_IWwW2lg8w4DD6GOFWLzAR2wTfKxHlvOkcuUQ,114
+lite_agent/processors/stream_chunk_processor.py,sha256=nMA_cW7FDpXwJvm4F8vFwBXmHHsSELQFcoNEjH3xvn8,4751
+lite_agent/stream_handlers/__init__.py,sha256=2GSiG0VUgcQlFMl6JkGAqikXMII1a43Hr-J5NIct6dk,115
+lite_agent/stream_handlers/litellm.py,sha256=NNMAl8Bvoc2xe-qWKtfqvJQA2yr3sz1IUU90rQ_9iBw,3976
+lite_agent/templates/handoffs_source_instructions.xml.j2,sha256=2XsXQlBzk38qbxGrfyt8y2b0KlZmsV_1xavLufcdkHc,428
+lite_agent/templates/handoffs_target_instructions.xml.j2,sha256=gSbWVYYcovPKbGpFc0kqGSJ5Y5UC3fOHyUmZfcrDgSE,356
+lite_agent/templates/wait_for_user_instructions.xml.j2,sha256=wXbcYD5Q1FaCGVBm3Hz_Cp7nnoK7KzloP0ao-jYMwPk,231
+lite_agent/types/__init__.py,sha256=8l2RL-55sRHQW-sTmtKkKzCQGLrENaJT7Cgy5iA5xCo,1767
+lite_agent/types/chunks.py,sha256=Ro5BtrrdsYGkKrEekIhs9vIrBM7HljtgOkHherH8B3k,1697
+lite_agent/types/messages.py,sha256=A66YVl2IYMMTlnEdGlbCXqMztSSMSjS9F2yyebBlKR0,3364
+lite_agent/types/tool_calls.py,sha256=Xnut8-2-Ld9vgA2GKJY6BbFlBaAv_n4W7vo7Jx21A-E,260
+lite_agent-0.3.0.dist-info/METADATA,sha256=l4SLUuFQlcrlr-CHOIaWqQ3WOyYnPVIsHjjxKjR_g4E,3455
+lite_agent-0.3.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+lite_agent-0.3.0.dist-info/RECORD,,

lite_agent-0.2.0.dist-info/RECORD

@@ -1,17 +0,0 @@
-lite_agent/__init__.py,sha256=T-jxY0aHwqPz5gs-ZzH1zvFgFdYVmOd15fnEc5A2H6U,229
-lite_agent/agent.py,sha256=UE06H1Huk4gHGovbh9DuAJAH2zC0hqf1YXgrVa59_Os,13546
-lite_agent/loggers.py,sha256=XkNkdqwD_nQGfhQJ-bBWT7koci_mMkNw3aBpyMhOICw,57
-lite_agent/message_transfers.py,sha256=nT7-tID20RK2yoN-rDiEE6sSclluSlhYSkayCzmPwk8,3984
-lite_agent/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lite_agent/runner.py,sha256=lO8Z5R2RxjjgYzDAeYQbO2niQeieXOEzqUEN484G7QU,21367
-lite_agent/processors/__init__.py,sha256=X78GKL_IWwW2lg8w4DD6GOFWLzAR2wTfKxHlvOkcuUQ,114
-lite_agent/processors/stream_chunk_processor.py,sha256=nMA_cW7FDpXwJvm4F8vFwBXmHHsSELQFcoNEjH3xvn8,4751
-lite_agent/stream_handlers/__init__.py,sha256=2GSiG0VUgcQlFMl6JkGAqikXMII1a43Hr-J5NIct6dk,115
-lite_agent/stream_handlers/litellm.py,sha256=NNMAl8Bvoc2xe-qWKtfqvJQA2yr3sz1IUU90rQ_9iBw,3976
-lite_agent/types/__init__.py,sha256=sc2cdX1tPisfQwu2-apZtUa3u_Q6WDEqzNglfXhwCJo,1295
-lite_agent/types/chunks.py,sha256=Ro5BtrrdsYGkKrEekIhs9vIrBM7HljtgOkHherH8B3k,1697
-lite_agent/types/messages.py,sha256=cmMcj_r1_R9Pgu6ixmBIq2uwwqi5KzIllxGoxpcxF3w,1531
-lite_agent/types/tool_calls.py,sha256=Xnut8-2-Ld9vgA2GKJY6BbFlBaAv_n4W7vo7Jx21A-E,260
-lite_agent-0.2.0.dist-info/METADATA,sha256=egjWZCc9UgPAHMCJqZrjBXlXrG39f9LoW9k6DzntMas,3455
-lite_agent-0.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-lite_agent-0.2.0.dist-info/RECORD,,