File size: 1,857 Bytes
95f8c4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | 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"}},
]
|