pydantic-ai-slim 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

pydantic_ai/models/bedrock.py

@@ -252,7 +252,7 @@ class BedrockConverseModel(Model):
 
     async def _process_response(self, response: ConverseResponseTypeDef) -> ModelResponse:
         items: list[ModelResponsePart] = []
-        if message := response['output'].get('message'):
+        if message := response['output'].get('message'):  # pragma: no branch
             for item in message['content']:
                 if text := item.get('text'):
                     items.append(TextPart(content=text))
@@ -271,7 +271,8 @@ class BedrockConverseModel(Model):
             response_tokens=response['usage']['outputTokens'],
             total_tokens=response['usage']['totalTokens'],
         )
-        return ModelResponse(items, usage=u, model_name=self.model_name)
+        vendor_id = response.get('ResponseMetadata', {}).get('RequestId', None)
+        return ModelResponse(items, usage=u, model_name=self.model_name, vendor_id=vendor_id)
 
     @overload
     async def _messages_create(
@@ -305,7 +306,7 @@ class BedrockConverseModel(Model):
         if not tools or not support_tools_choice:
             tool_choice: ToolChoiceTypeDef = {}
         elif not model_request_parameters.allow_text_output:
-            tool_choice = {'any': {}}
+            tool_choice = {'any': {}}  # pragma: no cover
         else:
             tool_choice = {'auto': {}}
 
@@ -492,7 +493,7 @@ class BedrockConverseModel(Model):
                         data = response.content
                         content.append({'document': {'name': name, 'format': item.format, 'source': {'bytes': data}}})
 
-                    elif item.kind == 'video-url':
+                    elif item.kind == 'video-url':  # pragma: no branch
                         format = item.media_type.split('/')[1]
                         assert format in ('mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp'), (
                             f'Unsupported video format: {format}'
@@ -535,13 +536,13 @@ class BedrockStreamedResponse(StreamedResponse):
             if 'messageStop' in chunk:
                 continue
             if 'metadata' in chunk:
-                if 'usage' in chunk['metadata']:
+                if 'usage' in chunk['metadata']:  # pragma: no branch
                     self._usage += self._map_usage(chunk['metadata'])
                 continue
             if 'contentBlockStart' in chunk:
                 index = chunk['contentBlockStart']['contentBlockIndex']
                 start = chunk['contentBlockStart']['start']
-                if 'toolUse' in start:
+                if 'toolUse' in start:  # pragma: no branch
                     tool_use_start = start['toolUse']
                     tool_id = tool_use_start['toolUseId']
                     tool_name = tool_use_start['name']
@@ -552,7 +553,7 @@ class BedrockStreamedResponse(StreamedResponse):
                         tool_call_id=tool_id,
                     )
                     if maybe_event:
-                        yield maybe_event
+                        yield maybe_event  # pragma: no cover
             if 'contentBlockDelta' in chunk:
                 index = chunk['contentBlockDelta']['contentBlockIndex']
                 delta = chunk['contentBlockDelta']['delta']
@@ -566,7 +567,7 @@ class BedrockStreamedResponse(StreamedResponse):
                         args=tool_use.get('input'),
                         tool_call_id=tool_id,
                     )
-                    if maybe_event:
+                    if maybe_event:  # pragma: no branch
                         yield maybe_event
 
     @property
@@ -602,4 +603,4 @@ class _AsyncIteratorWrapper(Generic[T]):
             if type(e.__cause__) is StopIteration:
                 raise StopAsyncIteration
             else:
-                raise e
+                raise e  # pragma: lax no cover

pydantic_ai/models/cohere.py

@@ -174,7 +174,7 @@ class CohereModel(Model):
         except ApiError as e:
             if (status_code := e.status_code) and status_code >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
