| import json |
|
|
| import pytest |
|
|
| from aistudio_api.api.responses import ( |
| to_gemini_parts, |
| to_gemini_usage_metadata, |
| ) |
|
|
|
|
|
|
| def test_to_gemini_usage_metadata_uses_visible_and_reasoning_tokens(): |
| assert to_gemini_usage_metadata( |
| { |
| "prompt_tokens": 9, |
| "completion_tokens": 316, |
| "total_tokens": 325, |
| "completion_tokens_details": {"reasoning_tokens": 290, "visible_tokens": 26}, |
| } |
| ).model_dump(mode="json") == { |
| "promptTokenCount": 9, |
| "candidatesTokenCount": 26, |
| "thoughtsTokenCount": 290, |
| "totalTokenCount": 325, |
| } |
|
|
|
|
| def test_to_gemini_parts_keeps_function_call_and_response_parts(): |
| parts = to_gemini_parts( |
| "", |
| function_calls=[{"name": "getWeather", "args": {"city": "Shanghai"}}], |
| function_responses=[{"name": "getWeather", "args": {"temperature": "24C"}}], |
| ) |
|
|
| assert [part.model_dump(mode="json", exclude_none=True) for part in parts] == [ |
| {"functionCall": {"name": "getWeather", "args": {"city": "Shanghai"}}}, |
| {"functionResponse": {"name": "getWeather", "response": {"temperature": "24C"}}}, |
| ] |
|
|
|
|
| def test_to_gemini_parts_can_emit_thought_part(): |
| assert [part.model_dump(mode="json", exclude_none=True) for part in to_gemini_parts("答案", thinking="思考")] == [ |
| {"text": "思考", "thought": True}, |
| {"text": "答案"}, |
| ] |
|
|
|
|
| def test_to_gemini_parts_can_emit_inline_image_data(): |
| from aistudio_api.domain.models import GeneratedImage |
|
|
| parts = to_gemini_parts( |
| "说明", |
| images=[GeneratedImage(mime="image/png", data=b"png-bytes", size=9)], |
| ) |
|
|
| assert [part.model_dump(mode="json", exclude_none=True) for part in parts] == [ |
| {"text": "说明"}, |
| {"inlineData": {"mimeType": "image/png", "data": "cG5nLWJ5dGVz"}}, |
| ] |
|
|
