Coverage for src/lite_agent/response_handlers/base.py: 83%

1"""Base response handler for unified streaming and non-streaming response processing.""" 

2 

3from abc import ABC, abstractmethod 

4from collections.abc import AsyncGenerator 

5from pathlib import Path 

6from typing import Any 

7 

8from lite_agent.types import AgentChunk 

9 

10 

11class ResponseHandler(ABC): 

12 """Base class for handling both streaming and non-streaming responses.""" 

13 

14 async def handle( 

15 self, 

16 response: Any, # noqa: ANN401 

17 *, 

18 streaming: bool, 

19 record_to: Path | None = None, 

20 ) -> AsyncGenerator[AgentChunk, None]: 

21 """Handle a response in either streaming or non-streaming mode. 

22 

23 Args: 

24 response: The LLM response object 

25 streaming: Whether to process as streaming or non-streaming 

26 record_to: Optional file path to record the conversation 

27 

28 Yields: 

29 AgentChunk: Processed chunks from the response 

30 """ 

31 if streaming: 

32 stream = self._handle_streaming(response, record_to) 

33 async for chunk in stream: 

34 yield chunk 

35 else: 

36 stream = self._handle_non_streaming(response, record_to) 

37 async for chunk in stream: 

38 yield chunk 

39 

40 @abstractmethod 

41 def _handle_streaming( 

42 self, 

43 response: Any, # noqa: ANN401 

44 record_to: Path | None = None, 

45 ) -> AsyncGenerator[AgentChunk, None]: 

46 """Handle streaming response.""" 

47 

48 @abstractmethod 

49 def _handle_non_streaming( 

50 self, 

51 response: Any, # noqa: ANN401 

52 record_to: Path | None = None, 

53 ) -> AsyncGenerator[AgentChunk, None]: 

54 """Handle non-streaming response."""