Call a REST API from an OpenAPI spec, with authentication

"""Replay the recorded Gemini turns against the real local release API, offline.

Run: python -m unittest -v test_assistant.py
"""
import asyncio
import json
import logging
import os
import secrets
import threading
import time
import unittest
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)
os.environ["SHIP_API_KEY"] = secrets.token_urlsafe(24)

import httpx  # noqa: E402
import uvicorn  # noqa: E402
from google.adk import Runner  # noqa: E402
from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.adk.sessions import InMemorySessionService  # noqa: E402
from google.genai import types  # noqa: E402

import assistant  # noqa: E402
import ship_api  # noqa: E402

RECORD = json.loads(Path(__file__).with_name("recorded-runs.json").read_text(encoding="utf-8"))
URL = "http://127.0.0.1:8767"
KEY = os.environ["SHIP_API_KEY"]


class ReplayLlm(BaseLlm):
    turns: list
    requests: list = []

    async def generate_content_async(self, llm_request, stream=False):
        self.requests.append(llm_request.model_dump_json(exclude_none=True))
        yield LlmResponse(content=types.Content.model_validate(self.turns.pop(0)))


def replay(tools, steps: list):
    model = ReplayLlm(model="replay", turns=[t for s in steps for t in s["model_turns"]], requests=[])
    runner = Runner(app_name="docs", agent=assistant.make_agent(tools).clone(update={"model": model}),
                    session_service=InMemorySessionService(), auto_create_session=True)

    async def go():
        events = []
        for step in steps:
            events.append(await assistant.ask(runner, "s1", step["question"]))
        await runner.close()
        return events
    return asyncio.run(go()), model


