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

lite_agent/runner.py

@@ -1,5 +1,4 @@
 import json
-import warnings
 from collections.abc import AsyncGenerator, Sequence
 from datetime import datetime, timedelta, timezone
 from os import PathLike
@@ -28,8 +27,9 @@ from lite_agent.types import (
     UserInput,
     UserTextContent,
 )
-from lite_agent.types.events import AssistantMessageEvent, FunctionCallOutputEvent, TimingEvent
+from lite_agent.types.events import AssistantMessageEvent, FunctionCallOutputEvent
 from lite_agent.utils.message_builder import MessageBuilder
+from lite_agent.utils.message_state_manager import MessageStateManager
 
 
 class Runner:
@@ -38,50 +38,40 @@ class Runner:
         self.messages: list[FlexibleRunnerMessage] = []
         self.api = api
         self.streaming = streaming
-        self._current_assistant_message: NewAssistantMessage | None = None
+        self._message_state_manager = MessageStateManager()
         self.usage = MessageUsage(input_tokens=0, output_tokens=0, total_tokens=0)
 
-    def _start_assistant_message(self, content: str = "", meta: AssistantMessageMeta | None = None) -> None:
+    async def _start_assistant_message(self, content: str = "", meta: AssistantMessageMeta | None = None) -> None:
         """Start a new assistant message."""
         # Create meta with model information if not provided
         if meta is None:
             meta = AssistantMessageMeta()
             if hasattr(self.agent.client, "model"):
                 meta.model = self.agent.client.model
-        self._current_assistant_message = NewAssistantMessage(
-            content=[AssistantTextContent(text=content)],
-            meta=meta,
-        )
+        await self._message_state_manager.start_message(content, meta)
 
-    def _ensure_current_assistant_message(self) -> NewAssistantMessage:
+    async def _ensure_current_assistant_message(self) -> NewAssistantMessage:
         """Ensure current assistant message exists and return it."""
-        if self._current_assistant_message is None:
-            self._start_assistant_message()
-        if self._current_assistant_message is None:
-            msg = "Failed to create current assistant message"
-            raise RuntimeError(msg)
-        return self._current_assistant_message
-
-    def _add_to_current_assistant_message(self, content_item: AssistantTextContent | AssistantToolCall | AssistantToolCallResult) -> None:
-        """Add content to the current assistant message."""
-        self._ensure_current_assistant_message().content.append(content_item)
+        return await self._message_state_manager.ensure_message_exists()
 
-    def _add_text_content_to_current_assistant_message(self, delta: str) -> None:
+    async def _add_to_current_assistant_message(self, content_item: AssistantTextContent | AssistantToolCall | AssistantToolCallResult) -> None:
+        """Add content to the current assistant message."""
+        if isinstance(content_item, AssistantTextContent):
+            await self._message_state_manager.add_text_delta(content_item.text)
+        elif isinstance(content_item, AssistantToolCall):
+            await self._message_state_manager.add_tool_call(content_item)
+        elif isinstance(content_item, AssistantToolCallResult):
+            await self._message_state_manager.add_tool_result(content_item)
+
+    async def _add_text_content_to_current_assistant_message(self, delta: str) -> None:
         """Add text delta to the current assistant message's text content."""
-        message = self._ensure_current_assistant_message()
-        # Find the first text content item and append the delta
-        for content_item in message.content:
-            if content_item.type == "text":
-                content_item.text += delta
-                return
-        # If no text content found, add new text content
-        message.content.append(AssistantTextContent(text=delta))
+        await self._message_state_manager.add_text_delta(delta)
 
-    def _finalize_assistant_message(self) -> None:
+    async def _finalize_assistant_message(self) -> None:
         """Finalize the current assistant message and add it to messages."""
-        if self._current_assistant_message is not None:
-            self.messages.append(self._current_assistant_message)
-            self._current_assistant_message = None
+        finalized_message = await self._message_state_manager.finalize_message()
+        if finalized_message is not None:
+            self.messages.append(finalized_message)
 
     def _add_tool_call_result(self, call_id: str, output: str, execution_time_ms: int | None = None) -> None:
         """Add a tool call result to the last assistant message, or create a new one if needed."""
