Coverage for src/lite_agent/response_handlers/responses.py: 36%

1"""Responses API response handler.""" 

2 

3from collections.abc import AsyncGenerator 

4from datetime import datetime, timezone 

5from pathlib import Path 

6from typing import Any 

7 

8from lite_agent.response_handlers.base import ResponseHandler 

9from lite_agent.stream_handlers import litellm_response_stream_handler 

10from lite_agent.types import AgentChunk 

11from lite_agent.types.events import AssistantMessageEvent, Usage, UsageEvent 

12from lite_agent.types.messages import AssistantMessageMeta, AssistantTextContent, AssistantToolCall, NewAssistantMessage 

13 

14 

15class ResponsesAPIHandler(ResponseHandler): 

16 """Handler for Responses API responses.""" 

17 

18 async def _handle_streaming( 

19 self, 

20 response: Any, # noqa: ANN401 

21 record_to: Path | None = None, 

22 ) -> AsyncGenerator[AgentChunk, None]: 

23 """Handle streaming responses API response.""" 

24 async for chunk in litellm_response_stream_handler(response, record_to): 

25 yield chunk 

26 

27 async def _handle_non_streaming( 

28 self, 

29 response: Any, # noqa: ANN401 

30 record_to: Path | None = None, # noqa: ARG002 

31 ) -> AsyncGenerator[AgentChunk, None]: 

32 """Handle non-streaming responses API response.""" 

33 # Convert ResponsesAPIResponse to chunks 

34 if hasattr(response, "output") and response.output: 

35 content_items = [] 

36 

37 for output_item in response.output: 

38 # Handle function tool calls 

39 if hasattr(output_item, "type") and output_item.type == "function_call": 

40 content_items.append( 

41 AssistantToolCall( 

42 call_id=output_item.call_id, 

43 name=output_item.name, 

44 arguments=output_item.arguments, 

45 ), 

46 ) 

47 # Handle text content (if exists) 

48 elif hasattr(output_item, "content") and output_item.content: 

49 content_text = "" 

50 for content_item in output_item.content: 

51 if hasattr(content_item, "text"): 

52 content_text += content_item.text 

53 

54 if content_text: 

55 content_items.append(AssistantTextContent(text=content_text)) 

56 

57 # Create assistant message if we have any content 

58 if content_items: 

59 # Extract model information from response 

60 model_name = getattr(response, "model", None) 

61 message = NewAssistantMessage( 

62 content=content_items, 

63 meta=AssistantMessageMeta( 

64 sent_at=datetime.now(timezone.utc), 

65 model=model_name, 

66 ), 

67 ) 

68 yield AssistantMessageEvent(message=message) 

69 

70 # Yield usage information if available 

71 if hasattr(response, "usage") and response.usage: 

72 usage = Usage( 

73 input_tokens=response.usage.input_tokens, 

74 output_tokens=response.usage.output_tokens, 

75 ) 

76 yield UsageEvent(usage=usage)