Coverage for projects/04-llm-adapter-shadow/src/llm_adapter/providers/mock.py: 94%

1"""Mock provider that can deterministically trigger failure modes.""" 

2 

3from __future__ import annotations 

4 

5import random 

6import time 

7from collections.abc import Iterable 

8 

9from ..errors import AdapterError, RateLimitError, RetriableError, TimeoutError 

10from ..provider_spi import ProviderRequest, ProviderResponse, ProviderSPI, TokenUsage 

11 

12ErrorSpec = tuple[type[AdapterError], str] 

13_ERROR_BY_MARKER: dict[str, ErrorSpec] = { 

14 "[TIMEOUT]": (TimeoutError, "simulated timeout"), 

15 "[RATELIMIT]": (RateLimitError, "simulated rate limit"), 

16 "[INVALID_JSON]": (RetriableError, "simulated invalid JSON"), 

17} 

18 

19 

20class MockProvider(ProviderSPI): 

21 """Very small provider implementation for exercising the adapter.""" 

22 

23 def __init__( 

24 self, 

25 name: str, 

26 base_latency_ms: int = 50, 

27 error_markers: Iterable[str] | None = None, 

28 ) -> None: 

29 self._name = name 

30 self.base_latency_ms = base_latency_ms 

31 if error_markers is None: 

32 self._error_markers: set[str] = set(_ERROR_BY_MARKER) 

33 else: 

34 self._error_markers = { 

35 marker for marker in error_markers if marker in _ERROR_BY_MARKER 

36 } 

37 

38 def name(self) -> str: 

39 return self._name 

40 

41 def capabilities(self) -> set[str]: 

42 return {"chat"} 

43 

44 def _maybe_raise_error(self, text: str) -> None: 

45 for marker in self._error_markers: 

46 if marker in text: 

47 exc_cls, message = _ERROR_BY_MARKER[marker] 

48 raise exc_cls(message) 

49 

50 def invoke(self, request: ProviderRequest) -> ProviderResponse: 

51 text = request.prompt 

52 self._maybe_raise_error(text) 

53 

54 latency = self.base_latency_ms + int(random.random() * 20) 

55 time.sleep(latency / 1000.0) 

56 

57 prompt_tokens = max(1, len(text) // 4) 

58 completion_tokens = 16 

59 

60 return ProviderResponse( 

61 text=f"echo({self._name}): {text}", 

62 token_usage=TokenUsage(prompt=prompt_tokens, completion=completion_tokens), 

63 latency_ms=latency, 

64 ) 

65 

66 

67__all__ = ["MockProvider"]