@@ -127,10 +117,12 @@ class Runner:
         # Check for transfer_to_agent calls first
         transfer_calls = [tc for tc in tool_calls if tc.function.name == ToolName.TRANSFER_TO_AGENT]
         if transfer_calls:
+            logger.info(f"Processing {len(transfer_calls)} transfer_to_agent calls")
             # Handle all transfer calls but only execute the first one
             for i, tool_call in enumerate(transfer_calls):
                 if i == 0:
                     # Execute the first transfer
+                    logger.info(f"Executing agent transfer: {tool_call.function.arguments}")
                     call_id, output = await self._handle_agent_transfer(tool_call)
                     # Generate function_call_output event if in includes
                     if "function_call_output" in includes:
@@ -205,8 +197,6 @@ class Runner:
                     last_message = cast("NewAssistantMessage", self.messages[-1])
                     last_message.content.append(tool_result)
 
-                # Note: For completion API compatibility, the conversion happens when sending to LLM
-
     async def _collect_all_chunks(self, stream: AsyncGenerator[AgentChunk, None]) -> list[AgentChunk]:
         """Collect all chunks from an async generator into a list."""
         return [chunk async for chunk in stream]
@@ -218,7 +208,6 @@ class Runner:
         includes: Sequence[AgentChunkType] | None = None,
         context: "Any | None" = None,  # noqa: ANN401
         record_to: PathLike | str | None = None,
-        agent_kwargs: dict[str, Any] | None = None,
     ) -> AsyncGenerator[AgentChunk, None]:
         """Run the agent and return a RunResponse object that can be asynchronously iterated for each chunk.
 
@@ -253,7 +242,7 @@ class Runner:
                 # Handle single message (BaseModel, TypedDict, or dict)
                 self.append_message(user_input)  # type: ignore[arg-type]
         logger.debug("Messages prepared, calling _run")
-        return self._run(max_steps, includes, self._normalize_record_path(record_to), context=context, agent_kwargs=agent_kwargs)
+        return self._run(max_steps, includes, self._normalize_record_path(record_to), context=context)
 
     async def _run(
         self,
@@ -261,7 +250,6 @@ class Runner:
         includes: Sequence[AgentChunkType],
         record_to: Path | None = None,
         context: Any | None = None,  # noqa: ANN401
-        agent_kwargs: dict[str, Any] | None = None,
     ) -> AsyncGenerator[AgentChunk, None]:
         """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}")
@@ -293,21 +281,13 @@ class Runner:
 
         while not is_finish() and steps < max_steps:
             logger.debug(f"Step {steps}: finish_reason={finish_reason}, is_finish()={is_finish()}")
-            # Convert to legacy format only when needed for LLM communication
-            # This allows us to keep the new format internally but ensures compatibility
-            # Extract agent kwargs for reasoning configuration
-            reasoning = None
-            if agent_kwargs:
-                reasoning = agent_kwargs.get("reasoning")
-
-            logger.debug(f"Using API: {self.api}, streaming: {self.streaming}")
+            logger.info(f"Making LLM request: API={self.api}, streaming={self.streaming}, messages={len(self.messages)}")
             match self.api:
                 case "completion":
                     logger.debug("Calling agent.completion")
                     resp = await self.agent.completion(
                         self.messages,
                         record_to_file=record_to,
-                        reasoning=reasoning,
                         streaming=self.streaming,
                     )
                 case "responses":
@@ -315,7 +295,6 @@ class Runner:
                     resp = await self.agent.responses(
                         self.messages,
                         record_to_file=record_to,
-                        reasoning=reasoning,
                         streaming=self.streaming,
                     )
                 case _:
@@ -331,48 +310,55 @@ class Runner:
                         logger.debug(f"Assistant message chunk: {len(chunk.message.content) if chunk.message.content else 0} content items")
                         # Start or update assistant message in new format
                         # If we already have a current assistant message, just update its metadata
