Coverage for src/lite_agent/response_handlers/completion.py: 55%

1"""Completion API response handler.""" 

2 

3from collections.abc import AsyncGenerator 

4from datetime import datetime, timezone 

5from pathlib import Path 

6from typing import Any 

7 

8from litellm import CustomStreamWrapper 

9 

10from lite_agent.response_handlers.base import ResponseHandler 

11from lite_agent.stream_handlers import litellm_completion_stream_handler 

12from lite_agent.types import AgentChunk 

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

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

15 

16 

17class CompletionResponseHandler(ResponseHandler): 

18 """Handler for Completion API responses.""" 

19 

20 async def _handle_streaming( 

21 self, 

22 response: Any, # noqa: ANN401 

23 record_to: Path | None = None, 

24 ) -> AsyncGenerator[AgentChunk, None]: 

25 """Handle streaming completion response.""" 

26 if isinstance(response, CustomStreamWrapper): 

27 async for chunk in litellm_completion_stream_handler(response, record_to): 

28 yield chunk 

29 else: 

30 msg = "Response is not a CustomStreamWrapper, cannot stream chunks." 

31 raise TypeError(msg) 

32 

33 async def _handle_non_streaming( 

34 self, 

35 response: Any, # noqa: ANN401 

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

37 ) -> AsyncGenerator[AgentChunk, None]: 

38 """Handle non-streaming completion response.""" 

39 # Convert completion response to chunks 

40 if hasattr(response, "choices") and response.choices: 

41 choice = response.choices[0] 

42 content_items = [] 

43 

44 # Add text content 

45 if choice.message and choice.message.content: 

46 content_items.append(AssistantTextContent(text=choice.message.content)) 

47 

48 # Handle tool calls 

49 if choice.message and choice.message.tool_calls: 

50 for tool_call in choice.message.tool_calls: 

51 content_items.append( # noqa: PERF401 

52 AssistantToolCall( 

53 call_id=tool_call.id, 

54 name=tool_call.function.name, 

55 arguments=tool_call.function.arguments, 

56 ), 

57 ) 

58 

59 # Always yield assistant message, even if content is empty for tool calls 

60 if choice.message and (content_items or choice.message.tool_calls): 

61 # Extract model information from response 

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

63 message = NewAssistantMessage( 

64 content=content_items, 

65 meta=AssistantMessageMeta( 

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

67 model=model_name, 

68 ), 

69 ) 

70 yield AssistantMessageEvent(message=message) 

71 

72 # Yield usage information if available 

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

74 usage = Usage( 

75 input_tokens=response.usage.prompt_tokens, 

76 output_tokens=response.usage.completion_tokens, 

77 ) 

78 yield UsageEvent(usage=usage)