lite-agent 0.6.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

lite_agent/agent.py

@@ -7,9 +7,21 @@ from funcall import Funcall
 from jinja2 import Environment, FileSystemLoader
 
 from lite_agent.client import BaseLLMClient, LiteLLMClient, ReasoningConfig
+from lite_agent.constants import CompletionMode, ToolName
 from lite_agent.loggers import logger
 from lite_agent.response_handlers import CompletionResponseHandler, ResponsesAPIHandler
-from lite_agent.types import AgentChunk, FunctionCallEvent, FunctionCallOutputEvent, RunnerMessages, ToolCall, message_to_llm_dict, system_message_to_llm_dict
+from lite_agent.types import (
+    AgentChunk,
+    AssistantTextContent,
+    AssistantToolCall,
+    AssistantToolCallResult,
+    FunctionCallEvent,
+    FunctionCallOutputEvent,
+    RunnerMessages,
+    ToolCall,
+    message_to_llm_dict,
+    system_message_to_llm_dict,
+)
 from lite_agent.types.messages import NewAssistantMessage, NewSystemMessage, NewUserMessage
 
 TEMPLATES_DIR = Path(__file__).parent / "templates"
@@ -32,10 +44,24 @@ class Agent:
         message_transfer: Callable[[RunnerMessages], RunnerMessages] | None = None,
         completion_condition: str = "stop",
         reasoning: ReasoningConfig = None,
+        stop_before_tools: list[str] | list[Callable] | None = None,
     ) -> None:
         self.name = name
         self.instructions = instructions
         self.reasoning = reasoning
+        # Convert stop_before_functions to function names
+        if stop_before_tools:
+            self.stop_before_functions = set()
+            for func in stop_before_tools:
+                if isinstance(func, str):
+                    self.stop_before_functions.add(func)
+                elif callable(func):
+                    self.stop_before_functions.add(func.__name__)
+                else:
+                    msg = f"stop_before_functions must contain strings or callables, got {type(func)}"
+                    raise TypeError(msg)
+        else:
+            self.stop_before_functions = set()
 
         if isinstance(model, BaseLLMClient):
             # If model is a BaseLLMClient instance, use it directly
@@ -54,7 +80,7 @@ class Agent:
         self.fc = Funcall(tools)
 
         # Add wait_for_user tool if completion condition is "call"
-        if completion_condition == "call":
+        if completion_condition == CompletionMode.CALL:
             self._add_wait_for_user_tool()
 
         # Set parent for handoff agents