class ReleaseApi(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.server = uvicorn.Server(uvicorn.Config(ship_api.app, host="127.0.0.1", port=8767, log_level="warning"))
        threading.Thread(target=cls.server.run, daemon=True).start()
        while not cls.server.started:
            time.sleep(0.05)
        cls.spec = assistant.load_spec(URL)

    @classmethod
    def tearDownClass(cls):
        cls.server.should_exit = True

    def setUp(self):
        ship_api.REQUESTS.clear()

    def tools(self, **kwargs):
        return asyncio.run(assistant.release_tools(self.spec, KEY, **kwargs).get_tools())

    def test_operations_become_snake_case_tools(self):
        self.assertEqual([t.name for t in self.tools()], ["list_releases", "get_release"])
        self.assertEqual([t.name for t in self.tools(tool_filter=None)],
                         ["list_releases", "get_release", "rollback_environment"])

    def test_operation_ids_in_the_filter_match_nothing(self):
        self.assertEqual(self.tools(tool_filter=["listReleases", "getRelease"]), [])
        self.assertEqual(RECORD["failures"]["camel_case_filter"], [])

    def test_the_key_is_a_header_and_absent_from_recorded_model_requests(self):
        events, model = replay(assistant.release_tools(self.spec, KEY), RECORD["run"]["steps"])
        self.assertEqual([h["key_sent"] for h in ship_api.REQUESTS], [True, True])
        self.assertTrue(all(KEY not in r for r in model.requests))
        result = events[0][-2].get_function_responses()[0].response
        self.assertEqual(result, RECORD["run"]["steps"][0]["results"][0])

    def test_each_call_carries_the_invocation_id(self):
        events, _ = replay(assistant.release_tools(self.spec, KEY), RECORD["run"]["steps"][:1])
        self.assertEqual(ship_api.REQUESTS[0]["request_id"], events[0][0].invocation_id)

    def test_a_wrong_key_comes_back_to_the_model_as_a_tool_error(self):
        steps = RECORD["failures"]["wrong_key"]["steps"]
        events, _ = replay(assistant.release_tools(self.spec, "not-the-key"), steps)
        error = [r.response for e in events[0] for r in e.get_function_responses()][0]["error"]
        self.assertIn("Status Code: 401", error)

    def test_a_spec_without_servers_fails_the_run(self):
        spec = {k: v for k, v in self.spec.items() if k != "servers"}
        with self.assertRaisesRegex(httpx.UnsupportedProtocol, "missing an 'http://' or 'https://'"):
            replay(assistant.release_tools(spec, KEY), RECORD["failures"]["no_servers"]["steps"])
        self.assertEqual(ship_api.REQUESTS, [])

    def test_the_model_was_never_offered_the_rollback_operation(self):
        names = [d["name"] for d in RECORD["run"]["declarations"]]
        self.assertEqual(names, ["list_releases", "get_release"])
        self.assertEqual(RECORD["run"]["steps"][2]["calls"], [])


if __name__ == "__main__":
    unittest.main()

"""Record the docs assistant calling the local release API through OpenAPIToolset, on Gemini.

Run from this folder with GOOGLE_API_KEY set: python record.py --record
Starts ship_api on 127.0.0.1:8766 with a random API key, which is never written to disk.
Writes recorded-runs.json and run-output.txt. The tests replay the model turns offline.
"""
import asyncio
import json
import os
import secrets
import threading
import time
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")
os.environ["SHIP_API_KEY"] = secrets.token_urlsafe(24)

import uvicorn  # noqa: E402
from google.adk import Runner  # noqa: E402
from google.adk.sessions import InMemorySessionService  # noqa: E402

import assistant  # noqa: E402
import ship_api  # noqa: E402

HERE = Path(__file__).resolve().parent
QUESTIONS = [
    "Which staging releases failed, and why?",
    "What is live in production right now?",
    "Roll production back to r39.",
]


def serve() -> uvicorn.Server:
    server = uvicorn.Server(uvicorn.Config(ship_api.app, host="127.0.0.1", port=8766, log_level="warning"))
    threading.Thread(target=server.run, daemon=True).start()
    while not server.started:
        time.sleep(0.05)
    return server


def recording(agent, turns: list, declarations: list):
    def before(callback_context, llm_request):
        if not declarations:
            declarations.extend(f.model_dump(mode="json", exclude_none=True)
                                for t in llm_request.config.tools for f in t.function_declarations)

    def after(callback_context, llm_response):
        if llm_response.content:
            content = llm_response.content.model_dump(mode="json", exclude_none=True)
            for part in content.get("parts", []):
                part.pop("thought_signature", None)
            turns.append(content)
    return agent.clone(update={"before_model_callback": before, "after_model_callback": after})


async def conversation(tools, questions: list) -> dict:
    turns, declarations = [], []
    runner = Runner(app_name="docs", agent=recording(assistant.make_agent(tools), turns, declarations),
                    session_service=InMemorySessionService(), auto_create_session=True)
    steps = []
    for question in questions:
        start_turns, start_requests = len(turns), len(ship_api.REQUESTS)
        error = None
        try:
            events = await assistant.ask(runner, "s1", question)
        except Exception as exc:  # recorded as a failure, not hidden
            events, error = [], f"{type(exc).__name__}: {exc}"
        answer = next(("".join(p.text or "" for p in e.content.parts if not p.thought).strip()
                       for e in reversed(events) if e.is_final_response() and e.content), "")
        steps.append({
            "question": question,
            "calls": [{"name": c.name, "args": c.args} for e in events for c in e.get_function_calls()],
            "results": [r.response for e in events for r in e.get_function_responses()],
            "http": ship_api.REQUESTS[start_requests:], "model_turns": turns[start_turns:],
            "answer": answer, "error": error,
        })
    await runner.close()
    return {"tools": [t.name for t in await tools.get_tools()], "declarations": declarations, "steps": steps}


async def main() -> None:
    server = serve()
    spec = assistant.load_spec()
    key = os.environ["SHIP_API_KEY"]
    main_run = await conversation(assistant.release_tools(spec, key), QUESTIONS)
    wrong_key = await conversation(assistant.release_tools(spec, "not-the-key"), QUESTIONS[:1])
    camel = assistant.release_tools(spec, key, tool_filter=["listReleases", "getRelease"])
    no_servers = await conversation(assistant.release_tools({k: v for k, v in spec.items() if k != "servers"}, key),
                                    QUESTIONS[:1])
    server.should_exit = True
    record = {
        "model": assistant.MODEL,
        "operations": {path: {m: op["operationId"] for m, op in ops.items()} for path, ops in spec["paths"].items()},
        "security_schemes": spec["components"]["securitySchemes"],
        "run": main_run,
        "failures": {"wrong_key": wrong_key, "camel_case_filter": [t.name for t in await camel.get_tools()],
                     "no_servers": no_servers},
    }
    text = json.dumps(record, indent=2)
    assert key not in text, "the API key must never be recorded"
    (HERE / "recorded-runs.json").write_text(text + "\n", encoding="utf-8")
    lines = [f"tools: {main_run['tools']}", ""]
    for label, run in [("main", main_run), ("wrong key", wrong_key), ("spec without servers", no_servers)]:
        for s in run["steps"]:
            lines.append(f"# {label}\n> {s['question']}")
            lines += [f"call   {c['name']} {json.dumps(c['args'])}" for c in s["calls"]]
            lines += [f"http   {h['method']} {h['path']} key_sent={h['key_sent']} request_id={h['request_id'][:10]}..."
                      for h in s["http"]]
            lines += [f"result {json.dumps(r)[:160]}" for r in s["results"]]
            lines.append(f"error  {s['error']}" if s["error"] else "answer " + s["answer"].replace("\n", " ")[:220])
            lines.append("")
    lines.append(f"# tool_filter with operationIds: {record['failures']['camel_case_filter']}")
    (HERE / "run-output.txt").write_text("\n".join(lines) + "\n", encoding="utf-8")
    print("\n".join(lines))


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Optional live-model recording; overwrites saved run files.")
    parser.add_argument("--record", action="store_true", help="explicitly allow live model calls and record replacement")
    args = parser.parse_args()
    if not args.record:
        parser.error("recording is opt-in: use --record; normal lesson commands do not need this script")
    asyncio.run(main())