test_gemini_native_adapter.py
1 """Tests for the native Google AI Studio Gemini adapter.""" 2 3 from __future__ import annotations 4 5 import json 6 from types import SimpleNamespace 7 8 import pytest 9 10 11 class DummyResponse: 12 def __init__(self, status_code=200, payload=None, headers=None, text=None): 13 self.status_code = status_code 14 self._payload = payload or {} 15 self.headers = headers or {} 16 self.text = text if text is not None else json.dumps(self._payload) 17 18 def json(self): 19 return self._payload 20 21 22 def test_build_native_request_preserves_thought_signature_on_tool_replay(): 23 from agent.gemini_native_adapter import build_gemini_request 24 25 request = build_gemini_request( 26 messages=[ 27 {"role": "system", "content": "Be helpful."}, 28 { 29 "role": "assistant", 30 "content": "", 31 "tool_calls": [ 32 { 33 "id": "call_1", 34 "type": "function", 35 "function": { 36 "name": "get_weather", 37 "arguments": '{"city": "Paris"}', 38 }, 39 "extra_content": { 40 "google": {"thought_signature": "sig-123"} 41 }, 42 } 43 ], 44 }, 45 ], 46 tools=[], 47 tool_choice=None, 48 ) 49 50 parts = request["contents"][0]["parts"] 51 assert parts[0]["functionCall"]["name"] == "get_weather" 52 assert parts[0]["thoughtSignature"] == "sig-123" 53 54 55 def test_build_native_request_uses_original_function_name_for_tool_result(): 56 from agent.gemini_native_adapter import build_gemini_request 57 58 request = build_gemini_request( 59 messages=[ 60 { 61 "role": "assistant", 62 "content": "", 63 "tool_calls": [ 64 { 65 "id": "call_1", 66 "type": "function", 67 "function": { 68 "name": "get_weather", 69 "arguments": '{"city": "Paris"}', 70 }, 71 } 72 ], 73 }, 74 { 75 "role": "tool", 76 "tool_call_id": "call_1", 77 "content": '{"forecast": "sunny"}', 78 }, 79 ], 80 tools=[], 81 tool_choice=None, 82 ) 83 84 tool_response = request["contents"][1]["parts"][0]["functionResponse"] 85 assert tool_response["name"] == "get_weather" 86 87 88 def test_build_native_request_strips_json_schema_only_fields_from_tool_parameters(): 89 from agent.gemini_native_adapter import build_gemini_request 90 91 request = build_gemini_request( 92 messages=[{"role": "user", "content": "Hello"}], 93 tools=[ 94 { 95 "type": "function", 96 "function": { 97 "name": "lookup_weather", 98 "description": "Weather lookup", 99 "parameters": { 100 "$schema": "https://json-schema.org/draft/2020-12/schema", 101 "type": "object", 102 "additionalProperties": False, 103 "properties": { 104 "city": { 105 "type": "string", 106 "$schema": "ignored", 107 "description": "City name", 108 } 109 }, 110 "required": ["city"], 111 }, 112 }, 113 } 114 ], 115 tool_choice=None, 116 ) 117 118 params = request["tools"][0]["functionDeclarations"][0]["parameters"] 119 assert "$schema" not in params 120 assert "additionalProperties" not in params 121 assert params["type"] == "object" 122 assert params["properties"]["city"] == { 123 "type": "string", 124 "description": "City name", 125 } 126 127 128 def test_translate_native_response_surfaces_reasoning_and_tool_calls(): 129 from agent.gemini_native_adapter import translate_gemini_response 130 131 payload = { 132 "candidates": [ 133 { 134 "content": { 135 "parts": [ 136 {"thought": True, "text": "thinking..."}, 137 {"functionCall": {"name": "search", "args": {"q": "hermes"}}}, 138 ] 139 }, 140 "finishReason": "STOP", 141 } 142 ], 143 "usageMetadata": { 144 "promptTokenCount": 10, 145 "candidatesTokenCount": 5, 146 "totalTokenCount": 15, 147 }, 148 } 149 150 response = translate_gemini_response(payload, model="gemini-2.5-flash") 151 choice = response.choices[0] 152 assert choice.finish_reason == "tool_calls" 153 assert choice.message.reasoning == "thinking..." 154 assert choice.message.tool_calls[0].function.name == "search" 155 assert json.loads(choice.message.tool_calls[0].function.arguments) == {"q": "hermes"} 156 157 158 def test_native_client_uses_x_goog_api_key_and_native_models_endpoint(monkeypatch): 159 from agent.gemini_native_adapter import GeminiNativeClient 160 161 recorded = {} 162 163 class DummyHTTP: 164 def post(self, url, json=None, headers=None, timeout=None): 165 recorded["url"] = url 166 recorded["json"] = json 167 recorded["headers"] = headers 168 return DummyResponse( 169 payload={ 170 "candidates": [ 171 { 172 "content": {"parts": [{"text": "hello"}]}, 173 "finishReason": "STOP", 174 } 175 ], 176 "usageMetadata": { 177 "promptTokenCount": 1, 178 "candidatesTokenCount": 1, 179 "totalTokenCount": 2, 180 }, 181 } 182 ) 183 184 def close(self): 185 return None 186 187 monkeypatch.setattr("agent.gemini_native_adapter.httpx.Client", lambda *a, **k: DummyHTTP()) 188 189 client = GeminiNativeClient(api_key="AIza-test", base_url="https://generativelanguage.googleapis.com/v1beta") 190 response = client.chat.completions.create( 191 model="gemini-2.5-flash", 192 messages=[{"role": "user", "content": "Hello"}], 193 ) 194 195 assert recorded["url"] == "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent" 196 assert recorded["headers"]["x-goog-api-key"] == "AIza-test" 197 assert "Authorization" not in recorded["headers"] 198 assert response.choices[0].message.content == "hello" 199 200 201 def test_native_http_error_keeps_status_and_retry_after(): 202 from agent.gemini_native_adapter import gemini_http_error 203 204 response = DummyResponse( 205 status_code=429, 206 headers={"Retry-After": "17"}, 207 payload={ 208 "error": { 209 "code": 429, 210 "message": "quota exhausted", 211 "status": "RESOURCE_EXHAUSTED", 212 "details": [ 213 { 214 "@type": "type.googleapis.com/google.rpc.ErrorInfo", 215 "reason": "RESOURCE_EXHAUSTED", 216 "metadata": {"service": "generativelanguage.googleapis.com"}, 217 } 218 ], 219 } 220 }, 221 ) 222 223 err = gemini_http_error(response) 224 assert getattr(err, "status_code", None) == 429 225 assert getattr(err, "retry_after", None) == 17.0 226 assert "quota exhausted" in str(err) 227 228 229 def test_native_client_accepts_injected_http_client(): 230 from agent.gemini_native_adapter import GeminiNativeClient 231 232 injected = SimpleNamespace(close=lambda: None) 233 client = GeminiNativeClient(api_key="AIza-test", http_client=injected) 234 assert client._http is injected 235 236 237 def test_native_client_rejects_empty_api_key_with_actionable_message(): 238 """Empty/whitespace api_key must raise at construction, not produce a cryptic 239 Google GFE 'Error 400 (Bad Request)!!1' HTML page on the first request.""" 240 from agent.gemini_native_adapter import GeminiNativeClient 241 242 for bad in ("", " ", None): 243 with pytest.raises(RuntimeError) as excinfo: 244 GeminiNativeClient(api_key=bad) # type: ignore[arg-type] 245 msg = str(excinfo.value) 246 assert "GOOGLE_API_KEY" in msg and "GEMINI_API_KEY" in msg 247 assert "aistudio.google.com" in msg 248 249 250 @pytest.mark.asyncio 251 async def test_async_native_client_streams_without_requiring_async_iterator_from_sync_client(): 252 from agent.gemini_native_adapter import AsyncGeminiNativeClient 253 254 chunk = SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="hi"), finish_reason=None)]) 255 sync_stream = iter([chunk]) 256 257 def _advance(iterator): 258 try: 259 return False, next(iterator) 260 except StopIteration: 261 return True, None 262 263 sync_client = SimpleNamespace( 264 api_key="AIza-test", 265 base_url="https://generativelanguage.googleapis.com/v1beta", 266 chat=SimpleNamespace(completions=SimpleNamespace(create=lambda **kwargs: sync_stream)), 267 _advance_stream_iterator=_advance, 268 close=lambda: None, 269 ) 270 271 async_client = AsyncGeminiNativeClient(sync_client) 272 stream = await async_client.chat.completions.create(stream=True) 273 collected = [] 274 async for item in stream: 275 collected.append(item) 276 assert collected == [chunk] 277 278 279 def test_stream_event_translation_emits_tool_call_delta_with_stable_index(): 280 from agent.gemini_native_adapter import translate_stream_event 281 282 tool_call_indices = {} 283 event = { 284 "candidates": [ 285 { 286 "content": { 287 "parts": [ 288 {"functionCall": {"name": "search", "args": {"q": "abc"}}} 289 ] 290 }, 291 "finishReason": "STOP", 292 } 293 ] 294 } 295 296 first = translate_stream_event(event, model="gemini-2.5-flash", tool_call_indices=tool_call_indices) 297 second = translate_stream_event(event, model="gemini-2.5-flash", tool_call_indices=tool_call_indices) 298 299 assert first[0].choices[0].delta.tool_calls[0].index == 0 300 assert second[0].choices[0].delta.tool_calls[0].index == 0 301 assert first[0].choices[0].delta.tool_calls[0].id == second[0].choices[0].delta.tool_calls[0].id 302 assert first[0].choices[0].delta.tool_calls[0].function.arguments == '{"q": "abc"}' 303 assert second[0].choices[0].delta.tool_calls[0].function.arguments == "" 304 assert first[-1].choices[0].finish_reason == "tool_calls" 305 306 307 def test_stream_event_translation_keeps_identical_calls_in_distinct_parts(): 308 from agent.gemini_native_adapter import translate_stream_event 309 310 event = { 311 "candidates": [ 312 { 313 "content": { 314 "parts": [ 315 {"functionCall": {"name": "search", "args": {"q": "abc"}}}, 316 {"functionCall": {"name": "search", "args": {"q": "abc"}}}, 317 ] 318 }, 319 "finishReason": "STOP", 320 } 321 ] 322 } 323 324 chunks = translate_stream_event(event, model="gemini-2.5-flash", tool_call_indices={}) 325 tool_chunks = [chunk for chunk in chunks if chunk.choices[0].delta.tool_calls] 326 assert tool_chunks[0].choices[0].delta.tool_calls[0].index == 0 327 assert tool_chunks[1].choices[0].delta.tool_calls[0].index == 1 328 assert tool_chunks[0].choices[0].delta.tool_calls[0].id != tool_chunks[1].choices[0].delta.tool_calls[0].id