Coverage for pydantic_ai_slim/pydantic_ai/models/fallback.py: 96.34%

1from __future__ import annotations as _annotations (empty)

2 

3from collections.abc import AsyncIterator (empty)

4from contextlib import AsyncExitStack, asynccontextmanager, suppress (empty)

5from dataclasses import dataclass, field (empty)

6from typing import TYPE_CHECKING, Callable (empty)

7 

8from opentelemetry.trace import get_current_span (empty)

9 

10from pydantic_ai.models.instrumented import InstrumentedModel (empty)

11 

12from ..exceptions import FallbackExceptionGroup, ModelHTTPError (empty)

13from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse, infer_model (empty)

14 

15if TYPE_CHECKING: (empty)

16 from ..messages import ModelMessage, ModelResponse 

17 from ..settings import ModelSettings 

18 from ..usage import Usage 

19 

20 

21@dataclass(init=False) (empty)

22class FallbackModel(Model): (empty)

23 """A model that uses one or more fallback models upon failure. 

24 

25 Apart from `__init__`, all methods are private or match those of the base class. 

26 """ 

27 

28 models: list[Model] (empty)

29 

30 _model_name: str = field(repr=False) (empty)

31 _fallback_on: Callable[[Exception], bool] (empty)

32 

33 def __init__( (empty)

34 self, 

35 default_model: Model | KnownModelName, 

36 *fallback_models: Model | KnownModelName, 

37 fallback_on: Callable[[Exception], bool] | tuple[type[Exception], ...] = (ModelHTTPError,), 

38 ): 

39 """Initialize a fallback model instance. 

40 

41 Args: 

42 default_model: The name or instance of the default model to use. 

43 fallback_models: The names or instances of the fallback models to use upon failure. 

44 fallback_on: A callable or tuple of exceptions that should trigger a fallback. 

45 """ 

46 self.models = [infer_model(default_model), *[infer_model(m) for m in fallback_models]] (empty)

47 

48 if isinstance(fallback_on, tuple): (empty)

49 self._fallback_on = _default_fallback_condition_factory(fallback_on) (empty)

50 else: 

51 self._fallback_on = fallback_on (empty)

52 

53 async def request( (empty)

54 self, 

55 messages: list[ModelMessage], 

56 model_settings: ModelSettings | None, 

57 model_request_parameters: ModelRequestParameters, 

58 ) -> tuple[ModelResponse, Usage]: 

59 """Try each model in sequence until one succeeds. 

60 

61 In case of failure, raise a FallbackExceptionGroup with all exceptions. 

62 """ 

63 exceptions: list[Exception] = [] (empty)

64 

65 for model in self.models: (empty)

66 try: (empty)

67 response, usage = await model.request(messages, model_settings, model_request_parameters) (empty)

68 except Exception as exc: (empty)

69 if self._fallback_on(exc): (empty)

70 exceptions.append(exc) (empty)

71 continue (empty)

72 raise exc (empty)

73 

74 self._set_span_attributes(model) (empty)

75 return response, usage (empty)

76 

77 raise FallbackExceptionGroup('All models from FallbackModel failed', exceptions) (empty)

78 

79 @asynccontextmanager (empty)

80 async def request_stream( (empty)

81 self, 

82 messages: list[ModelMessage], 

83 model_settings: ModelSettings | None, 

84 model_request_parameters: ModelRequestParameters, 

85 ) -> AsyncIterator[StreamedResponse]: 

86 """Try each model in sequence until one succeeds.""" 

87 exceptions: list[Exception] = [] (empty)

88 

89 for model in self.models: (empty)

90 async with AsyncExitStack() as stack: (empty)

91 try: (empty)

92 response = await stack.enter_async_context( (empty)

93 model.request_stream(messages, model_settings, model_request_parameters) 

94 ) 

95 except Exception as exc: (empty)

96 if self._fallback_on(exc): 96 ↛ 99line 96 didn't jump to line 99 because the condition on line 96 was always true(empty)

97 exceptions.append(exc) (empty)

98 continue (empty)

99 raise exc 

100 

101 self._set_span_attributes(model) (empty)

102 yield response (empty)

103 return (empty)

104 

105 raise FallbackExceptionGroup('All models from FallbackModel failed', exceptions) (empty)

106 

107 def _set_span_attributes(self, model: Model): (empty)

108 with suppress(Exception): (empty)

109 span = get_current_span() (empty)

110 if span.is_recording(): (empty)

111 attributes = getattr(span, 'attributes', {}) (empty)

112 if attributes.get('gen_ai.request.model') == self.model_name: 112 ↛ exitline 112 didn't jump to the function exit(empty)

113 span.set_attributes(InstrumentedModel.model_attributes(model)) (empty)

114 

115 @property (empty)

116 def model_name(self) -> str: (empty)

117 """The model name.""" 

118 return f'fallback:{",".join(model.model_name for model in self.models)}' (empty)

119 

120 @property (empty)

121 def system(self) -> str: (empty)

122 return f'fallback:{",".join(model.system for model in self.models)}' (empty)

123 

124 @property (empty)

125 def base_url(self) -> str | None: (empty)

126 return self.models[0].base_url (empty)

127 

128 

129def _default_fallback_condition_factory(exceptions: tuple[type[Exception], ...]) -> Callable[[Exception], bool]: (empty)

130 """Create a default fallback condition for the given exceptions.""" 

131 

132 def fallback_condition(exception: Exception) -> bool: (empty)

133 return isinstance(exception, exceptions) (empty)

134 

135 return fallback_condition (empty)