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

lite_agent/__init__.py

@@ -0,0 +1,8 @@
+"""Lite Agent - A lightweight AI agent framework."""
+
+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", "print_chat_history", "print_chat_summary"]

lite_agent/agent.py

@@ -1,36 +1,411 @@
-from collections.abc import AsyncGenerator, Callable
+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 open_agents.chunk_handler import AgentChunk, chunk_handler
-from open_agents.types import RunnerMessages
+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
+
+TEMPLATES_DIR = Path(__file__).parent / "templates"
+jinja_env = Environment(loader=FileSystemLoader(str(TEMPLATES_DIR)), autoescape=True)
+
+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__(self, *, model: str, name: str, instructions: str, tools: list[Callable] | None = None) -> None:
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        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
+        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)
-        self.model = model
 
-    def prepare_messages(self, messages: RunnerMessages) -> list[dict]:
-        return [
-            {
-                "role": "system",
-                "content": f"You are {self.name}. {self.instructions}",
+        # 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:
+                handoff_agent.parent = self
+            self._add_transfer_tools(handoffs)
+
+        # Add transfer_to_parent tool if this agent has a parent (for cases where parent is set externally)
+        if self.parent is not None:
+            self.add_transfer_to_parent_tool()
+
+    @property
+    def parent(self) -> Optional["Agent"]:
+        return self._parent
+
+    @parent.setter
+    def parent(self, value: Optional["Agent"]) -> None:
+        self._parent = value
+        if value is not None:
+            self.add_transfer_to_parent_tool()
+
+    def _add_transfer_tools(self, handoffs: list["Agent"]) -> None:
+        """Add transfer function for handoff agents using dynamic tools.
+
+        Creates a single 'transfer_to_agent' function that accepts a 'name' parameter
+        to specify which agent to transfer the conversation to.
+
+        Args:
+            handoffs: List of Agent objects that can be transferred to
+        """
+        # Collect all agent names for validation
+        agent_names = [agent.name for agent in handoffs]
+
+        def transfer_handler(name: str) -> str:
+            """Handler for transfer_to_agent function."""
+            if name in agent_names:
+                return f"Transferring to agent: {name}"
+
+            available_agents = ", ".join(agent_names)
+            return f"Agent '{name}' not found. Available agents: {available_agents}"
+
+        # Add single dynamic tool for all transfers
+        self.fc.add_dynamic_tool(
+            name="transfer_to_agent",
+            description="Transfer conversation to another agent.",
+            parameters={
+                "name": {
+                    "type": "string",
+                    "description": "The name of the agent to transfer to",
+                    "enum": agent_names,
+                },
             },
-            *messages,
+            required=["name"],
+            handler=transfer_handler,
+        )
+
+    def add_transfer_to_parent_tool(self) -> None:
+        """Add transfer_to_parent function for agents that have a parent.
+
+        This tool allows the agent to transfer back to its parent when:
+        - The current task is completed
+        - The agent cannot solve the current problem
+        - Escalation to a higher level is needed
+        """
+
+        def transfer_to_parent_handler() -> str:
+            """Handler for transfer_to_parent function."""
+            if self.parent:
+                return f"Transferring back to parent agent: {self.parent.name}"
+            return "No parent agent found"
+
+        # Add dynamic tool for parent transfer
+        self.fc.add_dynamic_tool(
+            name="transfer_to_parent",
+            description="Transfer conversation back to parent agent when current task is completed or cannot be solved by current agent",
+            parameters={},
+            required=[],
+            handler=transfer_to_parent_handler,
+        )
+
+    def add_handoff(self, agent: "Agent") -> None:
+        """Add a handoff agent after initialization.
+
+        This method allows adding handoff agents dynamically after the agent
+        has been constructed. It properly sets up parent-child relationships
+        and updates the transfer tools.
+
+        Args:
+            agent: The agent to add as a handoff target
+        """
+        # Add to handoffs list if not already present
+        if agent not in self.handoffs:
+            self.handoffs.append(agent)
+
+            # Set parent relationship
+            agent.parent = self
+
+            # Add transfer_to_parent tool to the handoff agent
+            agent.add_transfer_to_parent_tool()
+
+            # Remove existing transfer tool if it exists and recreate with all agents
+            try:
+                # Try to remove the existing transfer tool
+                if hasattr(self.fc, "remove_dynamic_tool"):
+                    self.fc.remove_dynamic_tool("transfer_to_agent")
+            except Exception as e:
+                # If removal fails, log and continue anyway
+                logger.debug(f"Failed to remove existing transfer tool: {e}")
+
+            # 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
+        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",
+                content=f"You are {self.name}. {instructions}",
+            ).model_dump(),
+            *converted_messages,
         ]
 
