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

lite_agent/__init__.py

@@ -1,8 +1,8 @@
 """Lite Agent - A lightweight AI agent framework."""
 
 from .agent import Agent
+from .chat_display import display_chat_summary, display_messages
 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", "print_chat_history", "print_chat_summary"]
+__all__ = ["Agent", "Runner", "consolidate_history_transfer", "display_chat_summary", "display_messages"]

lite_agent/agent.py

@@ -1,3 +1,4 @@
+import time
 from collections.abc import AsyncGenerator, Callable, Sequence
 from pathlib import Path
 from typing import Any, Optional
@@ -5,12 +6,12 @@ from typing import Any, Optional
 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
+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.types.messages import NewAssistantMessage, NewSystemMessage, NewUserMessage
 
 TEMPLATES_DIR = Path(__file__).parent / "templates"
 jinja_env = Environment(loader=FileSystemLoader(str(TEMPLATES_DIR)), autoescape=True)
@@ -162,41 +163,120 @@ class Agent:
             # Regenerate transfer tools to include the new agent
             self._add_transfer_tools(self.handoffs)
 
-    def prepare_completion_messages(self, messages: RunnerMessages) -> list[dict[str, str]]:
-        # Convert from responses format to completions format
+    def prepare_completion_messages(self, messages: RunnerMessages) -> list[dict]:
+        """Prepare messages for completions API (with conversion)."""
         converted_messages = self._convert_responses_to_completions_format(messages)
-
-        # Prepare instructions with handoff-specific additions
         instructions = self.instructions
-
-        # Add source instructions if this agent can handoff to others
         if self.handoffs:
             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_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(
-                role="system",
+            system_message_to_llm_dict(NewSystemMessage(
                 content=f"You are {self.name}. {instructions}",
-            ).model_dump(),
+            )),
             *converted_messages,
         ]
 
+    def prepare_responses_messages(self, messages: RunnerMessages) -> list[dict[str, Any]]:
+        """Prepare messages for responses API (no conversion, just add system message if needed)."""
+        instructions = self.instructions
+        if self.handoffs:
+            instructions = HANDOFFS_SOURCE_INSTRUCTIONS_TEMPLATE.render(extra_instructions=None) + "\n\n" + instructions
+        if self.parent:
+            instructions = HANDOFFS_TARGET_INSTRUCTIONS_TEMPLATE.render(extra_instructions=None) + "\n\n" + instructions
+        if self.completion_condition == "call":
+            instructions = WAIT_FOR_USER_INSTRUCTIONS_TEMPLATE.render(extra_instructions=None) + "\n\n" + instructions
+        res: list[dict[str, Any]] = [
+            {
+                "role": "system",
+                "content": f"You are {self.name}. {instructions}",
+            },
+        ]
+        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,
+                                },
+                            )
+            elif isinstance(message, NewSystemMessage):
+                res.append(
+                    {
+                        "role": "system",
+                        "content": message.content,
+                    },
+                )
+            elif isinstance(message, NewUserMessage):
+                contents = []
+                for item in message.content:
+                    match item.type:
+                        case "text":
+                            contents.append(
+                                {
+                                    "type": "input_text",
+                                    "text": item.text,
+                                },
+                            )
+                        case "image":
+                            contents.append(
+                                {
+                                    "type": "input_image",
+                                    "image_url": item.image_url,
+                                },
+                            )
+                        case "file":
+                            contents.append(
+                                {
+                                    "type": "input_file",
+                                    "file_id": item.file_id,
+                                    "file_name": item.file_name,
+                                },
+                            )
+                res.append(
+                    {
+                        "role": message.role,
+                        "content": contents,
+                    },
+                )
+            # Handle dict messages (legacy format)
+            elif isinstance(message, dict):
+                res.append(message)
+        return res
+
     async def completion(self, messages: RunnerMessages, record_to_file: Path | None = None) -> AsyncGenerator[AgentChunk, None]:
-        # Apply message transfer callback if provided
+        # Apply message transfer callback if provided - always use legacy format for LLM compatibility
         processed_messages = messages
         if self.message_transfer:
             logger.debug(f"Applying message transfer callback for agent {self.name}")
             processed_messages = self.message_transfer(messages)
 
+        # For completions API, use prepare_completion_messages
         self.message_histories = self.prepare_completion_messages(processed_messages)