-                        if self._current_assistant_message is not None:
+                        current_message = await self._message_state_manager.get_current_message()
+                        if current_message is not None:
                             # Preserve all existing metadata and only update specific fields
-                            original_meta = self._current_assistant_message.meta
-                            original_meta.sent_at = chunk.message.meta.sent_at
-                            if hasattr(chunk.message.meta, "latency_ms"):
-                                original_meta.latency_ms = chunk.message.meta.latency_ms
-                            if hasattr(chunk.message.meta, "output_time_ms"):
-                                original_meta.total_time_ms = chunk.message.meta.output_time_ms
+                            meta_updates = {"sent_at": chunk.message.meta.sent_at}
+                            # Only include fields of type datetime in meta_updates
+                            # Update int fields separately after update_meta
                             # Preserve other metadata fields like model, usage, etc.
                             for attr in ["model", "usage", "input_tokens", "output_tokens"]:
                                 if hasattr(chunk.message.meta, attr):
-                                    setattr(original_meta, attr, getattr(chunk.message.meta, attr))
+                                    meta_updates[attr] = getattr(chunk.message.meta, attr)
+                            await self._message_state_manager.update_meta(**meta_updates)
+                            # Now update int fields directly if present
+                            if hasattr(chunk.message.meta, "latency_ms"):
+                                await self._message_state_manager.update_meta(latency_ms=chunk.message.meta.latency_ms)
+                            if hasattr(chunk.message.meta, "output_time_ms"):
+                                await self._message_state_manager.update_meta(total_time_ms=chunk.message.meta.output_time_ms)
                         else:
-                            # For non-streaming mode, directly use the complete message from the response handler
-                            self._current_assistant_message = chunk.message
+                            # For non-streaming mode, start with complete message
+                            await self._start_assistant_message(meta=chunk.message.meta)
+                            # Add all content from the chunk message
+                            for content_item in chunk.message.content:
+                                await self._add_to_current_assistant_message(content_item)
 
                         # If model is None, try to get it from agent client
-                        if self._current_assistant_message is not None and self._current_assistant_message.meta.model is None and hasattr(self.agent.client, "model"):
-                            self._current_assistant_message.meta.model = self.agent.client.model
+                        current_message = await self._message_state_manager.get_current_message()
+                        if current_message is not None and current_message.meta.model is None and hasattr(self.agent.client, "model"):
+                            await self._message_state_manager.update_meta(model=self.agent.client.model)
                         # Only yield assistant_message chunk if it's in includes and has content
-                        if chunk.type in includes and self._current_assistant_message is not None:
+                        if chunk.type in includes and current_message is not None:
                             # Create a new chunk with the current assistant message content
                             updated_chunk = AssistantMessageEvent(
-                                message=self._current_assistant_message,
+                                message=current_message,
                             )
                             yield updated_chunk
                     case "content_delta":
                         # Accumulate text content to current assistant message
-                        self._add_text_content_to_current_assistant_message(chunk.delta)
+                        await self._add_text_content_to_current_assistant_message(chunk.delta)
                         # Always yield content_delta chunk if it's in includes
                         if chunk.type in includes:
                             yield chunk
                     case "function_call":
                         logger.debug(f"Function call: {chunk.name}({chunk.arguments or '{}'})")
                         # Add tool call to current assistant message
-                        # Keep arguments as string for compatibility with funcall library
                         tool_call = AssistantToolCall(
                             call_id=chunk.call_id,
                             name=chunk.name,
                             arguments=chunk.arguments or "{}",
                         )
-                        self._add_to_current_assistant_message(tool_call)
+                        await self._add_to_current_assistant_message(tool_call)
                         # Always yield function_call chunk if it's in includes
                         if chunk.type in includes:
                             yield chunk
@@ -390,9 +376,8 @@ class Runner:
                         target_message = None
 
                         # First check if we have a current assistant message
-                        if self._current_assistant_message is not None:
-                            target_message = self._current_assistant_message
-                        else:
+                        target_message = await self._message_state_manager.get_current_message()
+                        if target_message is None:
                             # Otherwise, look for the last assistant message in the list
                             for i in range(len(self.messages) - 1, -1, -1):
                                 current_message = self.messages[i]