-    async def stream_async(self, messages: RunnerMessages) -> AsyncGenerator[AgentChunk, None]:
-        self.message_histories = self.prepare_messages(messages)
-        tools = self.fc.get_tools(target="litellm")
-        resp = await litellm.acompletion(
-            model=self.model,
+    async def completion(self, messages: RunnerMessages, record_to_file: Path | None = None) -> AsyncGenerator[AgentChunk, None]:
+        # Apply message transfer callback if provided
+        processed_messages = messages
+        if self.message_transfer:
+            logger.debug(f"Applying message transfer callback for agent {self.name}")
+            processed_messages = self.message_transfer(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,
             tools=tools,
-            tool_choice="auto",
-            stream=True,
+            tool_choice="auto",  # TODO: make this configurable
+        )
+
+        # Ensure resp is a CustomStreamWrapper
+        if isinstance(resp, CustomStreamWrapper):
+            return litellm_stream_handler(resp, record_to=record_to_file)
+        msg = "Response is not a CustomStreamWrapper, cannot stream chunks."
+        raise TypeError(msg)
+
+    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)
+            if not tool_func:
+                logger.warning("Tool function %s not found in registry", tool_call.function.name)
+                continue
+            tool_meta = self.fc.get_tool_meta(tool_call.function.name)
+            if tool_meta["require_confirm"]:
+                logger.debug('Tool call "%s" requires confirmation', tool_call.id)
+                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
+        if not tool_calls:
+            return
+        if tool_calls:
+            for tool_call in tool_calls:
+                tool_func = self.fc.function_registry.get(tool_call.function.name)
+                if not tool_func:
+                    logger.warning("Tool function %s not found in registry", tool_call.function.name)
+                    continue
+
+            for tool_call in tool_calls:
+                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",
+                        tool_call_id=tool_call.id,
+                        name=tool_call.function.name,
+                        content=str(content),
+                    )
+                except Exception as e:  # noqa: PERF203
+                    logger.exception("Tool call %s failed", tool_call.id)
+                    yield ToolCallResultChunk(
+                        type="tool_call_result",
+                        tool_call_id=tool_call.id,
+                        name=tool_call.function.name,
+                        content=str(e),
+                    )
+
+    def _convert_responses_to_completions_format(self, messages: RunnerMessages) -> list[dict]:
+        """Convert messages from responses API format to completions API format."""
+        converted_messages = []
+        i = 0
+
+        while i < len(messages):
+            message = messages[i]
+            message_dict = message.model_dump() if isinstance(message, BaseModel) else message
+
+            message_type = message_dict.get("type")
+            role = message_dict.get("role")
+
+            if role == "assistant":
+                # Look ahead for function_call messages
+                tool_calls = []
+                j = i + 1
+
+                while j < len(messages):
+                    next_message = messages[j]
+                    next_dict = next_message.model_dump() if isinstance(next_message, BaseModel) else next_message
+
+                    if next_dict.get("type") == "function_call":
+                        tool_call = {
+                            "id": next_dict["function_call_id"],  # type: ignore
+                            "type": "function",
+                            "function": {
+                                "name": next_dict["name"],  # type: ignore
+                                "arguments": next_dict["arguments"],  # type: ignore
+                            },
+                            "index": len(tool_calls),
+                        }
+                        tool_calls.append(tool_call)
+                        j += 1
+                    else:
+                        break
+
+                # Create assistant message with tool_calls if any
+                assistant_msg = message_dict.copy()
+                if 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
+
+            elif message_type == "function_call_output":
+                # Convert to tool message
+                converted_messages.append(
+                    {
+                        "role": "tool",
+                        "tool_call_id": message_dict["call_id"],  # type: ignore
+                        "content": message_dict["output"],  # type: ignore
+                    },
+                )
+                i += 1
+
+            elif message_type == "function_call":
+                # This should have been processed with the assistant message
+                # Skip it if we encounter it standalone
+                i += 1
+
+            else:
+                # Regular message (user, system)
+                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.
+
+        Args:
+            message_transfer: A callback function that takes RunnerMessages as input
+                             and returns RunnerMessages as output. This function will be
+                             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,
         )
