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

lite_agent/response_handlers/base.py

@@ -1,4 +1,5 @@
 """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
@@ -12,35 +13,42 @@ class ResponseHandler(ABC):
 
     async def handle(
         self,
-        response: Any,
+        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:
-            async for chunk in self._handle_streaming(response, record_to):
+            stream = self._handle_streaming(response, record_to)
+            async for chunk in stream:
                 yield chunk
         else:
-            async for chunk in self._handle_non_streaming(response, record_to):
+            stream = self._handle_non_streaming(response, record_to)
+            async for chunk in stream:
                 yield chunk
 
     @abstractmethod
-    async def _handle_streaming(
-        self, response: Any, record_to: Path | None = None,
+    def _handle_streaming(
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,
     ) -> AsyncGenerator[AgentChunk, None]:
         """Handle streaming response."""
 
     @abstractmethod
-    async def _handle_non_streaming(
-        self, response: Any, record_to: Path | None = None,
+    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

@@ -1,4 +1,5 @@
 """Completion API response handler."""
+
 from collections.abc import AsyncGenerator
 from datetime import datetime, timezone
 from pathlib import Path
@@ -9,15 +10,17 @@ 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
-from lite_agent.types.messages import AssistantMessageMeta, AssistantTextContent, NewAssistantMessage
+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, record_to: Path | None = None,
+        self,
+        response: Any,  # noqa: ANN401
+        record_to: Path | None = None,
     ) -> AsyncGenerator[AgentChunk, None]:
         """Handle streaming completion response."""
         if isinstance(response, CustomStreamWrapper):
@@ -28,7 +31,9 @@ class CompletionResponseHandler(ResponseHandler):
             raise TypeError(msg)
 
     async def _handle_non_streaming(
-        self, response: Any, record_to: Path | None = None,
+        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
@@ -40,11 +45,34 @@ class CompletionResponseHandler(ResponseHandler):
             if choice.message and choice.message.content:
                 content_items.append(AssistantTextContent(text=choice.message.content))
 
-            # TODO: Handle tool calls in the future
+            # 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,
+                        ),
+                    )
 
-            if content_items:
+            # 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)),
+                    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

@@ -1,4 +1,5 @@
 """Responses API response handler."""
+
 from collections.abc import AsyncGenerator
 from datetime import datetime, timezone
 from pathlib import Path
@@ -7,36 +8,69 @@ 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
-from lite_agent.types.messages import AssistantMessageMeta, AssistantTextContent, NewAssistantMessage
+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, record_to: Path | None = None,
+        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, record_to: Path | None = None,
+        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:
-            for output_message in response.output:
-                if hasattr(output_message, "content") and output_message.content:
+            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_message.content:
+                    for content_item in output_item.content:
                         if hasattr(content_item, "text"):
                             content_text += content_item.text
 
                     if content_text:
-                        message = NewAssistantMessage(
-                            content=[AssistantTextContent(text=content_text)],
-                            meta=AssistantMessageMeta(sent_at=datetime.now(timezone.utc)),
-                        )
-                        yield AssistantMessageEvent(message=message)
+                        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)