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

lite_agent/__init__.py

@@ -0,0 +1,7 @@
+"""Lite Agent - A lightweight AI agent framework."""
+
+from .agent import Agent
+from .message_transfers import consolidate_history_transfer
+from .runner import Runner
+
+__all__ = ["Agent", "Runner", "consolidate_history_transfer"]

lite_agent/agent.py

@@ -1,36 +1,330 @@
-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 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.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.
+
+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>"""
 
 
 class Agent:
-    def __init__(self, *, model: str, name: str, instructions: str, tools: list[Callable] | None = None) -> None:
+    def __init__(  # noqa: PLR0913
+        self,
+        *,
+        model: str,
+        name: str,
+        instructions: str,
+        tools: list[Callable] | None = None,
+        handoffs: list["Agent"] | None = None,
+        message_transfer: Callable[[RunnerMessages], RunnerMessages] | None = None,
+    ) -> None:
         self.name = name
         self.instructions = instructions
-        self.fc = Funcall(tools)
         self.model = model
+        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)
 
-    def prepare_messages(self, messages: RunnerMessages) -> list[dict]:
-        return [
-            {
-                "role": "system",
-                "content": f"You are {self.name}. {self.instructions}",
+        # 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 + "\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
+
+        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")
+    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 litellm.acompletion(
             model=self.model,
             messages=self.message_histories,
             tools=tools,
-            tool_choice="auto",
+            tool_choice="auto",  # TODO: make this configurable
             stream=True,
         )
-        return chunk_handler(resp, self.fc)
+
+        # 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": "function",
+                            "function": {
+                                "name": next_dict["name"],
+                                "arguments": next_dict["arguments"],
+                            },
+                            "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
+
+                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"],
+                        "content": message_dict["output"],
+                    },
+                )
+                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_messages.append(message_dict)
+                i += 1
+
+        return converted_messages
+
+    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

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"]

lite_agent/processors/stream_chunk_processor.py

@@ -1,53 +1,61 @@
 import litellm
-from funcall import Funcall
-from litellm.types.utils import ChatCompletionDeltaToolCall, StreamingChoices
+from litellm.types.utils import ChatCompletionDeltaToolCall, ModelResponseStream, StreamingChoices
 
-from open_agents.loggers import logger
-from open_agents.types import AssistantMessage, ToolCall, ToolCallFunction
+from lite_agent.loggers import logger
+from lite_agent.types import AssistantMessage, ToolCall, ToolCallFunction
 
 
 class StreamChunkProcessor:
     """Processor for handling streaming responses"""
 
-    def __init__(self, fc: Funcall) -> None:
-        self.fc = fc
-        self.current_message: AssistantMessage = None
+    def __init__(self) -> None:
+        self._current_message: AssistantMessage | None = None
 
-    def initialize_message(self, chunk: litellm.ModelResponseStream, choice: StreamingChoices) -> None:
+    def initialize_message(self, chunk: ModelResponseStream, choice: StreamingChoices) -> None:
         """Initialize the message object"""
         delta = choice.delta
-        self.current_message = AssistantMessage(
+        if delta.role != "assistant":
+            logger.warning("Skipping chunk with role: %s", delta.role)
+            return
+        self._current_message = AssistantMessage(
             id=chunk.id,
             index=choice.index,
             role=delta.role,
             content="",
         )
-        logger.debug("Initialized new message: %s", self.current_message.id)
+        logger.debug('Initialized new message: "%s"', self._current_message.id)
 
     def update_content(self, content: str) -> None:
         """Update message content"""
-        if self.current_message and content:
-            self.current_message.content += content
+        if self._current_message and content:
+            self._current_message.content += content
 
     def _initialize_tool_calls(self, tool_calls: list[litellm.ChatCompletionMessageToolCall]) -> None:
         """Initialize tool calls"""
-        if not self.current_message:
+        if not self._current_message:
             return
 
-        self.current_message.tool_calls = []
+        self._current_message.tool_calls = []
         for call in tool_calls:
             logger.debug("Create new tool call: %s", call.id)
 
     def _update_tool_calls(self, tool_calls: list[litellm.ChatCompletionMessageToolCall]) -> None:
         """Update existing tool calls"""
-        if not self.current_message or not self.current_message.tool_calls:
+        if not self._current_message:
             return
-
-        for current_call, new_call in zip(self.current_message.tool_calls, tool_calls, strict=False):
-            if new_call.function.arguments:
+        if not hasattr(self._current_message, "tool_calls"):
+            self._current_message.tool_calls = []
+        if not self._current_message.tool_calls:
+            return
+        if not tool_calls:
+            return
+        for current_call, new_call in zip(self._current_message.tool_calls, tool_calls, strict=False):
+            if new_call.function.arguments and current_call.function.arguments:
                 current_call.function.arguments += new_call.function.arguments
-            if new_call.type:
+            if new_call.type and new_call.type == "function":
                 current_call.type = new_call.type
+            elif new_call.type:
+                logger.warning("Unexpected tool call type: %s", new_call.type)
 
     def update_tool_calls(self, tool_calls: list[ChatCompletionDeltaToolCall]) -> None:
         """Handle tool call updates"""
@@ -55,31 +63,44 @@ class StreamChunkProcessor:
             return
         for call in tool_calls:
             if call.id:
-                new_tool_call = ToolCall(
-                    id=call.id,
-                    type=call.type,
-                    function=ToolCallFunction(
-                        name=call.function.name or "",
-                        arguments=call.function.arguments,
-                    ),
-                    index=call.index,
-                )
-                if self.current_message.tool_calls is None:
-                    self.current_message.tool_calls = []
-                self.current_message.tool_calls.append(new_tool_call)
-            else:
-                existing_call = self.current_message.tool_calls[call.index]
+                if call.type == "function":
+                    new_tool_call = ToolCall(
+                        id=call.id,
+                        type=call.type,
+                        function=ToolCallFunction(
+                            name=call.function.name or "",
+                            arguments=call.function.arguments,
+                        ),
+                        index=call.index,
+                    )
+                    if self._current_message is not None:
+                        if self._current_message.tool_calls is None:
+                            self._current_message.tool_calls = []
+                        self._current_message.tool_calls.append(new_tool_call)
+                else:
+                    logger.warning("Unexpected tool call type: %s", call.type)
+            elif self._current_message is not None and self._current_message.tool_calls is not None and call.index is not None and 0 <= call.index < len(self._current_message.tool_calls):
+                existing_call = self._current_message.tool_calls[call.index]
                 if call.function.arguments:
+                    if existing_call.function.arguments is None:
+                        existing_call.function.arguments = ""
                     existing_call.function.arguments += call.function.arguments
+            else:
+                logger.warning("Cannot update tool call: current_message or tool_calls is None, or invalid index.")
 
-    def handle_usage_info(self, chunk: litellm.ModelResponseStream) -> litellm.Usage | None:
+    def handle_usage_info(self, chunk: ModelResponseStream) -> litellm.Usage | None:
         """Handle usage info, return whether this chunk should be skipped"""
-        usage = getattr(chunk, "usage", None)
-        if usage:
-            logger.debug("Model usage: %s", usage)
-        return usage
+        return getattr(chunk, "usage", None)
+
+    @property
+    def is_initialized(self) -> bool:
+        """Check if the current message is initialized"""
+        return self._current_message is not None
 
-    def finalize_message(self) -> AssistantMessage:
-        """Finalize message processing"""
-        logger.debug("Message finalized: %s", self.current_message)
-        return self.current_message
+    @property
+    def current_message(self) -> AssistantMessage:
+        """Get the current message being processed"""
+        if not self._current_message:
+            msg = "No current message initialized. Call initialize_message first."
+            raise ValueError(msg)
+        return self._current_message