-        return chunk_handler(resp, self.fc)

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/loggers.py

@@ -1,3 +1,3 @@
 import logging
 
-logger = logging.getLogger("easy_agent")
+logger = logging.getLogger("lite_agent")

lite_agent/message_transfers.py

@@ -0,0 +1,111 @@
+"""
+Predefined message transfer functions for lite-agent.
+
+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
+
+
+def consolidate_history_transfer(messages: RunnerMessages) -> RunnerMessages:
+    """Consolidate all message history into a single user message with XML format.
+
+    This message transfer function converts all message history into XML format
+    and creates a single user message asking what to do next. This is useful when
+    you want to summarize the entire conversation context in a single prompt.
+
+    Args:
+        messages: The original messages to be processed
+
+    Returns:
+        A single user message containing the consolidated history in XML format
+
+    Example:
+        >>> agent = Agent(
+        ...     model="gpt-4",
+        ...     name="HistoryAgent",
+        ...     instructions="You are a helpful assistant.",
+        ...     message_transfer=consolidate_history_transfer
+        ... )
+    """
+    if not messages:
+        return messages
+
+    # Convert messages to XML format
+    xml_content = ["<conversation_history>"]
+
+    for message in messages:
+        xml_content.extend(_process_message_to_xml(message))
+
+    xml_content.append("</conversation_history>")
+
+    # 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}]
+
+
+def _process_message_to_xml(message: dict | object) -> list[str]:
+    """Process a single message and convert it to XML format.
+
+    Args:
+        message: A single message to process
+
+    Returns:
+        List of XML strings representing the message
+    """
+    xml_lines = []
+
+    if isinstance(message, dict):
+        xml_lines.extend(_process_dict_message(message))
+    elif hasattr(message, "role"):
+        # Handle Pydantic model format messages
+        role = getattr(message, "role", "unknown")
+        content = getattr(message, "content", "")
+        if isinstance(content, str):
+            xml_lines.append(f"  <message role='{role}'>{content}</message>")
+    elif hasattr(message, "type"):
+        # Handle function call messages
+        xml_lines.extend(_process_function_message(message))
+
+    return xml_lines
+
+
+def _process_dict_message(message: dict) -> list[str]:
+    """Process dictionary format message to XML."""
+    xml_lines = []
+    role = message.get("role", "unknown")
+    content = message.get("content", "")
+    message_type = message.get("type")
+
+    if message_type == "function_call":
+        name = message.get("name", "unknown")
+        arguments = message.get("arguments", "")
+        xml_lines.append(f"  <function_call name='{name}' arguments='{arguments}' />")
+    elif message_type == "function_call_output":
+        call_id = message.get("call_id", "unknown")
+        output = message.get("output", "")
+        xml_lines.append(f"  <function_result call_id='{call_id}'>{output}</function_result>")
+    elif role in ["user", "assistant", "system"]:
+        xml_lines.append(f"  <message role='{role}'>{content}</message>")
+
+    return xml_lines
+
+
+def _process_function_message(message: dict | object) -> list[str]:
+    """Process function call message to XML."""
+    xml_lines = []
+    message_type = getattr(message, "type", "unknown")
+
+    if message_type == "function_call":
+        name = getattr(message, "name", "unknown")
+        arguments = getattr(message, "arguments", "")
+        xml_lines.append(f"  <function_call name='{name}' arguments='{arguments}' />")
+    elif message_type == "function_call_output":
+        call_id = getattr(message, "call_id", "unknown")
+        output = getattr(message, "output", "")
+        xml_lines.append(f"  <function_result call_id='{call_id}'>{output}</function_result>")
+
+    return xml_lines

lite_agent/processors/__init__.py

@@ -1,3 +1,3 @@
-from open_agents.processors.stream_chunk_processor import StreamChunkProcessor
+from lite_agent.processors.stream_chunk_processor import StreamChunkProcessor
 
 __all__ = ["StreamChunkProcessor"]