Coverage for pydantic_ai_slim/pydantic_ai/models/wrapper.py: 96.55%

1from __future__ import annotations (empty)

2 

3from collections.abc import AsyncIterator (empty)

4from contextlib import asynccontextmanager (empty)

5from dataclasses import dataclass (empty)

6from typing import Any (empty)

7 

8from ..messages import ModelMessage, ModelResponse (empty)

9from ..settings import ModelSettings (empty)

10from ..usage import Usage (empty)

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

12 

13 

14@dataclass(init=False) (empty)

15class WrapperModel(Model): (empty)

16 """Model which wraps another model. 

17 

18 Does nothing on its own, used as a base class. 

19 """ 

20 

21 wrapped: Model (empty)

22 """The underlying model being wrapped.""" (empty)

23 

24 def __init__(self, wrapped: Model | KnownModelName): (empty)

25 self.wrapped = infer_model(wrapped) (empty)

26 

27 async def request(self, *args: Any, **kwargs: Any) -> tuple[ModelResponse, Usage]: (empty)

28 return await self.wrapped.request(*args, **kwargs) (empty)

29 

30 @asynccontextmanager (empty)

31 async def request_stream( (empty)

32 self, 

33 messages: list[ModelMessage], 

34 model_settings: ModelSettings | None, 

35 model_request_parameters: ModelRequestParameters, 

36 ) -> AsyncIterator[StreamedResponse]: 

37 async with self.wrapped.request_stream(messages, model_settings, model_request_parameters) as response_stream: (empty)

38 yield response_stream (empty)

39 

40 @property (empty)

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

42 return self.wrapped.model_name (empty)

43 

44 @property (empty)

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

46 return self.wrapped.system (empty)

47 

48 def __getattr__(self, item: str): (empty)

49 return getattr(self.wrapped, item)