lite-agent 0.5.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

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

@@ -0,0 +1,11 @@
+"""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
+
+__all__ = [
+    "CompletionResponseHandler",
+    "ResponseHandler",
+    "ResponsesAPIHandler",
+]

lite_agent/response_handlers/base.py

@@ -0,0 +1,54 @@
+"""Base response handler for unified streaming and non-streaming response processing."""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from pathlib import Path
+from typing import Any
+
+from lite_agent.types import AgentChunk
+
+
+class ResponseHandler(ABC):
+    """Base class for handling both streaming and non-streaming responses."""
+
+    async def handle(
+        self,
+        response: Any,  # noqa: ANN401
+        *,
+        streaming: bool,
+        record_to: Path | None = None,
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Handle a response in either streaming or non-streaming mode.
+
+        Args:
+            response: The LLM response object
+            streaming: Whether to process as streaming or non-streaming
+            record_to: Optional file path to record the conversation
+
+        Yields:
+            AgentChunk: Processed chunks from the response
+        """
+        if streaming:
+            stream = self._handle_streaming(response, record_to)
+            async for chunk in stream:
+                yield chunk
+        else:
+            stream = self._handle_non_streaming(response, record_to)
+            async for chunk in stream:
+                yield chunk
+
+    @abstractmethod
+    def _handle_streaming(
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Handle streaming response."""
+
+    @abstractmethod
+    def _handle_non_streaming(
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Handle non-streaming response."""

lite_agent/response_handlers/completion.py

@@ -0,0 +1,78 @@
+"""Completion API response handler."""
+
+from collections.abc import AsyncGenerator
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from litellm import CustomStreamWrapper
+
+from lite_agent.response_handlers.base import ResponseHandler
+from lite_agent.stream_handlers import litellm_completion_stream_handler
+from lite_agent.types import AgentChunk
+from lite_agent.types.events import AssistantMessageEvent, Usage, UsageEvent
+from lite_agent.types.messages import AssistantMessageMeta, AssistantTextContent, AssistantToolCall, NewAssistantMessage
+
+
+class CompletionResponseHandler(ResponseHandler):
+    """Handler for Completion API responses."""
+
+    async def _handle_streaming(
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Handle streaming completion response."""
+        if isinstance(response, CustomStreamWrapper):
+            async for chunk in litellm_completion_stream_handler(response, record_to):
+                yield chunk
+        else:
+            msg = "Response is not a CustomStreamWrapper, cannot stream chunks."
+            raise TypeError(msg)
+
+    async def _handle_non_streaming(
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,  # noqa: ARG002
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Handle non-streaming completion response."""
+        # Convert completion response to chunks
+        if hasattr(response, "choices") and response.choices:
+            choice = response.choices[0]
+            content_items = []
+
+            # Add text content
+            if choice.message and choice.message.content:
+                content_items.append(AssistantTextContent(text=choice.message.content))
+
+            # Handle tool calls
+            if choice.message and choice.message.tool_calls:
+                for tool_call in choice.message.tool_calls:
+                    content_items.append(  # noqa: PERF401
+                        AssistantToolCall(
+                            call_id=tool_call.id,
+                            name=tool_call.function.name,
+                            arguments=tool_call.function.arguments,
+                        ),
+                    )
+
+            # Always yield assistant message, even if content is empty for tool calls
+            if choice.message and (content_items or choice.message.tool_calls):
+                # Extract model information from response
+                model_name = getattr(response, "model", None)
+                message = NewAssistantMessage(
+                    content=content_items,
+                    meta=AssistantMessageMeta(
+                        sent_at=datetime.now(timezone.utc),
+                        model=model_name,
+                    ),
+                )
+                yield AssistantMessageEvent(message=message)
+
+        # Yield usage information if available
+        if hasattr(response, "usage") and response.usage:
+            usage = Usage(
+                input_tokens=response.usage.prompt_tokens,
+                output_tokens=response.usage.completion_tokens,
+            )
+            yield UsageEvent(usage=usage)

lite_agent/response_handlers/responses.py

@@ -0,0 +1,76 @@
+"""Responses API response handler."""
+
+from collections.abc import AsyncGenerator
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from lite_agent.response_handlers.base import ResponseHandler
+from lite_agent.stream_handlers import litellm_response_stream_handler
+from lite_agent.types import AgentChunk
+from lite_agent.types.events import AssistantMessageEvent, Usage, UsageEvent
+from lite_agent.types.messages import AssistantMessageMeta, AssistantTextContent, AssistantToolCall, NewAssistantMessage
+
+
+class ResponsesAPIHandler(ResponseHandler):
+    """Handler for Responses API responses."""
+
+    async def _handle_streaming(
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Handle streaming responses API response."""
+        async for chunk in litellm_response_stream_handler(response, record_to):
+            yield chunk
+
+    async def _handle_non_streaming(
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,  # noqa: ARG002
+    ) -> AsyncGenerator[AgentChunk, None]:
+        """Handle non-streaming responses API response."""
+        # Convert ResponsesAPIResponse to chunks
+        if hasattr(response, "output") and response.output:
+            content_items = []
+
+            for output_item in response.output:
+                # Handle function tool calls
+                if hasattr(output_item, "type") and output_item.type == "function_call":
+                    content_items.append(
+                        AssistantToolCall(
+                            call_id=output_item.call_id,
+                            name=output_item.name,
+                            arguments=output_item.arguments,
+                        ),
+                    )
+                # Handle text content (if exists)
+                elif hasattr(output_item, "content") and output_item.content:
+                    content_text = ""
+                    for content_item in output_item.content:
+                        if hasattr(content_item, "text"):
+                            content_text += content_item.text
+
+                    if content_text:
+                        content_items.append(AssistantTextContent(text=content_text))
+
+            # Create assistant message if we have any content
+            if content_items:
+                # Extract model information from response
+                model_name = getattr(response, "model", None)
+                message = NewAssistantMessage(
+                    content=content_items,
+                    meta=AssistantMessageMeta(
+                        sent_at=datetime.now(timezone.utc),
+                        model=model_name,
+                    ),
+                )
+                yield AssistantMessageEvent(message=message)
+
+        # Yield usage information if available
+        if hasattr(response, "usage") and response.usage:
+            usage = Usage(
+                input_tokens=response.usage.input_tokens,
+                output_tokens=response.usage.output_tokens,
+            )
+            yield UsageEvent(usage=usage)