Python 88.3%
TypeScript 7.6%
Shell 4.1%
1"""Offline tests for examples/shared/provider-abstraction/llm_provider.py — mock transport, no network, no keys."""2from __future__ import annotations34import json5import os6import sys7from pathlib import Path89import pytest1011ROOT = Path(__file__).resolve().parents[2]12sys.path.insert(0, str(ROOT / "examples" / "shared" / "resilient-client"))13sys.path.insert(0, str(ROOT / "examples" / "shared" / "provider-abstraction"))14os.environ["ATLAS_DISABLE_LIVE_LOG"] = "1"1516import resilient_client as rc # noqa: E40217import llm_provider as lp # noqa: E402181920class Script:21 def __init__(self, responses):22 self.responses, self.requests = list(responses), []2324 def __call__(self, req: rc.Request) -> rc.Response:25 self.requests.append(req)26 return self.responses.pop(0)2728 def body(self, i=0):29 return json.loads(self.requests[i].body)303132def ok(body, headers=None):33 return rc.Response(200, headers or {}, json.dumps(body).encode())343536def sse(events):37 out = []38 for e in events:39 out += [f"event: {e['type']}\n".encode(), f"data: {json.dumps(e)}\n".encode(), b"\n"]40 return rc.Response(200, {}, lines=iter(out))414243def provider(name, script):44 return lp.get_provider(name, rc.ResilientClient(name, api_key="k", transport=script))454647TOOL = lp.ToolDef("get_weather", "Weather by city", {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]})484950# ---------------- request mapping ----------------5152def test_openai_body_mapping():53 s = Script([ok({"id": "resp_1", "status": "completed", "output": [], "usage": {}})])54 p = provider("openai", s)55 req = lp.GenerateRequest("gpt-5.4-nano", [lp.Message("user", "hi")], system="be brief", max_tokens=16, temperature=0.2,56 tools=[TOOL], tool_choice={"name": "get_weather"}, json_schema={"name": "out", "schema": {"type": "object"}},57 stop=["END"], extensions={"reasoning": {"effort": "low"}, "store": False, "previous_response_id": "resp_0"})58 p.generate(req)59 b = s.body()60 assert b["instructions"] == "be brief" and b["max_output_tokens"] == 16 and b["temperature"] == 0.261 assert b["input"] == [{"role": "user", "content": "hi"}]62 assert b["tools"][0] == {"type": "function", "name": "get_weather", "description": "Weather by city", "parameters": TOOL.parameters, "strict": False}63 assert b["tool_choice"] == {"type": "function", "name": "get_weather"}64 assert b["text"]["format"] == {"type": "json_schema", "name": "out", "schema": {"type": "object"}, "strict": True}65 assert "stop" not in b and "stop_sequences" not in b # lossy: Responses API has no stop sequences66 assert b["reasoning"] == {"effort": "low"} and b["store"] is False and b["previous_response_id"] == "resp_0"67 assert s.requests[0].headers["Authorization"] == "Bearer k"686970def test_anthropic_body_mapping():71 s = Script([ok({"id": "msg_1", "content": [], "stop_reason": "end_turn", "usage": {}})])72 p = provider("anthropic", s)73 req = lp.GenerateRequest("claude-haiku-4-5-20251001", [lp.Message("user", "hi")], system="be brief", max_tokens=16, temperature=0.2,74 tools=[TOOL], tool_choice="required", json_schema={"schema": {"type": "object"}}, stop=["END"],75 extensions={"thinking": {"type": "adaptive"}, "output_config": {"effort": "low", "format": {"type": "json_schema", "schema": {"type": "object"}}}},76 extra_headers={"anthropic-beta": "context-management-2025-06-27"})77 p.generate(req)78 b = s.body()79 assert b["system"] == "be brief" and b["max_tokens"] == 16 and b["messages"] == [{"role": "user", "content": "hi"}]80 assert b["tools"][0] == {"name": "get_weather", "description": "Weather by city", "input_schema": TOOL.parameters}81 assert b["tool_choice"] == {"type": "any"} and b["stop_sequences"] == ["END"]82 assert b["output_config"]["format"]["type"] == "json_schema" and b["output_config"]["effort"] == "low"83 assert b["thinking"] == {"type": "adaptive"}84 h = s.requests[0].headers85 assert h["x-api-key"] == "k" and h["anthropic-version"] == "2023-06-01" and h["anthropic-beta"] == "context-management-2025-06-27"868788def test_tool_choice_mapping_table():89 assert lp.AnthropicProvider(rc.ResilientClient("anthropic", api_key="k"))._tool_choice("auto") == {"type": "auto"}90 assert lp.AnthropicProvider(rc.ResilientClient("anthropic", api_key="k"))._tool_choice("none") == {"type": "none"}91 assert lp.OpenAIProvider(rc.ResilientClient("openai", api_key="k"))._tool_choice("required") == "required"929394# ---------------- response parsing ----------------9596def test_openai_parse_tool_call_and_usage():97 body = {"id": "resp_2", "model": "gpt-5.4-nano", "status": "completed",98 "output": [{"type": "function_call", "id": "fc_1", "call_id": "call_1", "name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}],99 "usage": {"input_tokens": 20, "output_tokens": 5, "total_tokens": 25, "input_tokens_details": {"cached_tokens": 8}}}100 r = provider("openai", Script([ok(body)])).generate(lp.GenerateRequest("gpt-5.4-nano", [lp.Message("user", "x")]))101 assert r.stop_reason == "tool_use" and r.tool_calls[0].id == "call_1" and r.tool_calls[0].arguments == {"city": "Paris"}102 assert r.usage.input_tokens == 20 and r.usage.cached_input_tokens == 8 and r.usage.total_tokens == 25103 assert r.response_id == "resp_2"104105106def test_openai_incomplete_max_tokens():107 body = {"status": "incomplete", "incomplete_details": {"reason": "max_output_tokens"}, "output": [{"type": "message", "content": [{"type": "output_text", "text": "par"}]}], "usage": {}}108 r = provider("openai", Script([ok(body)])).generate(lp.GenerateRequest("m", [lp.Message("user", "x")]))109 assert r.stop_reason == "max_tokens" and r.text == "par"110111112def test_anthropic_parse_tool_use_and_cache_usage():113 body = {"id": "msg_2", "model": "claude-haiku-4-5-20251001", "stop_reason": "tool_use",114 "content": [{"type": "text", "text": "Let me check."}, {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}],115 "usage": {"input_tokens": 12, "output_tokens": 7, "cache_read_input_tokens": 100, "cache_creation_input_tokens": 0}}116 r = provider("anthropic", Script([ok(body)])).generate(lp.GenerateRequest("m", [lp.Message("user", "x")]))117 assert r.stop_reason == "tool_use" and r.text == "Let me check." and r.tool_calls[0].id == "toolu_1"118 assert r.usage.input_tokens == 112 and r.usage.cached_input_tokens == 100 and r.usage.output_tokens == 7119120121# ---------------- tool loop (portable) ----------------122123def test_tool_loop_openai_round_trip_and_is_error():124 first = {"status": "completed", "output": [{"type": "function_call", "call_id": "c1", "name": "get_weather", "arguments": "{\"city\":\"Paris\"}"},125 {"type": "function_call", "call_id": "c2", "name": "missing", "arguments": "{}"}], "usage": {}}126 second = {"status": "completed", "output": [{"type": "message", "content": [{"type": "output_text", "text": "Sunny"}]}], "usage": {}}127 s = Script([ok(first), ok(second)])128 r = provider("openai", s).use_tools(lp.GenerateRequest("m", [lp.Message("user", "weather?")], tools=[TOOL]), {"get_weather": lambda city: f"{city}: 20C"})129 assert r.text == "Sunny"130 items = s.body(1)["input"]131 assert items[0] == {"role": "user", "content": "weather?"}132 assert items[1] == {"type": "function_call", "call_id": "c1", "name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}133 assert items[3] == {"type": "function_call_output", "call_id": "c1", "output": "Paris: 20C"}134 assert items[4]["output"].startswith("[TOOL ERROR] Unknown tool") # no is_error flag in Responses API → prefixed text135136137def test_tool_loop_anthropic_is_error_flag_and_exception_containment():138 first = {"stop_reason": "tool_use", "content": [{"type": "tool_use", "id": "t1", "name": "get_weather", "input": {"city": "Paris"}}], "usage": {}}139 second = {"stop_reason": "end_turn", "content": [{"type": "text", "text": "Unknown"}], "usage": {}}140 s = Script([ok(first), ok(second)])141142 def boom(city):143 raise RuntimeError("db down " + city)144145 r = provider("anthropic", s).use_tools(lp.GenerateRequest("m", [lp.Message("user", "weather?")], tools=[TOOL]), {"get_weather": boom})146 assert r.text == "Unknown"147 msgs = s.body(1)["messages"]148 assert msgs[1]["role"] == "assistant" and msgs[1]["content"][0]["type"] == "tool_use"149 block = msgs[2]["content"][0]150 assert msgs[2]["role"] == "user" and block["type"] == "tool_result" and block["tool_use_id"] == "t1" and block["is_error"] is True151 assert "db down Paris" in block["content"] and "Traceback" not in block["content"]152153154# ---------------- structured output ----------------155156def test_structured_output_both_and_refusal():157 oa = provider("openai", Script([ok({"status": "completed", "output": [{"type": "message", "content": [{"type": "output_text", "text": "{\"a\":1}"}]}], "usage": {}})]))158 assert oa.structured_output(lp.GenerateRequest("m", [lp.Message("user", "x")]), {"type": "object"}) == {"a": 1}159 an = provider("anthropic", Script([ok({"stop_reason": "end_turn", "content": [{"type": "text", "text": "{\"a\":2}"}], "usage": {}})]))160 assert an.structured_output(lp.GenerateRequest("m", [lp.Message("user", "x")]), {"type": "object"}) == {"a": 2}161 ref = provider("openai", Script([ok({"status": "completed", "output": [{"type": "message", "content": [{"type": "refusal", "refusal": "no"}]}], "usage": {}})]))162 with pytest.raises(ValueError):163 ref.structured_output(lp.GenerateRequest("m", [lp.Message("user", "x")]), {"type": "object"})164165166# ---------------- streaming normalisation ----------------167168def test_openai_stream_normalised():169 events = [170 {"type": "response.created", "sequence_number": 0, "response": {"id": "resp_9"}},171 {"type": "response.output_item.added", "sequence_number": 1, "output_index": 0, "item": {"id": "msg_1", "type": "message"}},172 {"type": "response.output_text.delta", "sequence_number": 2, "item_id": "msg_1", "delta": "O"},173 {"type": "response.output_text.delta", "sequence_number": 3, "item_id": "msg_1", "delta": "K"},174 {"type": "response.output_item.added", "sequence_number": 4, "output_index": 1, "item": {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": ""}},175 {"type": "response.function_call_arguments.delta", "sequence_number": 5, "item_id": "fc_1", "delta": "{\"city\":"},176 {"type": "response.function_call_arguments.delta", "sequence_number": 6, "item_id": "fc_1", "delta": "\"Paris\"}"},177 {"type": "response.function_call_arguments.done", "sequence_number": 7, "item_id": "fc_1", "arguments": "{\"city\":\"Paris\"}"},178 {"type": "response.completed", "sequence_number": 8, "response": {"id": "resp_9", "status": "completed", "output": [{"type": "function_call", "call_id": "call_1", "name": "get_weather", "arguments": "{}"}], "usage": {"input_tokens": 10, "output_tokens": 6, "total_tokens": 16}}},179 ]180 evs = list(provider("openai", Script([sse(events)])).stream(lp.GenerateRequest("m", [lp.Message("user", "x")])))181 types = [e.type for e in evs]182 assert types == ["raw", "raw", "text_delta", "text_delta", "tool_call_start", "tool_call_delta", "tool_call_delta", "tool_call_done", "usage", "done"]183 assert "".join(e.text for e in evs if e.type == "text_delta") == "OK"184 done_call = next(e for e in evs if e.type == "tool_call_done").tool_call185 assert done_call.name == "get_weather" and done_call.arguments == {"city": "Paris"}186 assert evs[-2].usage.total_tokens == 16 and evs[-1].stop_reason == "tool_use"187188189def test_anthropic_stream_normalised_with_input_json_delta():190 events = [191 {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 25, "output_tokens": 1}}},192 {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},193 {"type": "ping"},194 {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "OK"}},195 {"type": "content_block_stop", "index": 0},196 {"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}},197 {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": ""}},198 {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "{\"city\": \"Pa"}},199 {"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "ris\"}"}},200 {"type": "content_block_stop", "index": 1},201 {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 15}},202 {"type": "message_stop"},203 ]204 evs = list(provider("anthropic", Script([sse(events)])).stream(lp.GenerateRequest("m", [lp.Message("user", "x")])))205 types = [e.type for e in evs]206 assert types == ["raw", "raw", "raw", "text_delta", "raw", "tool_call_start", "tool_call_delta", "tool_call_delta", "tool_call_delta", "tool_call_done", "usage", "done"]207 tc = next(e for e in evs if e.type == "tool_call_done").tool_call208 assert tc.id == "toolu_1" and tc.arguments == {"city": "Paris"}209 usage = next(e for e in evs if e.type == "usage").usage210 assert usage.input_tokens == 25 and usage.output_tokens == 15211 assert evs[-1].stop_reason == "tool_use"212213214def test_anthropic_stream_error_event_surfaces():215 events = [{"type": "message_start", "message": {}}, {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}]216 evs = list(provider("anthropic", Script([sse(events)])).stream(lp.GenerateRequest("m", [lp.Message("user", "x")])))217 assert evs[-1].type == "error" and evs[-1].error["type"] == "overloaded_error"218219220# ---------------- count_tokens / web_search / upload bodies ----------------221222def test_count_tokens_bodies():223 s = Script([ok({"object": "response.input_tokens", "input_tokens": 42})])224 assert provider("openai", s).count_tokens(lp.GenerateRequest("m", [lp.Message("user", "x")], system="s")) == 42225 assert s.requests[0].url.endswith("/v1/responses/input_tokens") and "max_output_tokens" not in s.body()226 s2 = Script([ok({"input_tokens": 7})])227 assert provider("anthropic", s2).count_tokens(lp.GenerateRequest("m", [lp.Message("user", "x")])) == 7228 assert s2.requests[0].url.endswith("/v1/messages/count_tokens") and "max_tokens" not in s2.body()229230231def test_web_search_tool_injection():232 s = Script([ok({"status": "completed", "output": [], "usage": {}})])233 provider("openai", s).web_search(lp.GenerateRequest("m", [lp.Message("user", "x")]))234 assert s.body()["tools"] == [{"type": "web_search"}]235 s2 = Script([ok({"stop_reason": "end_turn", "content": [], "usage": {}})])236 provider("anthropic", s2).web_search(lp.GenerateRequest("m", [lp.Message("user", "x")]), max_uses=3)237 assert s2.body()["tools"] == [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}]238239240def test_upload_file_multipart(tmp_path):241 f = tmp_path / "a.txt"242 f.write_text("hello")243 s = Script([ok({"id": "file_1", "filename": "a.txt", "bytes": 5})])244 ref = provider("openai", s).upload_file(str(f))245 assert ref.id == "file_1" and s.requests[0].headers["Content-Type"].startswith("multipart/form-data") and b'name="purpose"' in s.requests[0].body246 s2 = Script([ok({"id": "file_2", "filename": "a.txt", "size_bytes": 5})])247 ref2 = provider("anthropic", s2).upload_file(str(f))248 assert ref2.bytes == 5 and s2.requests[0].headers["anthropic-beta"] == "files-api-2025-04-14"249250251# ======================================================================================252# xAI (Responses default + Chat Completions variant) — shapes from docs/xai/{responses,streaming,tool-loop,files}.md253# ======================================================================================254255def gemini_sse(chunks):256 """Gemini `?alt=sse` framing: `data:` lines only, no event names, no [DONE]."""257 out = []258 for c in chunks:259 out += [f"data: {json.dumps(c)}\r\n".encode(), b"\r\n"]260 return rc.Response(200, {}, lines=iter(out))261262263def data_only_sse(chunks, done=True):264 """OpenAI/xAI Chat Completions framing: `data: {...}` per chunk + `data: [DONE]`."""265 out = []266 for c in chunks:267 out += [f"data: {json.dumps(c)}\n".encode(), b"\n"]268 if done:269 out += [b"data: [DONE]\n", b"\n"]270 return rc.Response(200, {}, lines=iter(out))271272273def test_xai_responses_body_mapping_and_auth():274 s = Script([ok({"id": "r1", "status": "completed", "output": [], "usage": {}})])275 p = provider("xai", s)276 p.generate(lp.GenerateRequest("grok-4.3", [lp.Message("user", "hi")], system="s", max_tokens=16, tools=[TOOL], json_schema={"schema": {"type": "object"}},277 extensions={"previous_response_id": "resp_0", "reasoning": {"effort": "low"}, "store": False}))278 b = s.body()279 assert s.requests[0].url == "https://api.x.ai/v1/responses" and s.requests[0].headers["Authorization"] == "Bearer k"280 assert b["model"] == "grok-4.3" and b["instructions"] == "s" and b["max_output_tokens"] == 16 and b["previous_response_id"] == "resp_0"281 assert b["tools"][0]["type"] == "function" and b["text"]["format"]["type"] == "json_schema" and b["store"] is False282 assert "metadata" not in b and "background" not in b # xAI rejects both with 400 "Argument not supported"283284285def test_xai_responses_parse_reasoning_item_and_usage():286 body = {"id": "b778ce56", "object": "response", "model": "grok-4.3", "status": "completed",287 "output": [{"type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "The user requested a reply of \"OK.\""}], "status": "completed"},288 {"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": "OK.", "annotations": []}]}],289 "usage": {"input_tokens": 196, "input_tokens_details": {"cached_tokens": 192}, "output_tokens": 98, "output_tokens_details": {"reasoning_tokens": 96}, "total_tokens": 294}}290 r = provider("xai", Script([ok(body)])).generate(lp.GenerateRequest("grok-4.3", [lp.Message("user", "x")]))291 assert r.provider == "xai" and r.text == "OK." and r.stop_reason == "end" and r.reasoning.startswith("The user requested")292 assert r.usage.input_tokens == 196 and r.usage.cached_input_tokens == 192 and r.usage.output_tokens == 98 # output includes reasoning293294295def test_xai_responses_stream_reasoning_summary_and_single_delta_function_call():296 events = [297 {"type": "response.created", "sequence_number": 0, "response": {"id": "ee2a", "status": "in_progress"}},298 {"type": "response.output_item.added", "sequence_number": 2, "item": {"id": "rs_ee2a", "type": "reasoning", "summary": []}, "output_index": 0},299 {"type": "response.reasoning_summary_text.delta", "sequence_number": 4, "delta": "I need to calculate ", "item_id": "rs_ee2a", "summary_index": 0},300 {"type": "response.output_item.added", "sequence_number": 12, "item": {"id": "fc_ee2a_0", "type": "function_call", "call_id": "call-755d-0", "name": "get_weather", "arguments": ""}, "output_index": 1},301 {"type": "response.function_call_arguments.delta", "sequence_number": 13, "delta": "{\"city\":\"Paris\"}", "item_id": "fc_ee2a_0"}, # whole JSON in ONE delta (live)302 {"type": "response.function_call_arguments.done", "sequence_number": 14, "arguments": "{\"city\":\"Paris\"}", "item_id": "fc_ee2a_0", "name": "get_weather"},303 {"type": "response.completed", "sequence_number": 16, "response": {"id": "ee2a", "status": "completed", "output": [{"type": "function_call", "call_id": "call-755d-0", "name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}],304 "usage": {"input_tokens": 200, "output_tokens": 50, "total_tokens": 250}}},305 ]306 evs = list(provider("xai", Script([sse(events)])).stream(lp.GenerateRequest("grok-4.3", [lp.Message("user", "x")])))307 assert [e.type for e in evs] == ["raw", "raw", "reasoning_delta", "tool_call_start", "tool_call_delta", "tool_call_done", "usage", "done"]308 done = next(e for e in evs if e.type == "tool_call_done").tool_call309 assert done.name == "get_weather" and done.arguments == {"city": "Paris"} and evs[-1].stop_reason == "tool_use"310311312def test_xai_chat_body_tool_calls_round_trip_and_stop():313 first = {"id": "c1", "model": "grok-4.3", "choices": [{"index": 0, "finish_reason": "tool_calls", "message": {"role": "assistant", "content": None, "reasoning_content": "I'll use the tool",314 "tool_calls": [{"id": "call-755df5ce-0", "type": "function", "function": {"name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}}]}}],315 "usage": {"prompt_tokens": 196, "completion_tokens": 20, "total_tokens": 216, "prompt_tokens_details": {"cached_tokens": 192}, "completion_tokens_details": {"reasoning_tokens": 15}}}316 second = {"id": "c2", "model": "grok-4.3", "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Sunny in Paris."}}], "usage": {"prompt_tokens": 250, "completion_tokens": 5, "total_tokens": 255}}317 s = Script([ok(first), ok(second)])318 p = lp.get_provider("xai", rc.ResilientClient("xai", api_key="k", transport=s), api="chat")319 r = p.use_tools(lp.GenerateRequest("grok-4.3", [lp.Message("user", "weather?")], system="be brief", tools=[TOOL], stop=["END"]), {"get_weather": lambda city: f"{city}: 20C"})320 assert r.text == "Sunny in Paris." and r.stop_reason == "end" and s.requests[0].url.endswith("/v1/chat/completions")321 b0, b1 = s.body(0), s.body(1)322 assert b0["messages"][0] == {"role": "system", "content": "be brief"} and b0["tools"][0]["function"]["name"] == "get_weather" and b0["stop"] == ["END"] and b0["max_tokens"] == 1024323 assert b1["messages"][2]["tool_calls"][0]["function"]["arguments"] == "{\"city\":\"Paris\"}"324 assert b1["messages"][3] == {"role": "tool", "tool_call_id": "call-755df5ce-0", "content": "Paris: 20C"}325 first_result = p.parse_chat(first, 200, 1, None)326 assert first_result.stop_reason == "tool_use" and first_result.reasoning == "I'll use the tool" and first_result.usage.cached_input_tokens == 192327328329def test_xai_chat_stream_reasoning_content_whole_tool_call_and_usage_chunk():330 base = {"id": "29c58dae", "object": "chat.completion.chunk", "created": 0, "model": "grok-4.3", "system_fingerprint": "fp_x", "service_tier": "default"}331 chunks = [332 {**base, "choices": [{"index": 0, "delta": {"reasoning_content": "I'll use the get_weather", "role": "assistant"}}]},333 {**base, "choices": [{"index": 0, "delta": {"reasoning_content": " tool.", "tool_calls": [{"id": "call-755d-0", "function": {"name": "get_weather", "arguments": "{\"city\":\"Paris\"}"}, "index": 0, "type": "function"}]}}]},334 {**base, "choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]},335 {**base, "choices": [], "usage": {"prompt_tokens": 196, "completion_tokens": 2, "total_tokens": 319, "prompt_tokens_details": {"cached_tokens": 192}, "completion_tokens_details": {"reasoning_tokens": 121}}},336 ]337 p = lp.get_provider("xai", rc.ResilientClient("xai", api_key="k", transport=Script([data_only_sse(chunks)])), api="chat")338 evs = list(p.stream(lp.GenerateRequest("grok-4.3", [lp.Message("user", "x")])))339 assert [e.type for e in evs] == ["reasoning_delta", "reasoning_delta", "tool_call_start", "tool_call_delta", "tool_call_done", "usage", "done"]340 assert "".join(e.text for e in evs if e.type == "reasoning_delta") == "I'll use the get_weather tool."341 assert evs[4].tool_call.arguments == {"city": "Paris"} and evs[5].usage.cached_input_tokens == 192 and evs[6].stop_reason == "tool_use"342343344def test_xai_count_tokens_upload_and_web_search_filters(tmp_path):345 s = Script([ok({"token_ids": [{"token_id": 1, "string_token": "Reply", "token_bytes": [82]}, {"token_id": 2, "string_token": " with", "token_bytes": [32]},346 {"token_id": 3, "string_token": " OK", "token_bytes": [32]}, {"token_id": 4, "string_token": ".", "token_bytes": [46]}]})])347 assert provider("xai", s).count_tokens(lp.GenerateRequest("grok-4.3", [lp.Message("user", "Reply with OK.")])) == 4348 assert s.requests[0].url.endswith("/v1/tokenize-text") and s.body() == {"model": "grok-4.3", "text": "Reply with OK."}349 f = tmp_path / "a.txt"350 f.write_text("hello")351 s2 = Script([ok({"id": "file_8377", "object": "file", "bytes": 5, "filename": "a.txt", "purpose": ""})])352 ref = provider("xai", s2).upload_file(str(f))353 assert ref.provider == "xai" and ref.id == "file_8377" and ref.bytes == 5 and s2.requests[0].url.endswith("/v1/files")354 s3 = Script([ok({"status": "completed", "output": [{"type": "web_search_call", "id": "ws_1", "status": "completed", "action": {"type": "open_page", "url": "https://x.ai/news"}},355 {"type": "message", "content": [{"type": "output_text", "text": "News.", "annotations": [{"type": "url_citation", "url": "https://x.ai/news"}]}]}], "usage": {}})])356 r = provider("xai", s3).web_search(lp.GenerateRequest("grok-4.3", [lp.Message("user", "x")], extensions={"max_turns": 2}), allowed_domains=["x.ai"])357 assert s3.body()["tools"] == [{"type": "web_search", "allowed_domains": ["x.ai"]}] and s3.body()["max_turns"] == 2 and r.text == "News."358 with pytest.raises(ValueError):359 provider("xai", Script([])).web_search(lp.GenerateRequest("m", []), allowed_domains=["a.com"], excluded_domains=["b.com"])360361362def test_xai_bad_key_400_is_not_retried():363 s = Script([rc.Response(400, {}, json.dumps({"code": "invalid-argument", "error": "Incorrect API key provided. You can obtain an API key from https://console.x.ai."}).encode())])364 with pytest.raises(rc.NonRetryableError) as ei:365 provider("xai", s).generate(lp.GenerateRequest("grok-4.3", [lp.Message("user", "x")]))366 assert len(s.requests) == 1 and "incorrect API key" in str(ei.value)367368369# ======================================================================================370# Gemini — shapes from docs/gemini/{generate-content,streaming,tool-loop,files,token-counting}.md and tmp-live recordings371# ======================================================================================372373GEM_OK = {"candidates": [{"content": {"parts": [{"text": "OK.", "thoughtSignature": "El4K…"}], "role": "model"}, "finishReason": "STOP", "index": 0}],374 "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 5}], "serviceTier": "standard"},375 "modelVersion": "gemini-3.5-flash-lite", "responseId": "ZgWu"}376377378def test_gemini_body_mapping_and_auth_header():379 s = Script([ok(GEM_OK)])380 p = provider("gemini", s)381 r = p.generate(lp.GenerateRequest("gemini-3.5-flash-lite", [lp.Message("user", "hi")], system="be brief", max_tokens=16, temperature=0.2, tools=[TOOL],382 tool_choice={"name": "get_weather"}, json_schema={"schema": {"type": "object"}}, stop=["END"],383 extensions={"generationConfig": {"thinkingConfig": {"thinkingLevel": "LOW"}}, "safetySettings": [], "serviceTier": "flex"}))384 req0 = s.requests[0]385 assert req0.url == "https://generativelanguage.googleapis.com/v1beta/models/gemini-3.5-flash-lite:generateContent" and "key=" not in req0.url386 assert req0.headers["x-goog-api-key"] == "k" and req0.safe_headers()["x-goog-api-key"] == "***REDACTED***"387 b = s.body()388 assert b["contents"] == [{"role": "user", "parts": [{"text": "hi"}]}] and b["systemInstruction"] == {"parts": [{"text": "be brief"}]}389 assert b["generationConfig"] == {"maxOutputTokens": 16, "temperature": 0.2, "stopSequences": ["END"], "responseMimeType": "application/json",390 "responseJsonSchema": {"type": "object"}, "thinkingConfig": {"thinkingLevel": "LOW"}}391 assert b["tools"] == [{"functionDeclarations": [{"name": "get_weather", "description": "Weather by city", "parametersJsonSchema": TOOL.parameters}]}]392 assert b["toolConfig"] == {"functionCallingConfig": {"mode": "ANY", "allowedFunctionNames": ["get_weather"]}} and b["serviceTier"] == "flex"393 assert r.provider == "gemini" and r.text == "OK." and r.stop_reason == "end" and r.model == "gemini-3.5-flash-lite" and r.response_id == "ZgWu"394 assert r.usage.input_tokens == 5 and r.usage.output_tokens == 1 and r.assistant_content["parts"][0]["thoughtSignature"] == "El4K…"395396397def test_gemini_tool_choice_and_model_prefix():398 p = lp.GeminiProvider(rc.ResilientClient("gemini", api_key="k"))399 assert p._tool_config("auto") == {"functionCallingConfig": {"mode": "AUTO"}} and p._tool_config("required") == {"functionCallingConfig": {"mode": "ANY"}}400 assert p._tool_config("none") == {"functionCallingConfig": {"mode": "NONE"}}401 assert p._path("models/gemini-3.8-flash", "countTokens") == "/v1beta/models/gemini-3.8-flash:countTokens"402 assert lp.GeminiProvider(rc.ResilientClient("gemini", api_key="k"), api_version="v1")._path("g", "generateContent") == "/v1/models/g:generateContent"403404405def test_gemini_tool_loop_echoes_model_turn_verbatim_with_thought_signature():406 first = {"candidates": [{"content": {"role": "model", "parts": [407 {"functionCall": {"name": "get_weather", "args": {"city": "Paris"}, "id": "call_1"}, "thoughtSignature": "SIG1"},408 {"functionCall": {"name": "get_time", "args": {"city": "Tokyo"}, "id": "call_2"}}]}, "finishReason": "STOP", "index": 0}], # only the first parallel call carries a signature409 "usageMetadata": {"promptTokenCount": 40, "candidatesTokenCount": 12, "thoughtsTokenCount": 30, "totalTokenCount": 82}, "modelVersion": "gemini-3.5-flash-lite"}410 second = {"candidates": [{"content": {"role": "model", "parts": [{"text": "Rainy, 14:05.", "thoughtSignature": "SIG2"}]}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 90, "candidatesTokenCount": 6, "totalTokenCount": 96}}411 s = Script([ok(first), ok(second)])412 r = provider("gemini", s).use_tools(lp.GenerateRequest("gemini-3.5-flash-lite", [lp.Message("user", "weather+time?")], tools=[TOOL]),413 {"get_weather": lambda city: {"city": city, "temperature_c": 14}, "get_time": lambda city: "14:05"})414 assert r.text == "Rainy, 14:05." and r.stop_reason == "end"415 contents = s.body(1)["contents"]416 assert contents[1] == first["candidates"][0]["content"] # VERBATIM echo: thoughtSignature preserved, part order untouched417 assert contents[2]["role"] == "user" and [p_["functionResponse"]["name"] for p_ in contents[2]["parts"]] == ["get_weather", "get_time"] # one Content, one part per call418 assert contents[2]["parts"][0]["functionResponse"] == {"name": "get_weather", "response": {"city": "Paris", "temperature_c": 14}, "id": "call_1"}419 assert contents[2]["parts"][1]["functionResponse"]["response"] == {"result": "14:05"} # scalars wrapped into an object420 first_res = lp.GeminiProvider(rc.ResilientClient("gemini", api_key="k")).parse(first, 200, 1, None)421 assert first_res.stop_reason == "tool_use" and first_res.tool_calls[0].thought_signature == "SIG1" and first_res.tool_calls[1].thought_signature is None422 assert first_res.usage.output_tokens == 42 # candidates + thoughts (both billed)423424425def test_gemini_rebuilds_model_turn_from_tool_calls_when_no_raw_content():426 s = Script([ok(GEM_OK)])427 msgs = [lp.Message("user", "q"), lp.Message("assistant", "", tool_calls=[lp.ToolCall("call_9", "get_weather", {"city": "Rome"}, thought_signature="SIGX")]),428 lp.Message("tool", "{\"c\": 20}", tool_call_id="call_9", tool_name="get_weather"), lp.Message("tool", "boom", tool_call_id="call_10", tool_name="other", is_error=True)]429 provider("gemini", s).generate(lp.GenerateRequest("g", msgs))430 c = s.body()["contents"]431 assert c[1] == {"role": "model", "parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Rome"}, "id": "call_9"}, "thoughtSignature": "SIGX"}]}432 assert c[2]["parts"][0]["functionResponse"]["response"] == {"c": 20} and c[2]["parts"][1]["functionResponse"]["response"] == {"error": "boom"}433434435def test_gemini_stop_reasons_and_prompt_block():436 p = lp.GeminiProvider(rc.ResilientClient("gemini", api_key="k"))437 mk = lambda fr, parts=None: {"candidates": [{"content": {"parts": parts or [{"text": "x"}], "role": "model"}, "finishReason": fr, "index": 0}], "usageMetadata": {}} # noqa: E731438 assert p.parse(mk("MAX_TOKENS", [{"text": "", "thoughtSignature": "s"}]), 200, 1, None).stop_reason == "max_tokens"439 assert p.parse(mk("SAFETY"), 200, 1, None).stop_reason == "refusal" and p.parse(mk("RECITATION"), 200, 1, None).stop_reason == "refusal"440 assert p.parse(mk("MISSING_THOUGHT_SIGNATURE"), 200, 1, None).stop_reason == "other"441 blocked = p.parse({"promptFeedback": {"blockReason": "PROHIBITED_CONTENT"}, "usageMetadata": {"promptTokenCount": 9, "totalTokenCount": 9}}, 200, 1, None)442 assert blocked.stop_reason == "refusal" and blocked.error["blockReason"] == "PROHIBITED_CONTENT" and blocked.text == ""443 thought = p.parse(mk("STOP", [{"text": "**Plan**", "thought": True}, {"text": "42"}]), 200, 1, None)444 assert thought.text == "42" and thought.reasoning == "**Plan**"445446447def test_gemini_stream_alt_sse_text_thought_signature_and_usage():448 chunks = [449 {"candidates": [{"content": {"parts": [{"text": "**Thinking**", "thought": True}], "role": "model"}, "index": 0}], "usageMetadata": {"promptTokenCount": 13, "thoughtsTokenCount": 40, "totalTokenCount": 53}, "responseId": "bwWu"},450 {"candidates": [{"content": {"parts": [{"text": "1 "}], "role": "model"}, "index": 0}], "usageMetadata": {"promptTokenCount": 13, "candidatesTokenCount": 2, "thoughtsTokenCount": 40, "totalTokenCount": 55}, "responseId": "bwWu"},451 {"candidates": [{"content": {"parts": [{"text": "2 3"}], "role": "model"}, "index": 0}], "usageMetadata": {"promptTokenCount": 13, "candidatesTokenCount": 4, "thoughtsTokenCount": 40, "totalTokenCount": 57}, "responseId": "bwWu"},452 {"candidates": [{"content": {"parts": [{"text": "", "thoughtSignature": "El4KXAFp…"}], "role": "model"}, "finishReason": "STOP", "index": 0}],453 "usageMetadata": {"promptTokenCount": 13, "candidatesTokenCount": 4, "thoughtsTokenCount": 40, "totalTokenCount": 57, "serviceTier": "standard"}, "responseId": "bwWu"},454 ]455 evs = list(provider("gemini", Script([gemini_sse(chunks)])).stream(lp.GenerateRequest("gemini-3.5-flash-lite", [lp.Message("user", "count")])))456 assert [e.type for e in evs] == ["reasoning_delta", "text_delta", "text_delta", "raw", "usage", "done"]457 assert "".join(e.text for e in evs if e.type == "text_delta") == "1 2 3" and evs[0].text == "**Thinking**"458 assert evs[-2].usage.input_tokens == 13 and evs[-2].usage.output_tokens == 44 and evs[-1].stop_reason == "end" # last chunk's usage is authoritative459460461def test_gemini_stream_function_call_and_truncated_stream():462 chunks = [{"candidates": [{"content": {"parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}, "id": "call_7"}, "thoughtSignature": "SIG"}], "role": "model"}, "finishReason": "STOP", "index": 0}],463 "usageMetadata": {"promptTokenCount": 30, "candidatesTokenCount": 10, "totalTokenCount": 40}}]464 evs = list(provider("gemini", Script([gemini_sse(chunks)])).stream(lp.GenerateRequest("g", [lp.Message("user", "x")])))465 assert [e.type for e in evs] == ["tool_call_start", "tool_call_delta", "tool_call_done", "usage", "done"]466 assert evs[2].tool_call.id == "call_7" and evs[2].tool_call.arguments == {"city": "Paris"} and evs[2].tool_call.thought_signature == "SIG" and evs[-1].stop_reason == "tool_use"467 cut = list(provider("gemini", Script([gemini_sse([{"candidates": [{"content": {"parts": [{"text": "par"}], "role": "model"}, "index": 0}]}])])).stream(lp.GenerateRequest("g", [lp.Message("user", "x")])))468 assert cut[-1].type == "done" and cut[-1].stop_reason == "incomplete" # connection closed without finishReason469470471def test_gemini_count_tokens_wrapper_form():472 s = Script([ok({"totalTokens": 51, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 51}]})])473 n = provider("gemini", s).count_tokens(lp.GenerateRequest("gemini-3.5-flash-lite", [lp.Message("user", "The quick brown fox")], system="sys", tools=[TOOL]))474 assert n == 51 and s.requests[0].url.endswith("/v1beta/models/gemini-3.5-flash-lite:countTokens")475 b = s.body()476 assert set(b) == {"generateContentRequest"} and b["generateContentRequest"]["model"] == "models/gemini-3.5-flash-lite"477 assert b["generateContentRequest"]["systemInstruction"] == {"parts": [{"text": "sys"}]} and "tools" in b["generateContentRequest"] and "generationConfig" not in b["generateContentRequest"]478479480def test_gemini_structured_output_and_truncation():481 good = {"candidates": [{"content": {"parts": [{"text": "{\"ok\": true}", "thoughtSignature": "s"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {}}482 assert provider("gemini", Script([ok(good)])).structured_output(lp.GenerateRequest("g", [lp.Message("user", "x")]), {"type": "object"}) == {"ok": True}483 trunc = {"candidates": [{"content": {"parts": [{"text": "{\"ok\": true, \"word\": \"OK"}], "role": "model"}, "finishReason": "MAX_TOKENS"}], "usageMetadata": {}}484 with pytest.raises(ValueError):485 provider("gemini", Script([ok(trunc)])).structured_output(lp.GenerateRequest("g", [lp.Message("user", "x")]), {"type": "object"})486487488def test_gemini_resumable_upload_two_steps(tmp_path):489 f = tmp_path / "probe.txt"490 f.write_text("hello world")491 s = Script([rc.Response(200, {"x-goog-upload-url": "https://generativelanguage.googleapis.com/upload/v1beta/files?upload_id=abc&upload_protocol=resumable", "x-goog-upload-status": "active"}, b""),492 ok({"file": {"name": "files/w7x6", "displayName": "probe.txt", "mimeType": "text/plain", "sizeBytes": "11", "uri": "https://generativelanguage.googleapis.com/v1beta/files/w7x6", "state": "ACTIVE"}})])493 ref = provider("gemini", s).upload_file(str(f))494 start, fin = s.requests495 assert start.url.endswith("/upload/v1beta/files") and start.headers["X-Goog-Upload-Protocol"] == "resumable" and start.headers["X-Goog-Upload-Command"] == "start"496 assert start.headers["X-Goog-Upload-Header-Content-Length"] == "11" and start.headers["X-Goog-Upload-Header-Content-Type"] == "text/plain" and json.loads(start.body) == {"file": {"display_name": "probe.txt"}}497 assert fin.url.endswith("/upload/v1beta/files?upload_id=abc&upload_protocol=resumable") and fin.headers["X-Goog-Upload-Command"] == "upload, finalize" and fin.headers["X-Goog-Upload-Offset"] == "0" and fin.body == b"hello world"498 assert ref.provider == "gemini" and ref.id == "https://generativelanguage.googleapis.com/v1beta/files/w7x6" and ref.bytes == 11499 s2 = Script([ok(GEM_OK)])500 provider("gemini", s2).generate(lp.GenerateRequest("g", [lp.Message("user", "secret?", file_ids=[ref.id])]))501 assert s2.body()["contents"][0]["parts"][1] == {"fileData": {"fileUri": ref.id}}502503504GEMINI_429_LIMIT_0 = {"error": {"code": 429, "message": "You exceeded your current quota, please check your plan and billing details. ... * Quota exceeded for metric: generativelanguage.googleapis.com/generate_content_free_tier_requests, limit: 0, model: gemini-3.1-pro ... Please retry in 50.868302469s.",505 "status": "RESOURCE_EXHAUSTED", "details": [{"@type": "type.googleapis.com/google.rpc.Help"}, {"@type": "type.googleapis.com/google.rpc.QuotaFailure", "violations": [{"quotaMetric": "generativelanguage.googleapis.com/generate_content_free_tier_requests", "quotaId": "GenerateRequestsPerDayPerProjectPerModel-FreeTier", "quotaDimensions": {"model": "gemini-3.1-pro", "location": "global"}}]}]}}506507508def test_gemini_web_search_restricted_returns_typed_result_not_exception():509 s = Script([rc.Response(429, {"x-gemini-service-tier": "standard"}, json.dumps(GEMINI_429_LIMIT_0).encode())])510 r = provider("gemini", s).web_search(lp.GenerateRequest("gemini-3.5-flash-lite", [lp.Message("user", "today?")]))511 assert s.body()["tools"] == [{"googleSearch": {}}]512 assert r.stop_reason == "restricted" and r.http_status == 429 and r.error["status"] == "RESOURCE_EXHAUSTED" and r.text == "" and len(s.requests) == 1513 s2 = Script([rc.Response(400, {}, json.dumps({"error": {"code": 400, "message": "Precondition check failed.", "status": "FAILED_PRECONDITION"}}).encode())])514 assert provider("gemini", s2).web_search(lp.GenerateRequest("g", [lp.Message("user", "x")])).stop_reason == "restricted"515 s3 = Script([rc.Response(400, {}, json.dumps({"error": {"code": 400, "message": "Invalid JSON payload received. Unknown name \"x\"", "status": "INVALID_ARGUMENT"}}).encode())])516 with pytest.raises(rc.NonRetryableError): # genuine request bugs still raise517 provider("gemini", s3).web_search(lp.GenerateRequest("g", [lp.Message("user", "x")]))518519520def test_get_provider_four_names():521 for name in ("openai", "anthropic", "xai", "gemini"):522 assert lp.get_provider(name, rc.ResilientClient(name, api_key="k")).name == name523