@@ -422,9 +407,9 @@ class Runner:
                             yield chunk
                     case "timing":
                         # Update timing information in current assistant message
-                        if self._current_assistant_message is not None:
-                            self._current_assistant_message.meta.latency_ms = chunk.timing.latency_ms
-                            self._current_assistant_message.meta.total_time_ms = chunk.timing.output_time_ms
+                        current_message = await self._message_state_manager.get_current_message()
+                        if current_message is not None:
+                            await self._message_state_manager.update_meta(latency_ms=chunk.timing.latency_ms, total_time_ms=chunk.timing.output_time_ms)
                         # Also try to update the last assistant message if no current message
                         elif self.messages and isinstance(self.messages[-1], NewAssistantMessage):
                             last_message = cast("NewAssistantMessage", self.messages[-1])
@@ -437,7 +422,7 @@ class Runner:
                         yield chunk
 
             # Finalize assistant message so it can be found in pending function calls
-            self._finalize_assistant_message()
+            await self._finalize_assistant_message()
 
             # Check for pending tool calls after processing current assistant message
             pending_tool_calls = self._find_pending_tool_calls()
@@ -463,36 +448,6 @@ class Runner:
         require_confirm_tools = await self.agent.list_require_confirm_tools(tool_calls)
         return bool(require_confirm_tools)
 
