SPB Git forge

spb/doc-api

Public
2commits 1branches 0releases
15.7 MBsize
maindefault branch
13 days agolast push
Python 88.3% TypeScript 7.6% Shell 4.1%
5.3 KB · 92 lines python
Raw Blame History
1"""Smoke tests for the Gemini Live API (BidiGenerateContent WebSocket) and ephemeral tokens (API Atlas).23Cheap (always run, free): POST /v1alpha/auth_tokens (and v1beta) → {name}; the token is never printed.4Realtime (RUN_REALTIME_TESTS=true, ≈ $0.003 per turn on gemini-2.5-flash-native-audio-latest): TEXT modality rejected (close 1007),5AUDIO turn message sequence (setupComplete → serverContent… → turnComplete + usageMetadata), toolCall with a NON_BLOCKING declaration,6`toolConfig` rejected in setup, ephemeral token accepted only on the BidiGenerateContentConstrained method (plain method → close 1008).7Reference run: 2026-09-18 (tmp-live/gemini-tools/j*.json).8"""9from __future__ import annotations1011import asyncio12import json13import os14from datetime import datetime, timedelta, timezone1516import pytest1718from scripts.live import log_request1920MODEL = "gemini-2.5-flash-native-audio-latest"21WS = "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.{ver}.GenerativeService.{method}"222324def _token_body(**extra):25    now = datetime.now(timezone.utc)26    return {"uses": 1, "expireTime": (now + timedelta(minutes=5)).strftime("%Y-%m-%dT%H:%M:%SZ"),27            "newSessionExpireTime": (now + timedelta(minutes=2)).strftime("%Y-%m-%dT%H:%M:%SZ"), **extra}282930def test_auth_token_create(gemini):31    for ver in ("v1alpha", "v1beta"):32        st, r, _ = gemini("POST", f"/{ver}/auth_tokens", _token_body(), note=f"test_live auth_tokens {ver}")33        assert st == 200 and list(r) == ["name"] and r["name"].startswith("auth_tokens/")343536async def _turn(url, headers, setup_extra, prompt="Reply with OK.", modality="AUDIO"):37    import websockets38    seq, msgs = [], []39    try:40        async with websockets.connect(url, additional_headers=headers, max_size=2**24) as ws:41            await ws.send(json.dumps({"setup": {"model": f"models/{MODEL}", "generationConfig": {"responseModalities": [modality]}, **setup_extra}}))42            first = json.loads(await asyncio.wait_for(ws.recv(), 30)); msgs.append(first); seq.append(",".join(first))43            await ws.send(json.dumps({"clientContent": {"turns": [{"role": "user", "parts": [{"text": prompt}]}], "turnComplete": True}}))44            for _ in range(80):45                m = json.loads(await asyncio.wait_for(ws.recv(), 30)); msgs.append(m); seq.append(",".join(m))46                if m.get("serverContent", {}).get("turnComplete"):47                    break48        return {"ok": True, "seq": seq, "msgs": msgs}49    except websockets.ConnectionClosed as e:50        return {"ok": False, "close_code": e.rcvd.code if e.rcvd else None, "reason": (e.rcvd.reason if e.rcvd else ""), "seq": seq, "msgs": msgs}51    finally:52        log_request("gemini", "WSS", url.split("?")[0].replace("wss://generativelanguage.googleapis.com", ""), 101, 0.003, "test_live turn")535455@pytest.mark.run_realtime_tests56def test_text_modality_rejected_and_audio_turn(gemini):57    hdr = {"x-goog-api-key": os.environ["GEMINI_API_KEY"]}58    url = WS.format(ver="v1beta", method="BidiGenerateContent")59    r = asyncio.run(_turn(url, hdr, {}, modality="TEXT"))60    assert not r["ok"] and r["close_code"] == 1007 and "response modalities" in r["reason"]61    r = asyncio.run(_turn(url, hdr, {"outputAudioTranscription": {}}))62    assert r["ok"] and r["seq"][0] == "setupComplete"63    last = r["msgs"][-1]64    assert last["serverContent"]["turnComplete"] is True and "usageMetadata" in last65    assert any("outputTranscription" in m.get("serverContent", {}) for m in r["msgs"])66    assert any(d["modality"] == "AUDIO" for d in last["usageMetadata"].get("responseTokensDetails", []))676869@pytest.mark.run_realtime_tests70def test_tool_call_and_tool_config_rejected(gemini):71    hdr = {"x-goog-api-key": os.environ["GEMINI_API_KEY"]}72    url = WS.format(ver="v1beta", method="BidiGenerateContent")73    tools = [{"functionDeclarations": [{"name": "get_weather", "description": "Weather for a city.", "behavior": "NON_BLOCKING",74                                        "parameters": {"type": "OBJECT", "properties": {"city": {"type": "STRING"}}, "required": ["city"]}}]}]75    r = asyncio.run(_turn(url, hdr, {"tools": tools, "toolConfig": {"functionCallingConfig": {"mode": "ANY"}}}))76    assert not r["ok"] and r["close_code"] == 1007 and "toolConfig" in r["reason"]77    r = asyncio.run(_turn(url, hdr, {"tools": tools, "outputAudioTranscription": {}}, prompt="Call get_weather for Paris now."))78    calls = [m["toolCall"]["functionCalls"] for m in r["msgs"] if "toolCall" in m]79    assert calls and calls[0][0]["name"] == "get_weather" and calls[0][0]["id"]808182@pytest.mark.run_realtime_tests83def test_ephemeral_token_constrained_method(gemini):84    st, tok, _ = gemini("POST", "/v1alpha/auth_tokens", _token_body(bidiGenerateContentSetup={"model": f"models/{MODEL}"}), note="test_live token for ws")85    assert st == 20086    hdr = {"Authorization": f"Token {tok['name']}"}87    r = asyncio.run(_turn(WS.format(ver="v1alpha", method="BidiGenerateContent"), hdr, {}))88    assert not r["ok"] and r["close_code"] == 100889    st, tok, _ = gemini("POST", "/v1alpha/auth_tokens", _token_body(), note="test_live token for constrained ws")90    r = asyncio.run(_turn(WS.format(ver="v1alpha", method="BidiGenerateContentConstrained"), {"Authorization": f"Token {tok['name']}"}, {"outputAudioTranscription": {}}))91    assert r["ok"] and r["msgs"][-1]["serverContent"]["turnComplete"]92