-            raise
+            raise  # pragma: lax no cover
 
     def _process_response(self, response: ChatResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
@@ -185,7 +185,7 @@ class CohereModel(Model):
             choice = response.message.content[0]
             parts.append(TextPart(choice.text))
         for c in response.message.tool_calls or []:
-            if c.function and c.function.name and c.function.arguments:
+            if c.function and c.function.name and c.function.arguments:  # pragma: no branch
                 parts.append(
                     ToolCallPart(
                         tool_name=c.function.name,
@@ -269,7 +269,7 @@ class CohereModel(Model):
                 )
             elif isinstance(part, RetryPromptPart):
                 if part.tool_name is None:
-                    yield UserChatMessageV2(role='user', content=part.model_response())
+                    yield UserChatMessageV2(role='user', content=part.model_response())  # pragma: no cover
                 else:
                     yield ToolChatMessageV2(
                         role='tool',
@@ -287,7 +287,7 @@ def _map_usage(response: ChatResponse) -> usage.Usage:
     else:
         details: dict[str, int] = {}
         if u.billed_units is not None:
-            if u.billed_units.input_tokens:
+            if u.billed_units.input_tokens:  # pragma: no branch
                 details['input_tokens'] = int(u.billed_units.input_tokens)
             if u.billed_units.output_tokens:
                 details['output_tokens'] = int(u.billed_units.output_tokens)

pydantic_ai/models/fallback.py

@@ -96,7 +96,7 @@ class FallbackModel(Model):
                     if self._fallback_on(exc):
                         exceptions.append(exc)
                         continue
-                    raise exc
+                    raise exc  # pragma: no cover
 
                 self._set_span_attributes(model)
                 yield response
@@ -109,7 +109,7 @@ class FallbackModel(Model):
             span = get_current_span()
             if span.is_recording():
                 attributes = getattr(span, 'attributes', {})
-                if attributes.get('gen_ai.request.model') == self.model_name:
+                if attributes.get('gen_ai.request.model') == self.model_name:  # pragma: no branch
                     span.set_attributes(InstrumentedModel.model_attributes(model))
 
     @property

pydantic_ai/models/function.py

@@ -106,7 +106,7 @@ class FunctionModel(Model):
             response = response_
         response.model_name = self._model_name
         # Add usage data if not already present
-        if not response.usage.has_values():
+        if not response.usage.has_values():  # pragma: no branch
             response.usage = _estimate_usage(chain(messages, [response]))
             response.usage.requests = 1
         return response

pydantic_ai/models/gemini.py

@@ -137,8 +137,8 @@ class GeminiModel(Model):
 
     @property
     def base_url(self) -> str:
-        assert self._url is not None, 'URL not initialized'
-        return self._url
+        assert self._url is not None, 'URL not initialized'  # pragma: no cover
+        return self._url  # pragma: no cover
 
     async def request(
         self,
@@ -201,7 +201,7 @@ class GeminiModel(Model):
         elif tools:
             return _tool_config([t['name'] for t in tools['function_declarations']])
         else:
-            return _tool_config([])
+            return _tool_config([])  # pragma: no cover
 
     @asynccontextmanager
     async def _make_request(
@@ -257,17 +257,21 @@ class GeminiModel(Model):
                 await r.aread()
                 if status_code >= 400:
                     raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=r.text)
-                raise UnexpectedModelBehavior(f'Unexpected response from gemini {status_code}', r.text)
+                raise UnexpectedModelBehavior(  # pragma: no cover
+                    f'Unexpected response from gemini {status_code}', r.text
+                )
             yield r
 
     def _process_response(self, response: _GeminiResponse) -> ModelResponse:
         if len(response['candidates']) != 1:
-            raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')
+            raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')  # pragma: no cover
         if 'content' not in response['candidates'][0]:
             if response['candidates'][0].get('finish_reason') == 'SAFETY':
                 raise UnexpectedModelBehavior('Safety settings triggered', str(response))
             else:
-                raise UnexpectedModelBehavior('Content field missing from Gemini response', str(response))
+                raise UnexpectedModelBehavior(  # pragma: no cover
+                    'Content field missing from Gemini response', str(response)
+                )
         parts = response['candidates'][0]['content']['parts']
         usage = _metadata_as_usage(response)
         usage.requests = 1
@@ -285,7 +289,7 @@ class GeminiModel(Model):
                 _ensure_decodeable(content),
                 experimental_allow_partial='trailing-strings',
             )
-            if responses:
+            if responses:  # pragma: no branch
                 last = responses[-1]
                 if last['candidates'] and last['candidates'][0].get('content', {}).get('parts'):
                     start_response = last
@@ -314,14 +318,14 @@ class GeminiModel(Model):
                         message_parts.append(_response_part_from_response(part.tool_name, part.model_response_object()))
                     elif isinstance(part, RetryPromptPart):
                         if part.tool_name is None:
-                            message_parts.append(_GeminiTextPart(text=part.model_response()))
+                            message_parts.append(_GeminiTextPart(text=part.model_response()))  # pragma: no cover
                         else:
                             response = {'call_error': part.model_response()}
                             message_parts.append(_response_part_from_response(part.tool_name, response))
                     else:
                         assert_never(part)
 
-                if message_parts:
+                if message_parts:  # pragma: no branch
                     contents.append(_GeminiContent(role='user', parts=message_parts))
             elif isinstance(m, ModelResponse):
                 contents.append(_content_model_response(m))
@@ -372,7 +376,7 @@ class ApiKeyAuth:
 
     async def headers(self) -> dict[str, str]:
         # https://cloud.google.com/docs/authentication/api-keys-use#using-with-rest
-        return {'X-Goog-Api-Key': self.api_key}
+        return {'X-Goog-Api-Key': self.api_key}  # pragma: no cover
 
 
 @dataclass
@@ -388,7 +392,7 @@ class GeminiStreamedResponse(StreamedResponse):
         async for gemini_response in self._get_gemini_responses():
             candidate = gemini_response['candidates'][0]
             if 'content' not in candidate:
-                raise UnexpectedModelBehavior('Streamed response has no content field')
+                raise UnexpectedModelBehavior('Streamed response has no content field')  # pragma: no cover
             gemini_part: _GeminiPartUnion
             for gemini_part in candidate['content']['parts']:
                 if 'text' in gemini_part:
@@ -407,10 +411,10 @@ class GeminiStreamedResponse(StreamedResponse):
                         args=gemini_part['function_call']['args'],
                         tool_call_id=None,
                     )
-                    if maybe_event is not None:
+                    if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
                 else:
-                    assert 'function_response' in gemini_part, f'Unexpected part: {gemini_part}'
+                    assert 'function_response' in gemini_part, f'Unexpected part: {gemini_part}'  # pragma: no cover
 
     async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]:
         # This method exists to ensure we only yield completed items, so we don't need to worry about
@@ -438,7 +442,7 @@ class GeminiStreamedResponse(StreamedResponse):
                 yield r
 
         # Now yield the final response, which should be complete
-        if gemini_responses:
+        if gemini_responses:  # pragma: no branch
             r = gemini_responses[-1]
             self._usage += _metadata_as_usage(r)
             yield r
@@ -601,7 +605,7 @@ def _process_response_from_parts(
             items.append(TextPart(content=part['text']))
         elif 'function_call' in part:
             items.append(ToolCallPart(tool_name=part['function_call']['name'], args=part['function_call']['args']))
-        elif 'function_response' in part:
+        elif 'function_response' in part:  # pragma: no cover
             raise UnexpectedModelBehavior(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
             )
@@ -631,13 +635,13 @@ class _GeminiFunctionResponse(TypedDict):
 
 
 def _part_discriminator(v: Any) -> str:
-    if isinstance(v, dict):
+    if isinstance(v, dict):  # pragma: no branch
         if 'text' in v:
             return 'text'
         elif 'inlineData' in v:
-            return 'inline_data'
+            return 'inline_data'  # pragma: no cover
         elif 'fileData' in v:
-            return 'file_data'
+            return 'file_data'  # pragma: no cover
         elif 'functionCall' in v or 'function_call' in v:
             return 'function_call'
         elif 'functionResponse' in v or 'function_response' in v:
@@ -748,10 +752,10 @@ class _GeminiUsageMetaData(TypedDict, total=False):
 def _metadata_as_usage(response: _GeminiResponse) -> usage.Usage:
     metadata = response.get('usage_metadata')
     if metadata is None:
-        return usage.Usage()
+        return usage.Usage()  # pragma: no cover
     details: dict[str, int] = {}
     if cached_content_token_count := metadata.get('cached_content_token_count'):
-        details['cached_content_token_count'] = cached_content_token_count
+        details['cached_content_token_count'] = cached_content_token_count  # pragma: no cover
     return usage.Usage(
         request_tokens=metadata.get('prompt_token_count', 0),
         response_tokens=metadata.get('candidates_token_count', 0),
@@ -866,10 +870,10 @@ class _GeminiJsonSchema(WalkJsonSchema):
                     unique_items.append(item)
             if len(unique_items) > 1:  # pragma: no cover
                 schema['items'] = {'anyOf': unique_items}
-            elif len(unique_items) == 1:
+            elif len(unique_items) == 1:  # pragma: no branch
                 schema['items'] = unique_items[0]
             schema.setdefault('minItems', len(prefix_items))
-            if items is None:
+            if items is None:  # pragma: no branch
                 schema.setdefault('maxItems', len(prefix_items))
 
         return schema