@@ -99,7 +125,7 @@ class Agent:
 
         # Add single dynamic tool for all transfers
         self.fc.add_dynamic_tool(
-            name="transfer_to_agent",
+            name=ToolName.TRANSFER_TO_AGENT,
             description="Transfer conversation to another agent.",
             parameters={
                 "name": {
@@ -129,7 +155,7 @@ class Agent:
 
         # Add dynamic tool for parent transfer
         self.fc.add_dynamic_tool(
-            name="transfer_to_parent",
+            name=ToolName.TRANSFER_TO_PARENT,
             description="Transfer conversation back to parent agent when current task is completed or cannot be solved by current agent",
             parameters={},
             required=[],
@@ -160,7 +186,7 @@ class Agent:
             try:
                 # Try to remove the existing transfer tool
                 if hasattr(self.fc, "remove_dynamic_tool"):
-                    self.fc.remove_dynamic_tool("transfer_to_agent")
+                    self.fc.remove_dynamic_tool(ToolName.TRANSFER_TO_AGENT)
             except Exception as e:
                 # If removal fails, log and continue anyway
                 logger.debug(f"Failed to remove existing transfer tool: {e}")
@@ -205,31 +231,30 @@ class Agent:
         for message in messages:
             if isinstance(message, NewAssistantMessage):
                 for item in message.content:
-                    match item.type:
-                        case "text":
-                            res.append(
-                                {
-                                    "role": "assistant",
-                                    "content": item.text,
-                                },
-                            )
-                        case "tool_call":
-                            res.append(
-                                {
-                                    "type": "function_call",
-                                    "call_id": item.call_id,
-                                    "name": item.name,
-                                    "arguments": item.arguments,
-                                },
-                            )
-                        case "tool_call_result":
-                            res.append(
-                                {
-                                    "type": "function_call_output",
-                                    "call_id": item.call_id,
-                                    "output": item.output,
-                                },
-                            )
+                    if isinstance(item, AssistantTextContent):
+                        res.append(
+                            {
+                                "role": "assistant",
+                                "content": item.text,
+                            },
+                        )
+                    elif isinstance(item, AssistantToolCall):
+                        res.append(
+                            {
+                                "type": "function_call",
+                                "call_id": item.call_id,
+                                "name": item.name,
+                                "arguments": item.arguments,
+                            },
+                        )
+                    elif isinstance(item, AssistantToolCallResult):
+                        res.append(
+                            {
+                                "type": "function_call_output",
+                                "call_id": item.call_id,
+                                "output": item.output,
+                            },
+                        )
             elif isinstance(message, NewSystemMessage):
                 res.append(
                     {
@@ -269,9 +294,6 @@ class Agent:
                         "content": contents,
                     },
                 )
-            # Handle dict messages (legacy format)
-            elif isinstance(message, dict):
-                res.append(message)
         return res
 
     async def completion(
@@ -279,6 +301,7 @@ class Agent:
         messages: RunnerMessages,
         record_to_file: Path | None = None,
         reasoning: ReasoningConfig = None,
+        *,
         streaming: bool = True,
     ) -> AsyncGenerator[AgentChunk, None]:
         # Apply message transfer callback if provided - always use legacy format for LLM compatibility
@@ -301,13 +324,14 @@ class Agent:
 
         # Use response handler for unified processing
         handler = CompletionResponseHandler()
-        return handler.handle(resp, streaming, record_to_file)
+        return handler.handle(resp, streaming=streaming, record_to=record_to_file)
 
     async def responses(
         self,
         messages: RunnerMessages,
         record_to_file: Path | None = None,
         reasoning: ReasoningConfig = None,
+        *,
         streaming: bool = True,
     ) -> AsyncGenerator[AgentChunk, None]:
         # Apply message transfer callback if provided - always use legacy format for LLM compatibility
@@ -328,20 +352,29 @@ class Agent:
         )
         # Use response handler for unified processing
         handler = ResponsesAPIHandler()
-        return handler.handle(resp, streaming, record_to_file)
+        return handler.handle(resp, streaming=streaming, record_to=record_to_file)
 
     async def list_require_confirm_tools(self, tool_calls: Sequence[ToolCall] | None) -> Sequence[ToolCall]:
         if not tool_calls:
             return []
         results = []
         for tool_call in tool_calls:
-            tool_func = self.fc.function_registry.get(tool_call.function.name)
+            function_name = tool_call.function.name
+
+            # Check if function is in dynamic stop_before_functions list
+            if function_name in self.stop_before_functions:
+                logger.debug('Tool call "%s" requires confirmation (stop_before_functions)', tool_call.id)
+                results.append(tool_call)
+                continue
+
+            # Check decorator-based require_confirmation
+            tool_func = self.fc.function_registry.get(function_name)
             if not tool_func:
-                logger.warning("Tool function %s not found in registry", tool_call.function.name)
+                logger.warning("Tool function %s not found in registry", function_name)
                 continue
-            tool_meta = self.fc.get_tool_meta(tool_call.function.name)
+            tool_meta = self.fc.get_tool_meta(function_name)
             if tool_meta["require_confirm"]:
-                logger.debug('Tool call "%s" requires confirmation', tool_call.id)
+                logger.debug('Tool call "%s" requires confirmation (decorator)', tool_call.id)
                 results.append(tool_call)
         return results
 
@@ -396,10 +429,79 @@ class Agent:
             role = message_dict.get("role")
 
             if role == "assistant":
-                # Look ahead for function_call messages
+                # For NewAssistantMessage, extract directly from the message object
                 tool_calls = []
+                tool_results = []
+
+                if isinstance(message, NewAssistantMessage):
+                    # Process content directly from NewAssistantMessage
+                    for item in message.content:
+                        if item.type == "tool_call":
+                            tool_call = {
+                                "id": item.call_id,
+                                "type": "function",
+                                "function": {
+                                    "name": item.name,
+                                    "arguments": item.arguments,
+                                },
+                                "index": len(tool_calls),
+                            }
+                            tool_calls.append(tool_call)
+                        elif item.type == "tool_call_result":
+                            # Collect tool call results to be added as separate tool messages
+                            tool_results.append({
+                                "call_id": item.call_id,
+                                "output": item.output,
+                            })
+
+                    # Create assistant message with only text content and tool calls
+                    text_content = " ".join([item.text for item in message.content if item.type == "text"])
+                    message_dict = {
+                        "role": "assistant",
+                        "content": text_content if text_content else None,
+                    }
+                    if tool_calls:
+                        message_dict["tool_calls"] = tool_calls
+                else:
+                    # Legacy handling for dict messages
+                    content = message_dict.get("content", [])
+                    # Handle both string and array content
+                    if isinstance(content, list):
+                        # Extract tool_calls and tool_call_results from content array and filter out non-text content
+                        filtered_content = []
+                        for item in content:
+                            if isinstance(item, dict):
+                                if item.get("type") == "tool_call":
+                                    tool_call = {
+                                        "id": item.get("call_id", ""),
+                                        "type": "function",
+                                        "function": {
+                                            "name": item.get("name", ""),
+                                            "arguments": item.get("arguments", "{}"),
+                                        },
+                                        "index": len(tool_calls),
+                                    }
+                                    tool_calls.append(tool_call)
+                                elif item.get("type") == "tool_call_result":
+                                    # Collect tool call results to be added as separate tool messages
+                                    tool_results.append({
+                                        "call_id": item.get("call_id", ""),
+                                        "output": item.get("output", ""),
+                                    })
+                                elif item.get("type") == "text":
+                                    filtered_content.append(item)
+
+                        # Update content to only include text items
+                        if filtered_content:
+                            message_dict = message_dict.copy()
+                            message_dict["content"] = filtered_content
+                        elif tool_calls:
+                            # If we have tool_calls but no text content, set content to None per OpenAI API spec
+                            message_dict = message_dict.copy()
+                            message_dict["content"] = None
+
+                # Look ahead for function_call messages (legacy support)
                 j = i + 1
-
                 while j < len(messages):
                     next_message = messages[j]
                     next_dict = message_to_llm_dict(next_message) if isinstance(next_message, (NewUserMessage, NewSystemMessage, NewAssistantMessage)) else next_message
@@ -419,12 +521,33 @@ class Agent:
                     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
+                # For legacy dict messages, create assistant message with tool_calls if any
+                if not isinstance(message, NewAssistantMessage):
+                    assistant_msg = message_dict.copy()
+                    if tool_calls:
+                        assistant_msg["tool_calls"] = tool_calls  # type: ignore
+
+                    # Convert content format for OpenAI API compatibility
+                    content = assistant_msg.get("content", [])
+                    if isinstance(content, list):
+                        # Extract text content and convert to string using list comprehension
+                        text_parts = [item.get("text", "") for item in content if isinstance(item, dict) and item.get("type") == "text"]
+                        assistant_msg["content"] = " ".join(text_parts) if text_parts else None
+
+                    message_dict = assistant_msg
+
+                converted_messages.append(message_dict)
+
+                # Add tool messages for any tool_call_results found in the assistant message
+                converted_messages.extend([
+                    {
+                        "role": "tool",
+                        "tool_call_id": tool_result["call_id"],
+                        "content": tool_result["output"],
+                    }
+                    for tool_result in tool_results
+                ])
 
-                converted_messages.append(assistant_msg)
                 i = j  # Skip the function_call messages we've processed
 
             elif message_type == "function_call_output":
@@ -536,10 +659,73 @@ class Agent:
 
         # Add dynamic tool for task completion
         self.fc.add_dynamic_tool(
-            name="wait_for_user",
+            name=ToolName.WAIT_FOR_USER,
             description="Call this function when you have completed your assigned task or need more information from the user.",
             parameters={},
             required=[],
             handler=wait_for_user_handler,
         )
 
+    def set_stop_before_functions(self, functions: list[str] | list[Callable]) -> None:
+        """Set the list of functions that require confirmation before execution.
+
+        Args:
+            functions: List of function names (str) or callable objects
+        """
+        self.stop_before_functions = set()
+        for func in functions:
+            if isinstance(func, str):
+                self.stop_before_functions.add(func)
+            elif callable(func):
+                self.stop_before_functions.add(func.__name__)
+            else:
+                msg = f"stop_before_functions must contain strings or callables, got {type(func)}"
+                raise TypeError(msg)
+        logger.debug(f"Set stop_before_functions to: {self.stop_before_functions}")
+
+    def add_stop_before_function(self, function: str | Callable) -> None:
+        """Add a function to the stop_before_functions list.
+
+        Args:
+            function: Function name (str) or callable object to add
+        """
+        if isinstance(function, str):
+            function_name = function
+        elif callable(function):
+            function_name = function.__name__
+        else:
+            msg = f"function must be a string or callable, got {type(function)}"
+            raise TypeError(msg)
+
+        self.stop_before_functions.add(function_name)
+        logger.debug(f"Added '{function_name}' to stop_before_functions")
+
+    def remove_stop_before_function(self, function: str | Callable) -> None:
+        """Remove a function from the stop_before_functions list.
+
+        Args:
+            function: Function name (str) or callable object to remove
+        """
+        if isinstance(function, str):
+            function_name = function
+        elif callable(function):
+            function_name = function.__name__
+        else:
+            msg = f"function must be a string or callable, got {type(function)}"
+            raise TypeError(msg)
+
+        self.stop_before_functions.discard(function_name)
+        logger.debug(f"Removed '{function_name}' from stop_before_functions")
+
+    def clear_stop_before_functions(self) -> None:
+        """Clear all function names from the stop_before_functions list."""
+        self.stop_before_functions.clear()
+        logger.debug("Cleared all stop_before_functions")
+
+    def get_stop_before_functions(self) -> set[str]:
+        """Get the current set of function names that require confirmation.
+
+        Returns:
+            Set of function names
+        """
+        return self.stop_before_functions.copy()