-    async def run_continue_until_complete(
-        self,
-        max_steps: int = 20,
-        includes: list[AgentChunkType] | None = None,
-        record_to: PathLike | str | None = None,
-    ) -> list[AgentChunk]:
-        """Deprecated: Use run_until_complete(None) instead."""
-        warnings.warn(
-            "run_continue_until_complete is deprecated. Use run_until_complete(None) instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        resp = self.run_continue_stream(max_steps, includes, record_to=record_to)
-        return await self._collect_all_chunks(resp)
-
-    def run_continue_stream(
-        self,
-        max_steps: int = 20,
-        includes: list[AgentChunkType] | None = None,
-        record_to: PathLike | str | None = None,
-        context: "Any | None" = None,  # noqa: ANN401
-    ) -> AsyncGenerator[AgentChunk, None]:
-        """Deprecated: Use run(None) instead."""
-        warnings.warn(
-            "run_continue_stream is deprecated. Use run(None) instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        return self._run_continue_stream(max_steps, includes, record_to=record_to, context=context)
-
     async def _run_continue_stream(
         self,
         max_steps: int = 20,

lite_agent/types/__init__.py

@@ -17,7 +17,7 @@ from .events import (
     Usage as EventUsage,
 )
 from .messages import (
-    # Legacy compatibility aliases (temporary)
+    # Message wrapper classes
     AgentAssistantMessage,
     AgentSystemMessage,
     AgentUserMessage,

lite_agent/types/messages.py

@@ -211,7 +211,7 @@ class UserMessageContentItemImageURL(BaseModel):
     image_url: UserMessageContentItemImageURLImageURL
 
 
-# Legacy compatibility wrapper classes
+# Message wrapper classes for backward compatibility
 class AgentUserMessage(NewUserMessage):
     def __init__(
         self,
@@ -250,6 +250,7 @@ class AgentAssistantMessage(NewAssistantMessage):
         )
 
 
+# AgentSystemMessage is now an alias to NewSystemMessage
 AgentSystemMessage = NewSystemMessage
 RunnerMessage = NewMessage
 

lite_agent/utils/advanced_message_builder.py

@@ -0,0 +1,201 @@
+"""Advanced message builder with fluent interface for complex message construction."""
+
+from datetime import datetime, timezone
+from typing import Any
+
+from lite_agent.loggers import logger
+from lite_agent.types import (
+    AssistantMessageMeta,
+    AssistantTextContent,
+    AssistantToolCall,
+    AssistantToolCallResult,
+    MessageMeta,
+    MessageUsage,
+    NewAssistantMessage,
+    NewSystemMessage,
+    NewUserMessage,
+    UserImageContent,
+    UserTextContent,
+)
+
+
+class FluentMessageBuilder:
+    """Fluent interface for building complex messages step by step."""
+
+    def __init__(self):
+        self._reset()
+
+    def _reset(self) -> None:
+        """Reset builder state."""
+        self._message_type = None
+        self._content_items = []
+        self._meta = None
+
+    def user_message(self) -> "FluentMessageBuilder":
+        """Start building a user message."""
+        self._reset()
+        self._message_type = "user"
+        self._meta = MessageMeta()
+        logger.debug("Started building user message")
+        return self
+
+    def assistant_message(self, model: str | None = None) -> "FluentMessageBuilder":
+        """Start building an assistant message."""
+        self._reset()
+        self._message_type = "assistant"
+        self._meta = AssistantMessageMeta(model=model)
+        logger.debug(f"Started building assistant message (model: {model})")
+        return self
+
+    def system_message(self) -> "FluentMessageBuilder":
+        """Start building a system message."""
+        self._reset()
+        self._message_type = "system"
+        self._meta = MessageMeta()
+        self._content = ""
+        logger.debug("Started building system message")
+        return self
+
+    def add_text(self, text: str) -> "FluentMessageBuilder":
+        """Add text content."""
+        if self._message_type == "user":
+            self._content_items.append(UserTextContent(text=text))
+        elif self._message_type == "assistant":
+            self._content_items.append(AssistantTextContent(text=text))
+        elif self._message_type == "system":
+            self._content = text
+        else:
+            msg = "Message type not set. Call user_message(), assistant_message(), or system_message() first."
+            raise ValueError(msg)
+
+        logger.debug(f"Added text content (length: {len(text)})")
+        return self
+
+    def add_image(self, image_url: str | None = None, file_id: str | None = None, detail: str = "auto") -> "FluentMessageBuilder":
+        """Add image content to user message."""
+        if self._message_type != "user":
+            msg = "Images can only be added to user messages"
+            raise ValueError(msg)
+
+        self._content_items.append(UserImageContent(image_url=image_url, file_id=file_id, detail=detail))
+        logger.debug(f"Added image content (url: {bool(image_url)}, file_id: {bool(file_id)})")
+        return self
+
+    def add_tool_call(self, call_id: str, name: str, arguments: dict[str, Any] | str) -> "FluentMessageBuilder":
+        """Add tool call to assistant message."""
+        if self._message_type != "assistant":
+            msg = "Tool calls can only be added to assistant messages"
+            raise ValueError(msg)
+
+        self._content_items.append(AssistantToolCall(call_id=call_id, name=name, arguments=arguments))
+        logger.debug(f"Added tool call: {name} (call_id: {call_id})")
+        return self
+
+    def add_tool_result(self, call_id: str, output: str, execution_time_ms: int | None = None) -> "FluentMessageBuilder":
+        """Add tool call result to assistant message."""
+        if self._message_type != "assistant":
+            msg = "Tool results can only be added to assistant messages"
+            raise ValueError(msg)
+
+        self._content_items.append(AssistantToolCallResult(call_id=call_id, output=output, execution_time_ms=execution_time_ms))
+        logger.debug(f"Added tool result for call: {call_id}")
+        return self
+
+    def with_timestamp(self, timestamp: datetime | None = None) -> "FluentMessageBuilder":
+        """Set message timestamp."""
+        if self._meta is None:
+            msg = "Message type not set. Call user_message(), assistant_message(), or system_message() first."
+            raise ValueError(msg)
+        if timestamp is None:
+            timestamp = datetime.now(timezone.utc)
+        self._meta.sent_at = timestamp
+        return self
+
+    def with_usage(self, input_tokens: int | None = None, output_tokens: int | None = None) -> "FluentMessageBuilder":
+        """Set usage information (assistant messages only)."""
+        if self._message_type != "assistant":
+            msg = "Usage information can only be set for assistant messages"
+            raise ValueError(msg)
+
+        if self._meta is None:
+            msg = "Message type not set. Call user_message(), assistant_message(), or system_message() first."
+            raise ValueError(msg)
+
+        if not hasattr(self._meta, "usage"):
+            self._meta.usage = MessageUsage()
+
+        if self._meta.usage is not None:
+            if input_tokens is not None:
+                self._meta.usage.input_tokens = input_tokens
+            if output_tokens is not None:
+                self._meta.usage.output_tokens = output_tokens
+            if input_tokens is not None and output_tokens is not None:
+                self._meta.usage.total_tokens = input_tokens + output_tokens
+
+        return self
+
+    def with_timing(self, latency_ms: int | None = None, total_time_ms: int | None = None) -> "FluentMessageBuilder":
+        """Set timing information (assistant messages only)."""
+        if self._message_type != "assistant":
+            msg = "Timing information can only be set for assistant messages"
+            raise ValueError(msg)
+
+        if self._meta is None:
+            msg = "Message type not set. Call user_message(), assistant_message(), or system_message() first."
+            raise ValueError(msg)
+
+        if latency_ms is not None:
+            self._meta.latency_ms = latency_ms  # type: ignore[attr-defined]
+        if total_time_ms is not None:
+            self._meta.total_time_ms = total_time_ms  # type: ignore[attr-defined]
+
+        return self
+
+    def build(self) -> NewUserMessage | NewAssistantMessage | NewSystemMessage:
+        """Build the final message."""
+        if self._message_type == "user":
+            message = NewUserMessage(content=self._content_items, meta=self._meta)
+        elif self._message_type == "assistant":
+            message = NewAssistantMessage(content=self._content_items, meta=self._meta)
+        elif self._message_type == "system":
+            message = NewSystemMessage(content=self._content, meta=self._meta)
+        else:
+            msg = "Message type not set"
+            raise ValueError(msg)
+
+        logger.debug(f"Built {self._message_type} message with {len(getattr(self, '_content_items', []))} content items")
+        return message
+
+
+class MessageBuilderFactory:
+    """Factory for creating common message types quickly."""
+
+    @staticmethod
+    def create_simple_user_message(text: str) -> NewUserMessage:
+        """Create a simple user text message."""
+        return FluentMessageBuilder().user_message().add_text(text).build()
+
+    @staticmethod
+    def create_simple_assistant_message(text: str, model: str | None = None) -> NewAssistantMessage:
+        """Create a simple assistant text message."""
+        return FluentMessageBuilder().assistant_message(model).add_text(text).build()
+
+    @staticmethod
+    def create_system_message(text: str) -> NewSystemMessage:
+        """Create a system message."""
+        return FluentMessageBuilder().system_message().add_text(text).build()
+
+    @staticmethod
+    def create_user_message_with_image(text: str, image_url: str) -> NewUserMessage:
+        """Create a user message with text and image."""
+        return FluentMessageBuilder().user_message().add_text(text).add_image(image_url=image_url).build()
+
+    @staticmethod
+    def create_assistant_with_tool_call(text: str, call_id: str, tool_name: str, arguments: dict[str, Any], model: str | None = None) -> NewAssistantMessage:
+        """Create an assistant message with text and a tool call."""
+        return FluentMessageBuilder().assistant_message(model).add_text(text).add_tool_call(call_id, tool_name, arguments).build()
+
+    @staticmethod
+    def create_assistant_with_tool_result(call_id: str, result: str, execution_time_ms: int | None = None, model: str | None = None) -> NewAssistantMessage:
+        """Create an assistant message with just a tool result."""
+        return FluentMessageBuilder().assistant_message(model).add_tool_result(call_id, result, execution_time_ms).build()

lite_agent/utils/message_builder.py

@@ -190,7 +190,7 @@ class MessageBuilder:
                 # Fallback for other content types
                 assistant_content_items = [AssistantTextContent(text=str(content))]
 
-        # Handle tool calls if present (legacy format)
+        # Handle tool calls if present
         if "tool_calls" in message:
             for tool_call in message.get("tool_calls", []):
                 try: