142 lines
4.9 KiB
Python
142 lines
4.9 KiB
Python
"""Record one scenario through a locally constructed application."""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from importlib import import_module
|
|
from pathlib import Path
|
|
from typing import cast
|
|
|
|
import httpx
|
|
import httpx2
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
from live import ScenarioResult, build_request
|
|
from recording import Httpx2RecordingTransport, HttpxRecordingTransport, write_cassettes
|
|
from scenario import SCENARIO_PATH, Scenario, load_scenarios
|
|
from wire import BACKEND_ROOT, validate_event
|
|
|
|
EVAL_ROOT = Path(__file__).parent
|
|
|
|
|
|
def main() -> int:
|
|
"""Record the selected scenario and persist raw plus safe cassettes."""
|
|
arguments = _parse_arguments()
|
|
scenarios = load_scenarios(arguments.scenarios)
|
|
selected = next((item for item in scenarios if item.key == arguments.scenario), None)
|
|
if selected is None:
|
|
raise SystemExit(f"Unknown scenario: {arguments.scenario}")
|
|
spotify_client_id = _required_environment("SPOTIFY_CLIENT_ID")
|
|
spotify_refresh_token = _required_environment("SPOTIFY_SEED_REFRESH_TOKEN")
|
|
anthropic_api_key = _required_environment("ANTHROPIC_API_KEY")
|
|
|
|
spotify_transport = Httpx2RecordingTransport(httpx2.AsyncHTTPTransport())
|
|
anthropic_transport = HttpxRecordingTransport(httpx.AsyncHTTPTransport())
|
|
anthropic_http_client = httpx.AsyncClient(transport=anthropic_transport)
|
|
app = _create_recording_app(
|
|
spotify_client_id,
|
|
spotify_refresh_token,
|
|
anthropic_api_key,
|
|
spotify_transport,
|
|
anthropic_http_client,
|
|
)
|
|
try:
|
|
with TestClient(app, base_url=arguments.base_url) as client:
|
|
_record_selected(client, scenarios, selected, spotify_transport)
|
|
finally:
|
|
raw_root, redacted_root = write_cassettes(
|
|
arguments.fixture_dir,
|
|
selected.key,
|
|
{
|
|
"spotify": spotify_transport.interactions,
|
|
"anthropic": anthropic_transport.interactions,
|
|
},
|
|
)
|
|
print(f"Raw cassettes: {raw_root}")
|
|
print(f"Redacted cassettes: {redacted_root}")
|
|
return 0
|
|
|
|
|
|
def _record_selected(
|
|
client: TestClient,
|
|
scenarios: tuple[Scenario, ...],
|
|
selected: Scenario,
|
|
spotify_transport: Httpx2RecordingTransport,
|
|
) -> None:
|
|
if selected.after is not None:
|
|
parent_scenario = next(item for item in scenarios if item.key == selected.after)
|
|
parent = _send_turn(client, parent_scenario, None)
|
|
search_count = _search_count(spotify_transport)
|
|
_send_turn(client, selected, parent)
|
|
if _search_count(spotify_transport) != search_count:
|
|
raise RuntimeError("Refinement issued new Spotify searches")
|
|
else:
|
|
_send_turn(client, selected, None)
|
|
|
|
|
|
def _send_turn(
|
|
client: TestClient,
|
|
scenario: Scenario,
|
|
parent: ScenarioResult | None,
|
|
) -> ScenarioResult:
|
|
response = client.post("/api/recommendations", json=build_request(scenario, parent))
|
|
response.raise_for_status()
|
|
result = ScenarioResult(scenario)
|
|
for line in response.text.splitlines():
|
|
event = validate_event(line)
|
|
result.events.append(cast(dict[str, object], event.model_dump(mode="json")))
|
|
if not result.events or result.events[-1].get("type") != "done":
|
|
raise RuntimeError(f"Scenario {scenario.key} did not complete")
|
|
return result
|
|
|
|
|
|
def _create_recording_app(
|
|
spotify_client_id: str,
|
|
spotify_refresh_token: str,
|
|
anthropic_api_key: str,
|
|
spotify_transport: Httpx2RecordingTransport,
|
|
anthropic_http_client: httpx.AsyncClient,
|
|
) -> FastAPI:
|
|
backend_path = str(BACKEND_ROOT)
|
|
if backend_path not in sys.path:
|
|
sys.path.insert(0, backend_path)
|
|
config_module = import_module("app.config")
|
|
main_module = import_module("app.main")
|
|
application_settings = config_module.Settings(
|
|
app_mode=config_module.AppMode.LIVE,
|
|
spotify_client_id=spotify_client_id,
|
|
spotify_seed_refresh_token=spotify_refresh_token,
|
|
anthropic_api_key=anthropic_api_key,
|
|
)
|
|
return cast(
|
|
FastAPI,
|
|
main_module.create_app(
|
|
application_settings=application_settings,
|
|
http_transport=spotify_transport,
|
|
anthropic_http_client=anthropic_http_client,
|
|
),
|
|
)
|
|
|
|
|
|
def _search_count(transport: Httpx2RecordingTransport) -> int:
|
|
return sum("/search" in interaction.url for interaction in transport.interactions)
|
|
|
|
|
|
def _required_environment(name: str) -> str:
|
|
value = os.environ.get(name)
|
|
if not value:
|
|
raise SystemExit(f"{name} is required for recording")
|
|
return value
|
|
|
|
|
|
def _parse_arguments() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--scenario", required=True)
|
|
parser.add_argument("--base-url", required=True)
|
|
parser.add_argument("--scenarios", type=Path, default=SCENARIO_PATH)
|
|
parser.add_argument("--fixture-dir", type=Path, default=EVAL_ROOT / "fixtures")
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|