lite-agent 0.6.0__py3-none-any.whl → 0.8.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,42 @@ class Agent:
             role = message_dict.get("role")
 
             if role == "assistant":
-                # Look ahead for function_call messages
+                # Extract tool_calls from content if present
                 tool_calls = []
+                content = message_dict.get("content", [])
+
+                # Handle both string and array content
+                if isinstance(content, list):
+                    # Extract tool_calls 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") == "text":
+                                filtered_content.append(item)
+                            # Skip tool_call_result - they should be handled by separate function_call_output messages
+
+                    # 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
@@ -424,6 +489,13 @@ class Agent:
                 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
+
                 converted_messages.append(assistant_msg)
                 i = j  # Skip the function_call messages we've processed
 
@@ -536,10 +608,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()

lite_agent/chat_display.py

@@ -26,6 +26,8 @@ from lite_agent.types import (
     AgentSystemMessage,
     AgentUserMessage,
     AssistantMessageMeta,
+    AssistantToolCall,
+    AssistantToolCallResult,
     BasicMessageMeta,
     FlexibleRunnerMessage,
     LLMResponseMeta,
@@ -228,9 +230,9 @@ def _update_message_counts(message: FlexibleRunnerMessage, counts: dict[str, int
         counts["Assistant"] += 1
         # Count tool calls and outputs within the assistant message
         for content_item in message.content:
-            if content_item.type == "tool_call":
+            if isinstance(content_item, AssistantToolCall):
                 counts["Function Call"] += 1
-            elif content_item.type == "tool_call_result":
+            elif isinstance(content_item, AssistantToolCallResult):
                 counts["Function Output"] += 1
     elif isinstance(message, NewSystemMessage):
         counts["System"] += 1
@@ -295,10 +297,18 @@ def _process_object_meta(meta: BasicMessageMeta | LLMResponseMeta | AssistantMes
     """处理对象类型的 meta 数据。"""
     # LLMResponseMeta 和 AssistantMessageMeta 都有这些字段
     if isinstance(meta, (LLMResponseMeta, AssistantMessageMeta)):
-        if hasattr(meta, "input_tokens") and meta.input_tokens is not None:
-            total_input += int(meta.input_tokens)
-        if hasattr(meta, "output_tokens") and meta.output_tokens is not None:
-            total_output += int(meta.output_tokens)
+        # For AssistantMessageMeta, use the structured usage field
+        if isinstance(meta, AssistantMessageMeta) and meta.usage is not None:
+            if meta.usage.input_tokens is not None:
+                total_input += int(meta.usage.input_tokens)
+            if meta.usage.output_tokens is not None:
+                total_output += int(meta.usage.output_tokens)
+        # For LLMResponseMeta, use the flat fields
+        elif isinstance(meta, LLMResponseMeta):
+            if hasattr(meta, "input_tokens") and meta.input_tokens is not None:
+                total_input += int(meta.input_tokens)
+            if hasattr(meta, "output_tokens") and meta.output_tokens is not None:
+                total_output += int(meta.output_tokens)
         if hasattr(meta, "latency_ms") and meta.latency_ms is not None:
             total_latency += int(meta.latency_ms)
         if hasattr(meta, "output_time_ms") and meta.output_time_ms is not None:
@@ -363,11 +373,9 @@ def display_chat_summary(messages: RunnerMessages, *, console: Console | None =
         messages: 要汇总的消息列表
         console: Rich Console 实例，如果为 None 则创建新的
     """
-    if console is None:
-        console = Console()
-
+    active_console = console or Console()
     summary_table = build_chat_summary_table(messages)
-    console.print(summary_table)
+    active_console.print(summary_table)
 
 
 def display_messages(
@@ -577,9 +585,9 @@ def _display_assistant_message_compact_v2(message: AgentAssistantMessage, contex
             meta_parts.append(f"Latency:{message.meta.latency_ms}ms")
         if message.meta.output_time_ms is not None:
             meta_parts.append(f"Output:{message.meta.output_time_ms}ms")
-        if message.meta.input_tokens is not None and message.meta.output_tokens is not None:
-            total_tokens = message.meta.input_tokens + message.meta.output_tokens
-            meta_parts.append(f"Tokens:↑{message.meta.input_tokens}↓{message.meta.output_tokens}={total_tokens}")
+        if message.meta.usage and message.meta.usage.input_tokens is not None and message.meta.usage.output_tokens is not None:
+            total_tokens = message.meta.usage.input_tokens + message.meta.usage.output_tokens
+            meta_parts.append(f"Tokens:↑{message.meta.usage.input_tokens}↓{message.meta.usage.output_tokens}={total_tokens}")
 
         if meta_parts:
             meta_info = f" [dim]({' | '.join(meta_parts)})[/dim]"

lite_agent/client.py

@@ -100,6 +100,7 @@ class BaseLLMClient(abc.ABC):
         tools: list[ChatCompletionToolParam] | None = None,
         tool_choice: str = "auto",
         reasoning: ReasoningConfig = None,
+        *,
         streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # noqa: ANN401
@@ -112,6 +113,7 @@ class BaseLLMClient(abc.ABC):
         tools: list[FunctionToolParam] | None = None,
         tool_choice: Literal["none", "auto", "required"] = "auto",
         reasoning: ReasoningConfig = None,
+        *,
         streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # noqa: ANN401
@@ -136,6 +138,7 @@ class LiteLLMClient(BaseLLMClient):
         tools: list[ChatCompletionToolParam] | None = None,
         tool_choice: str = "auto",
         reasoning: ReasoningConfig = None,
+        *,
         streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # noqa: ANN401
@@ -187,6 +190,7 @@ class LiteLLMClient(BaseLLMClient):
         tools: list[FunctionToolParam] | None = None,
         tool_choice: Literal["none", "auto", "required"] = "auto",
         reasoning: ReasoningConfig = None,
+        *,
         streaming: bool = True,
         **kwargs: Any,  # noqa: ANN401
     ) -> Any:  # type: ignore[return]  # noqa: ANN401

lite_agent/constants.py

@@ -0,0 +1,30 @@
+from typing import Literal
+
+
+class CompletionMode:
+    """Agent completion modes."""
+
+    STOP: Literal["stop"] = "stop"  # Traditional completion until model decides to stop
+    CALL: Literal["call"] = "call"  # Completion until specific tool is called
+
+
+class ToolName:
+    """System tool names."""
+
+    TRANSFER_TO_AGENT = "transfer_to_agent"
+    TRANSFER_TO_PARENT = "transfer_to_parent"
+    WAIT_FOR_USER = "wait_for_user"
+
+
+class StreamIncludes:
+    """Default stream includes configuration."""
+
+    DEFAULT_INCLUDES = (
+        "completion_raw",
+        "usage",
+        "function_call",
+        "function_call_output",
+        "content_delta",
+        "function_call_delta",
+        "assistant_message",
+    )

lite_agent/message_transfers.py

@@ -5,7 +5,7 @@ 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
+from lite_agent.types import NewUserMessage, RunnerMessages, UserTextContent
 
 
 def consolidate_history_transfer(messages: RunnerMessages) -> RunnerMessages:
@@ -43,8 +43,8 @@ def consolidate_history_transfer(messages: RunnerMessages) -> RunnerMessages:
     # 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}]
+    # Return a single user message using NewMessage format
+    return [NewUserMessage(content=[UserTextContent(text=consolidated_content)])]
 
 
 def _process_message_to_xml(message: dict | object) -> list[str]:

lite_agent/processors/completion_event_processor.py

@@ -26,6 +26,7 @@ from lite_agent.types import (
     ToolCallFunction,
     UsageEvent,
 )
+from lite_agent.utils.metrics import TimingMetrics
 
 
 class CompletionEventProcessor:
@@ -71,21 +72,18 @@ class CompletionEventProcessor:
             if not self.yielded_content:
                 self.yielded_content = True
                 end_time = datetime.now(timezone.utc)
-                latency_ms = None
-                output_time_ms = None
-                # latency_ms: 从开始准备输出到 LLM 输出第一个字符的时间差
-                if self._start_time and self._first_output_time:
-                    latency_ms = int((self._first_output_time - self._start_time).total_seconds() * 1000)
-                # output_time_ms: 从输出第一个字符到输出完成的时间差
-                if self._first_output_time and self._output_complete_time:
-                    output_time_ms = int((self._output_complete_time - self._first_output_time).total_seconds() * 1000)
+                latency_ms = TimingMetrics.calculate_latency_ms(self._start_time, self._first_output_time)
+                output_time_ms = TimingMetrics.calculate_output_time_ms(self._first_output_time, self._output_complete_time)
 
                 usage = MessageUsage(
                     input_tokens=self._usage_data.get("input_tokens"),
                     output_tokens=self._usage_data.get("output_tokens"),
                 )
+                # Extract model information from chunk
+                model_name = getattr(chunk, "model", None)
                 meta = AssistantMessageMeta(
                     sent_at=end_time,
+                    model=model_name,
                     latency_ms=latency_ms,
                     total_time_ms=output_time_ms,
                     usage=usage,
@@ -152,21 +150,18 @@ class CompletionEventProcessor:
             if not self.yielded_content:
                 self.yielded_content = True
                 end_time = datetime.now(timezone.utc)
-                latency_ms = None
-                output_time_ms = None
-                # latency_ms: 从开始准备输出到 LLM 输出第一个字符的时间差
-                if self._start_time and self._first_output_time:
-                    latency_ms = int((self._first_output_time - self._start_time).total_seconds() * 1000)
-                # output_time_ms: 从输出第一个字符到输出完成的时间差
-                if self._first_output_time and self._output_complete_time:
-                    output_time_ms = int((self._output_complete_time - self._first_output_time).total_seconds() * 1000)
+                latency_ms = TimingMetrics.calculate_latency_ms(self._start_time, self._first_output_time)
+                output_time_ms = TimingMetrics.calculate_output_time_ms(self._first_output_time, self._output_complete_time)
 
                 usage = MessageUsage(
                     input_tokens=self._usage_data.get("input_tokens"),
                     output_tokens=self._usage_data.get("output_tokens"),
                 )
+                # Extract model information from chunk
+                model_name = getattr(chunk, "model", None)
                 meta = AssistantMessageMeta(
                     sent_at=end_time,
+                    model=model_name,
                     latency_ms=latency_ms,
                     total_time_ms=output_time_ms,
                     usage=usage,
@@ -199,10 +194,9 @@ class CompletionEventProcessor:
             results.append(UsageEvent(usage=EventUsage(input_tokens=usage["prompt_tokens"], output_tokens=usage["completion_tokens"])))
 
             # Then yield timing event if we have timing data
-            if self._start_time and self._first_output_time and self._output_complete_time:
-                latency_ms = int((self._first_output_time - self._start_time).total_seconds() * 1000)
-                output_time_ms = int((self._output_complete_time - self._first_output_time).total_seconds() * 1000)
-
+            latency_ms = TimingMetrics.calculate_latency_ms(self._start_time, self._first_output_time)
+            output_time_ms = TimingMetrics.calculate_output_time_ms(self._first_output_time, self._output_complete_time)
+            if latency_ms is not None and output_time_ms is not None:
                 results.append(
                     TimingEvent(
                         timing=Timing(

lite_agent/processors/response_event_processor.py

@@ -22,12 +22,14 @@ from lite_agent.types import (
     ContentDeltaEvent,
     EventUsage,
     FunctionCallEvent,
+    MessageUsage,
     NewAssistantMessage,
     ResponseRawEvent,
     Timing,
     TimingEvent,
     UsageEvent,
 )
+from lite_agent.utils.metrics import TimingMetrics
 
 
 class ResponseEventProcessor:
@@ -111,21 +113,26 @@ class ResponseEventProcessor:
                 content = item.get("content", [])
                 if content and isinstance(content, list) and len(content) > 0:
                     end_time = datetime.now(timezone.utc)
-                    latency_ms = None
-                    output_time_ms = None
-                    # latency_ms: 从开始准备输出到 LLM 输出第一个字符的时间差
-                    if self._start_time and self._first_output_time:
-                        latency_ms = int((self._first_output_time - self._start_time).total_seconds() * 1000)
-                    # output_time_ms: 从输出第一个字符到输出完成的时间差
-                    if self._first_output_time and self._output_complete_time:
-                        output_time_ms = int((self._output_complete_time - self._first_output_time).total_seconds() * 1000)
-
+                    latency_ms = TimingMetrics.calculate_latency_ms(self._start_time, self._first_output_time)
+                    output_time_ms = TimingMetrics.calculate_output_time_ms(self._first_output_time, self._output_complete_time)
+
+                    # Extract model information from event
+                    model_name = getattr(event, "model", None)
+                    # Debug: check if event has model info in different location
+                    if hasattr(event, "response") and hasattr(event.response, "model"):
+                        model_name = getattr(event.response, "model", None)
+                    # Create usage information
+                    usage = MessageUsage(
+                        input_tokens=self._usage_data.get("input_tokens"),
+                        output_tokens=self._usage_data.get("output_tokens"),
+                        total_tokens=(self._usage_data.get("input_tokens") or 0) + (self._usage_data.get("output_tokens") or 0),
+                    )
                     meta = AssistantMessageMeta(
                         sent_at=end_time,
+                        model=model_name,
                         latency_ms=latency_ms,
                         output_time_ms=output_time_ms,
-                        input_tokens=self._usage_data.get("input_tokens"),
-                        output_tokens=self._usage_data.get("output_tokens"),
+                        usage=usage,
                     )
                     return [
                         AssistantMessageEvent(
@@ -173,10 +180,9 @@ class ResponseEventProcessor:
                 )
 
                 # Then yield timing event if we have timing data
-                if self._start_time and self._first_output_time and self._output_complete_time:
-                    latency_ms = int((self._first_output_time - self._start_time).total_seconds() * 1000)
-                    output_time_ms = int((self._output_complete_time - self._first_output_time).total_seconds() * 1000)
-
+                latency_ms = TimingMetrics.calculate_latency_ms(self._start_time, self._first_output_time)
+                output_time_ms = TimingMetrics.calculate_output_time_ms(self._first_output_time, self._output_complete_time)
+                if latency_ms is not None and output_time_ms is not None:
                     results.append(
                         TimingEvent(
                             timing=Timing(

lite_agent/response_handlers/__init__.py

@@ -1,4 +1,5 @@
 """Response handlers for unified streaming and non-streaming processing."""
+
 from lite_agent.response_handlers.base import ResponseHandler
 from lite_agent.response_handlers.completion import CompletionResponseHandler
 from lite_agent.response_handlers.responses import ResponsesAPIHandler