+
         tools = self.fc.get_tools(target="completion")
         resp = await self.client.completion(
             messages=self.message_histories,
@@ -206,10 +286,27 @@ class Agent:
 
         # Ensure resp is a CustomStreamWrapper
         if isinstance(resp, CustomStreamWrapper):
-            return litellm_stream_handler(resp, record_to=record_to_file)
+            return litellm_completion_stream_handler(resp, record_to=record_to_file)
         msg = "Response is not a CustomStreamWrapper, cannot stream chunks."
         raise TypeError(msg)
 
+    async def responses(self, messages: RunnerMessages, record_to_file: Path | None = None) -> AsyncGenerator[AgentChunk, None]:
+        # Apply message transfer callback if provided - always use legacy format for LLM compatibility
+        processed_messages = messages
+        if self.message_transfer:
+            logger.debug(f"Applying message transfer callback for agent {self.name}")
+            processed_messages = self.message_transfer(messages)
+
+        # For responses API, use prepare_responses_messages (no conversion)
+        self.message_histories = self.prepare_responses_messages(processed_messages)
+        tools = self.fc.get_tools()
+        resp = await self.client.responses(
+            messages=self.message_histories,
+            tools=tools,
+            tool_choice="auto",  # TODO: make this configurable
+        )
+        return litellm_response_stream_handler(resp, record_to=record_to_file)
+
     async def list_require_confirm_tools(self, tool_calls: Sequence[ToolCall] | None) -> Sequence[ToolCall]:
         if not tool_calls:
             return []
@@ -225,7 +322,7 @@ class Agent:
                 results.append(tool_call)
         return results
 
-    async def handle_tool_calls(self, tool_calls: Sequence[ToolCall] | None, context: Any | None = None) -> AsyncGenerator[ToolCallChunk | ToolCallResultChunk, None]:  # noqa: ANN401
+    async def handle_tool_calls(self, tool_calls: Sequence[ToolCall] | None, context: Any | None = None) -> AsyncGenerator[FunctionCallEvent | FunctionCallOutputEvent, None]:  # noqa: ANN401
         if not tool_calls:
             return
         if tool_calls:
@@ -236,26 +333,31 @@ class Agent:
                     continue
 
             for tool_call in tool_calls:
+                yield FunctionCallEvent(
+                    call_id=tool_call.id,
+                    name=tool_call.function.name,
+                    arguments=tool_call.function.arguments or "",
+                )
+                start_time = time.time()
                 try:
-                    yield ToolCallChunk(
-                        type="tool_call",
-                        name=tool_call.function.name,
-                        arguments=tool_call.function.arguments or "",
-                    )
                     content = await self.fc.call_function_async(tool_call.function.name, tool_call.function.arguments or "", context)
-                    yield ToolCallResultChunk(
-                        type="tool_call_result",
+                    end_time = time.time()
+                    execution_time_ms = int((end_time - start_time) * 1000)
+                    yield FunctionCallOutputEvent(
                         tool_call_id=tool_call.id,
                         name=tool_call.function.name,
                         content=str(content),
+                        execution_time_ms=execution_time_ms,
                     )
-                except Exception as e:  # noqa: PERF203
+                except Exception as e:
                     logger.exception("Tool call %s failed", tool_call.id)
-                    yield ToolCallResultChunk(
-                        type="tool_call_result",
+                    end_time = time.time()
+                    execution_time_ms = int((end_time - start_time) * 1000)
+                    yield FunctionCallOutputEvent(
                         tool_call_id=tool_call.id,
                         name=tool_call.function.name,
                         content=str(e),
+                        execution_time_ms=execution_time_ms,
                     )
 
     def _convert_responses_to_completions_format(self, messages: RunnerMessages) -> list[dict]:
@@ -265,7 +367,7 @@ class Agent:
 
         while i < len(messages):
             message = messages[i]
-            message_dict = message.model_dump() if isinstance(message, BaseModel) else message
+            message_dict = message_to_llm_dict(message) if isinstance(message, (NewUserMessage, NewSystemMessage, NewAssistantMessage)) else message
 
             message_type = message_dict.get("type")
             role = message_dict.get("role")
@@ -277,11 +379,11 @@ class Agent:
 
                 while j < len(messages):
                     next_message = messages[j]
-                    next_dict = next_message.model_dump() if isinstance(next_message, BaseModel) else next_message
+                    next_dict = message_to_llm_dict(next_message) if isinstance(next_message, (NewUserMessage, NewSystemMessage, NewAssistantMessage)) else next_message
 
                     if next_dict.get("type") == "function_call":
                         tool_call = {
-                            "id": next_dict["function_call_id"],  # type: ignore
+                            "id": next_dict["call_id"],  # type: ignore
                             "type": "function",
                             "function": {
                                 "name": next_dict["name"],  # type: ignore
@@ -340,44 +442,52 @@ class Agent:
 
         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)
+            # Convert Pydantic objects to dict first
+            if hasattr(item, "model_dump"):
+                item_dict = item.model_dump()
+            elif hasattr(item, "dict"):  # For older Pydantic versions
+                item_dict = item.dict()
+            elif isinstance(item, dict):
+                item_dict = item
             else:
                 # Handle non-dict items (shouldn't happen, but just in case)
                 converted_content.append(item)
+                continue
+
+            item_type = item_dict.get("type")
+            if item_type in ["input_text", "text"]:
+                # Convert ResponseInputText or new text format to completion API format
+                converted_content.append(
+                    {
+                        "type": "text",
+                        "text": item_dict["text"],
+                    },
+                )
+            elif item_type in ["input_image", "image"]:
+                # Convert ResponseInputImage to completion API format
+                if item_dict.get("file_id"):
+                    msg = "File ID input is not supported for Completion API"
+                    raise ValueError(msg)
+
+                if not item_dict.get("image_url"):
+                    msg = "ResponseInputImage must have either file_id or image_url"
+                    raise ValueError(msg)
+
+                # Build image_url object with detail inside
+                image_data = {"url": item_dict["image_url"]}
+                detail = item_dict.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_dict)
 
         return converted_content