Skip to content

Commit e07b58f

Browse files
feat: added finish_reason
1 parent b03cd97 commit e07b58f

File tree

1 file changed

+14
-3
lines changed

1 file changed

+14
-3
lines changed

pydantic_ai_slim/pydantic_ai/models/gemini.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -274,10 +274,15 @@ def _process_response(self, response: _GeminiResponse) -> ModelResponse:
274274
)
275275
parts = response['candidates'][0]['content']['parts']
276276
vendor_id = response.get('vendor_id', None)
277+
vendor_details = {'finish_reason': response['candidates'][0].get('finish_reason')}
277278
usage = _metadata_as_usage(response)
278279
usage.requests = 1
279280
return _process_response_from_parts(
280-
parts, response.get('model_version', self._model_name), usage, vendor_id=vendor_id
281+
parts,
282+
response.get('model_version', self._model_name),
283+
usage,
284+
vendor_id=vendor_id,
285+
vendor_details=vendor_details,
281286
)
282287

283288
async def _process_streamed_response(self, http_response: HTTPResponse) -> StreamedResponse:
@@ -600,7 +605,11 @@ def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart
600605

601606

602607
def _process_response_from_parts(
603-
parts: Sequence[_GeminiPartUnion], model_name: GeminiModelName, usage: usage.Usage, vendor_id: str | None
608+
parts: Sequence[_GeminiPartUnion],
609+
model_name: GeminiModelName,
610+
usage: usage.Usage,
611+
vendor_id: str | None,
612+
vendor_details: dict[str, Any] | None,
604613
) -> ModelResponse:
605614
items: list[ModelResponsePart] = []
606615
for part in parts:
@@ -612,7 +621,9 @@ def _process_response_from_parts(
612621
raise UnexpectedModelBehavior(
613622
f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
614623
)
615-
return ModelResponse(parts=items, usage=usage, model_name=model_name, vendor_id=vendor_id)
624+
return ModelResponse(
625+
parts=items, usage=usage, model_name=model_name, vendor_id=vendor_id, vendor_details=vendor_details
626+
)
616627

617628

618629
class _GeminiFunctionCall(TypedDict):

0 commit comments

Comments
 (0)