diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 368a059..62acb56 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -35,14 +35,12 @@ make test # Unit tests with coverage report uv run pytest tests/unit --cov=lago_agent_sdk --cov-report=term-missing - -# Integration tests (require credentials — see env vars in each test) -AWS_BEARER_TOKEN_BEDROCK="..." \ -MISTRAL_API_KEY="..." \ -LAGO_API_URL="..." LAGO_API_KEY="..." LAGO_EXTERNAL_SUBSCRIPTION_ID="..." \ -uv run pytest tests/integration -q ``` +There is no committed live-provider test tier. Adapter behaviour is pinned by +captured real responses under `tests/unit/adapters/fixtures/`, which is what the +unit tests assert against; re-capture a fixture rather than hand-editing one. + ## Linting and type checks ```bash @@ -72,7 +70,6 @@ uv lock --upgrade-package X # bump a single package - `src/lago_agent_sdk/lago_client.py` — thin HTTP client to `/events/batch` - `tests/unit/` — unit tests, organized to mirror `src/` - `tests/unit/adapters/fixtures/` — captured real provider responses, used by adapter tests -- `tests/integration/` — live tests, gated on credential env vars ## Adding a provider @@ -82,7 +79,6 @@ uv lock --upgrade-package X # bump a single package 4. Update `detector.py` to recognize the client class. 5. Update `sdk.py::wrap()` to dispatch to the new wrapper. 6. Add unit tests against the captured fixtures. -7. Add a live integration test gated on the provider's API key env var. ## Pull request checklist diff --git a/README.md b/README.md index 2389c66..104167c 100644 --- a/README.md +++ b/README.md @@ -272,17 +272,6 @@ pip install -e '.[dev]' pytest ``` -Run live integration tests (requires real credentials): - -```bash -AWS_BEARER_TOKEN_BEDROCK="..." \ -MISTRAL_API_KEY="..." \ -LAGO_API_URL="https://api.getlago.com/api/v1/" \ -LAGO_API_KEY="..." \ -LAGO_EXTERNAL_SUBSCRIPTION_ID="sub_..." \ -pytest tests/integration -``` - ## Security Found a vulnerability? See [SECURITY.md](SECURITY.md). diff --git a/pyproject.toml b/pyproject.toml index 3daae9d..b170c14 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ dev = [ "ruff>=0.6", "mypy>=1.10", "types-requests>=2.31", - # every provider SDK (so unit + live integration tests can import them) + # every provider SDK (so the unit + wrapper tests can import them) "boto3>=1.34", "mistralai>=2.0", "anthropic>=0.30", diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/integration/test_lago_reconciliation.py b/tests/integration/test_lago_reconciliation.py deleted file mode 100644 index d6e64fc..0000000 --- a/tests/integration/test_lago_reconciliation.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Live Lago reconciliation — emit N events, poll current_usage, verify exact match. - -This is the ONLY test that proves Lago *accepts* what the SDK emits. Every other -integration test points at an in-process mock, so a wrong metric code, a missing -`dynamic` charge model, or a rejected `precise_total_amount_cents` would pass -there and only surface in production. - -Skipped unless LAGO_API_URL, LAGO_API_KEY, and LAGO_EXTERNAL_SUBSCRIPTION_ID are -set. For a local dev Lago behind a self-signed cert (Traefik's default), set -LAGO_VERIFY_SSL=false — the SDK has `LagoConfig.verify_ssl` for exactly that, and -this test honours the same switch on its own reads. `truststore` is an -alternative if the cert is in the OS trust store. -""" - -from __future__ import annotations - -import os -import time - -import pytest -import requests - -from lago_agent_sdk import CanonicalUsage, LagoSDK - -try: - import truststore - - truststore.inject_into_ssl() -except Exception: # noqa: BLE001 - pass - -API_URL = (os.environ.get("LAGO_API_URL") or "").rstrip("/") -API_KEY = os.environ.get("LAGO_API_KEY") or "" -SUB_ID = os.environ.get("LAGO_EXTERNAL_SUBSCRIPTION_ID") or "" -CUST_ID = os.environ.get("LAGO_EXTERNAL_CUSTOMER_ID") or "cust_demo" -# Mirrors LagoConfig.verify_ssl: a local dev instance on a self-signed cert is a -# real, common setup, and without this BOTH halves of this test fail on SSL — the -# SDK's POST and this module's own GET. -VERIFY_SSL = (os.environ.get("LAGO_VERIFY_SSL") or "true").strip().lower() not in ( - "0", - "false", - "no", -) - -pytestmark = pytest.mark.skipif( - not (API_URL and API_KEY and SUB_ID), - reason="LAGO_API_URL / LAGO_API_KEY / LAGO_EXTERNAL_SUBSCRIPTION_ID not set", -) - - -def _read_usage() -> dict[str, float]: - r = requests.get( - f"{API_URL}/customers/{CUST_ID}/current_usage", - params={"external_subscription_id": SUB_ID}, - headers={"Authorization": f"Bearer {API_KEY}"}, - timeout=15, - verify=VERIFY_SSL, - ) - r.raise_for_status() - out: dict[str, float] = {} - for c in r.json().get("customer_usage", {}).get("charges_usage", []) or []: - code = c.get("billable_metric", {}).get("code", "") - out[code] = float(c.get("units", 0) or 0) - return out - - -def test_emit_then_reconcile_with_live_lago(): - """Send 5 known-shape events; assert input/output totals incremented correctly.""" - sdk = LagoSDK( - api_key=API_KEY, - api_url=API_URL, - default_subscription_id=SUB_ID, - verify_ssl=VERIFY_SSL, - ) - - before = _read_usage() - in_before = before.get("llm_input_tokens", 0.0) - out_before = before.get("llm_output_tokens", 0.0) - - # Emit 5 events with stable values for arithmetic - for _ in range(5): - sdk.emit( - CanonicalUsage( - input=100, - output=200, - model="claude-sonnet-4-6", - provider="anthropic", - api="bedrock_invoke", - ) - ) - - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=3.0) - - # Lago is async — poll for up to 30s - deadline = time.time() + 30 - after = before - while time.time() < deadline: - after = _read_usage() - in_delta = after.get("llm_input_tokens", 0.0) - in_before - out_delta = after.get("llm_output_tokens", 0.0) - out_before - if in_delta >= 500 and out_delta >= 1000: - break - time.sleep(1.0) - - in_delta = after.get("llm_input_tokens", 0.0) - in_before - out_delta = after.get("llm_output_tokens", 0.0) - out_before - assert in_delta == 500, f"input delta {in_delta} != 500 — events lost or duplicated" - assert out_delta == 1000, f"output delta {out_delta} != 1000 — events lost or duplicated" diff --git a/tests/integration/test_live_anthropic.py b/tests/integration/test_live_anthropic.py deleted file mode 100644 index 73c4e35..0000000 --- a/tests/integration/test_live_anthropic.py +++ /dev/null @@ -1,146 +0,0 @@ -"""End-to-end Anthropic integration test — live API + mocked Lago. - -Skipped unless ANTHROPIC_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("ANTHROPIC_API_KEY"), - reason="ANTHROPIC_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def test_live_anthropic_messages_create_emits_to_lago() -> None: - from anthropic import Anthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - client.messages.create( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "native" - assert e["properties"]["provider"] == "anthropic" - finally: - server.shutdown() - - -def test_live_anthropic_streaming_emits_from_final_delta() -> None: - from anthropic import Anthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - for _ in client.messages.create( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - stream=True, - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -def test_live_anthropic_messages_stream_context_manager() -> None: - from anthropic import Anthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - with client.messages.stream( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - ) as stream: - for _ in stream.text_stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -@pytest.mark.asyncio -async def test_live_async_anthropic_messages_stream_context_manager_emits() -> None: - """Live regression test for the async messages.stream(...) context manager. - - Bug: __aexit__ called the sync _emit_final, which invoked - get_final_message() without await. On AsyncMessageStream that method is - a coroutine, so the un-awaited object fell through to the adapter as {} - → zero usage emitted, plus a "coroutine was never awaited" RuntimeWarning. - """ - from anthropic import AsyncAnthropic - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])) - async with client.messages.stream( - model="claude-haiku-4-5-20251001", - max_tokens=20, - messages=[{"role": "user", "content": "Say hi"}], - ) as stream: - async for _ in stream.text_stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() diff --git a/tests/integration/test_live_bedrock.py b/tests/integration/test_live_bedrock.py deleted file mode 100644 index e972324..0000000 --- a/tests/integration/test_live_bedrock.py +++ /dev/null @@ -1,91 +0,0 @@ -"""End-to-end integration test — live Bedrock REST + mocked Lago endpoint. - -Skipped unless `AWS_BEARER_TOKEN_BEDROCK` is set. Mocks Lago so no real -events are sent. Verifies that wrapping the bearer-token REST flow -produces correctly-shaped events at the Lago HTTP boundary. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest -import requests - -from lago_agent_sdk import LagoSDK -from lago_agent_sdk.adapters import extract_bedrock_converse - -REGION = "eu-west-1" -PROMPT = "One sentence about dolphins." - -pytestmark = pytest.mark.skipif( - not os.environ.get("AWS_BEARER_TOKEN_BEDROCK"), - reason="AWS_BEARER_TOKEN_BEDROCK not set — skipping live Bedrock integration", -) - - -class _MockLagoHandler(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - length = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(length).decode("utf-8") - self.server.received_payloads.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): # silence - return - - -def _start_mock_lago() -> tuple[HTTPServer, str]: - server = HTTPServer(("127.0.0.1", 0), _MockLagoHandler) - server.received_payloads = [] # type: ignore[attr-defined] - t = threading.Thread(target=server.serve_forever, daemon=True) - t.start() - return server, f"http://127.0.0.1:{server.server_port}" - - -def _bearer_call_converse(api_key: str, model_id: str) -> dict: - url = f"https://bedrock-runtime.{REGION}.amazonaws.com/model/{model_id}/converse" - body = { - "messages": [{"role": "user", "content": [{"text": PROMPT}]}], - "inferenceConfig": {"maxTokens": 50}, - } - r = requests.post( - url, - headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, - json=body, - timeout=60, - ) - r.raise_for_status() - return r.json() - - -def test_live_converse_to_mocked_lago(): - api_key = os.environ["AWS_BEARER_TOKEN_BEDROCK"] - server, base_url = _start_mock_lago() - try: - sdk = LagoSDK(api_key="lago_dummy", api_url=base_url, default_subscription_id="sub_int") - model_id = "eu.amazon.nova-lite-v1:0" - # Use the bearer-token REST surface (works without IAM creds in env) - resp = _bearer_call_converse(api_key, model_id) - usage = extract_bedrock_converse(resp, model_id=model_id) - sdk.emit(usage) - assert sdk.flush(timeout=5.0) - sdk.shutdown(timeout=2.0) - - assert len(server.received_payloads) >= 1 # type: ignore[attr-defined] - events = [e for p in server.received_payloads for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["external_subscription_id"] == "sub_int" - assert e["properties"]["api"] == "bedrock_converse" - assert e["properties"]["provider"] == "amazon" - finally: - server.shutdown() diff --git a/tests/integration/test_live_gemini.py b/tests/integration/test_live_gemini.py deleted file mode 100644 index 4ac5de6..0000000 --- a/tests/integration/test_live_gemini.py +++ /dev/null @@ -1,154 +0,0 @@ -"""End-to-end Gemini integration test — live API + mocked Lago. - -Skipped unless GEMINI_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("GEMINI_API_KEY"), - reason="GEMINI_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def _collect_events(server) -> list[dict]: - return [e for p in server.received for e in p["events"]] - - -def _codes(events) -> set[str]: - return {e["code"] for e in events} - - -def test_live_gemini_generate_content_emits_to_lago() -> None: - from google import genai - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - client.models.generate_content( - model="gemini-2.5-flash", - contents="Say hi", - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "native" - assert e["properties"]["provider"] == "gemini" - finally: - server.shutdown() - - -def test_live_gemini_streaming_captures_usage_from_final_chunk() -> None: - from google import genai - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - for _ in client.models.generate_content_stream( - model="gemini-2.5-flash", - contents="Count from 1 to 3.", - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -def test_live_gemini_thinking_emits_reasoning() -> None: - """Gemini 2.5 emits thoughts_token_count → llm_reasoning_tokens event.""" - from google import genai - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - client.models.generate_content( - model="gemini-2.5-flash", - contents="What is 17 * 23? Show your reasoning step by step.", - ) - assert sdk.flush(timeout=15.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - # Gemini 2.5 reasons even without explicit thinking_config - assert "llm_reasoning_tokens" in codes - finally: - server.shutdown() - - -def test_live_gemini_tool_use_emits_tool_calls() -> None: - from google import genai - from google.genai import types as genai_types - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(genai.Client(api_key=os.environ["GEMINI_API_KEY"])) - weather_fn = genai_types.FunctionDeclaration( - name="get_weather", - description="Get the current weather for a city.", - parameters=genai_types.Schema( - type="OBJECT", - properties={"city": genai_types.Schema(type="STRING")}, - required=["city"], - ), - ) - client.models.generate_content( - model="gemini-2.5-flash", - contents="What's the weather in Tokyo?", - config=genai_types.GenerateContentConfig( - tools=[genai_types.Tool(function_declarations=[weather_fn])], - tool_config=genai_types.ToolConfig( - function_calling_config=genai_types.FunctionCallingConfig(mode="ANY"), - ), - ), - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - assert "llm_tool_calls" in _codes(events) - finally: - server.shutdown() diff --git a/tests/integration/test_live_mistral.py b/tests/integration/test_live_mistral.py deleted file mode 100644 index 72fe916..0000000 --- a/tests/integration/test_live_mistral.py +++ /dev/null @@ -1,120 +0,0 @@ -"""End-to-end Mistral integration test — live API + mocked Lago. - -Skipped unless MISTRAL_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("MISTRAL_API_KEY"), - reason="MISTRAL_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def test_live_mistral_chat_complete_emits_to_lago(): - from mistralai.client import Mistral - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Mistral(api_key=os.environ["MISTRAL_API_KEY"])) - client.chat.complete( - model="mistral-small-latest", - messages=[{"role": "user", "content": "Say hi"}], - max_tokens=20, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "native" - assert e["properties"]["provider"] == "mistral" - finally: - server.shutdown() - - -def test_live_mistral_chat_stream_emits_to_lago(): - from mistralai.client import Mistral - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Mistral(api_key=os.environ["MISTRAL_API_KEY"])) - for _ in client.chat.stream( - model="mistral-small-latest", - messages=[{"role": "user", "content": "Say hi"}], - max_tokens=20, - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -@pytest.mark.asyncio -async def test_live_mistral_chat_stream_async_emits_to_lago() -> None: - """Live regression test for chat.stream_async. - - Bug: the wrapper iterated `original_stream_async(*args, **kwargs)` without - awaiting it. In mistralai v2 this method is `async def`, so calling it - returns a coroutine — `async for` raises "got coroutine" TypeError. - """ - from mistralai.client import Mistral - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(Mistral(api_key=os.environ["MISTRAL_API_KEY"])) - stream = await client.chat.stream_async( - model="mistral-small-latest", - messages=[{"role": "user", "content": "Say hi"}], - max_tokens=20, - ) - async for _ in stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() diff --git a/tests/integration/test_live_openai.py b/tests/integration/test_live_openai.py deleted file mode 100644 index 04c5a8f..0000000 --- a/tests/integration/test_live_openai.py +++ /dev/null @@ -1,244 +0,0 @@ -"""End-to-end OpenAI integration test — live API + mocked Lago. - -Skipped unless OPENAI_API_KEY is set. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import pytest - -from lago_agent_sdk import LagoSDK - -pytestmark = pytest.mark.skipif( - not os.environ.get("OPENAI_API_KEY"), - reason="OPENAI_API_KEY not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def _collect_events(server) -> list[dict]: - return [e for p in server.received for e in p["events"]] - - -def _codes(events) -> set[str]: - return {e["code"] for e in events} - - -# -------------------------------------------------------------------------- -# Chat Completions -# -------------------------------------------------------------------------- -def test_live_openai_chat_completions_create_emits_to_lago() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - client.chat.completions.create( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Say hi"}], - max_completion_tokens=20, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "chat_completions" - assert e["properties"]["provider"] == "openai" - finally: - server.shutdown() - - -def test_live_openai_chat_completions_streaming_emits_from_final_chunk() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - # Note: stream_options.include_usage is auto-injected by the wrapper - for _ in client.chat.completions.create( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "Say hi"}], - max_completion_tokens=20, - stream=True, - ): - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - finally: - server.shutdown() - - -def test_live_openai_chat_completions_tool_use_emits_tool_calls() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - client.chat.completions.create( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "What's the weather in Tokyo?"}], - tools=[ - { - "type": "function", - "function": { - "name": "get_weather", - "description": "Get the current weather for a city.", - "parameters": { - "type": "object", - "properties": {"city": {"type": "string"}}, - "required": ["city"], - }, - }, - } - ], - tool_choice={"type": "function", "function": {"name": "get_weather"}}, - max_completion_tokens=200, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - assert "llm_tool_calls" in _codes(events) - finally: - server.shutdown() - - -def test_live_openai_reasoning_model_emits_reasoning_tokens() -> None: - """o-series models populate completion_tokens_details.reasoning_tokens. - First provider to actually expose this metric. - - Asserted against what the PROVIDER reported, not against the model choosing to - reason. `o4-mini` spends a variable number of reasoning tokens on the same - prompt — measured 0 on some calls and non-zero on others, minutes apart — and - since the SDK only emits non-zero fields, a hardcoded assertion made this test - a coin flip. It failed and passed on identical input in both repos, alternating - between them, which is exactly the kind of noise that hides a real regression. - - The SDK's contract is "emit reasoning tokens WHEN the provider reports them", - so that is what this checks; a call the model answered without reasoning has - nothing to assert and skips. - """ - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - resp = client.chat.completions.create( - model="o4-mini", - messages=[{"role": "user", "content": "What is 17 * 23? Just the number."}], - max_completion_tokens=2000, - ) - reported = int( - getattr(getattr(resp.usage, "completion_tokens_details", None), "reasoning_tokens", 0) or 0 - ) - assert sdk.flush(timeout=30.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - if reported == 0: - pytest.skip("o4-mini reported reasoning_tokens=0 for this call — nothing to emit") - assert "llm_reasoning_tokens" in codes # ← the key win for OpenAI - emitted = {e["code"]: int(float(e["properties"]["value"])) for e in events} - assert emitted["llm_reasoning_tokens"] == reported - finally: - server.shutdown() - - -# -------------------------------------------------------------------------- -# Responses API -# -------------------------------------------------------------------------- -def test_live_openai_responses_create_emits_to_lago() -> None: - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - client.responses.create( - model="gpt-4o-mini", - input="Say hi", - max_output_tokens=20, - ) - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "responses" - assert e["properties"]["provider"] == "openai" - finally: - server.shutdown() - - -def test_live_openai_responses_create_with_stream_emits_to_lago() -> None: - """Live regression test for two bugs in the Responses API streaming path: - - 1. The wrapper must NOT inject `stream_options.include_usage` — Responses - rejects that param and the call would fail with HTTP 400. - 2. The wrapper must extract usage from `event.response.usage` on the - terminal `response.completed` event (not from a top-level `event.usage`). - """ - from openai import OpenAI - - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = sdk.wrap(OpenAI(api_key=os.environ["OPENAI_API_KEY"])) - stream = client.responses.create( - model="gpt-4o-mini", - input="Say hi", - max_output_tokens=20, - stream=True, - ) - # Drain — also verifies the customer's call wasn't broken by injection. - for _ in stream: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = _collect_events(server) - codes = _codes(events) - assert "llm_input_tokens" in codes - assert "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "responses" - finally: - server.shutdown() diff --git a/tests/integration/test_live_pricing.py b/tests/integration/test_live_pricing.py deleted file mode 100644 index c03b104..0000000 --- a/tests/integration/test_live_pricing.py +++ /dev/null @@ -1,58 +0,0 @@ -"""Live pricing test — hits the real OpenRouter + AWS Bedrock bulk APIs. - -Skipped unless LAGO_LIVE_PRICING=1 (it makes real network calls, no keys needed -since both sources are public). Validates that the real fetchers build tables -and that known models resolve to sane USD-per-token prices — in particular it -exercises the AWS Bedrock offer-file parser against the live schema. -""" - -from __future__ import annotations - -import os -from decimal import Decimal - -import pytest - -from lago_agent_sdk.pricing import HttpPricingFetcher, lookup_bedrock, lookup_openrouter - -pytestmark = pytest.mark.skipif( - os.environ.get("LAGO_LIVE_PRICING") != "1", - reason="LAGO_LIVE_PRICING != 1 (live network test)", -) - - -def test_openrouter_live_table_and_known_models() -> None: - table = HttpPricingFetcher(timeout=30).fetch_openrouter() - exact = table["exact"] - assert len(exact) > 50, "expected a substantial OpenRouter model list" - - # A few well-known models should resolve with a positive input price. - resolved = 0 - for provider, model in [ - ("openai", "gpt-4o"), - ("anthropic", "claude-3.5-sonnet"), - ("google", "gemini-2.5-flash"), - ]: - mp = lookup_openrouter(table, provider, model) - if mp is not None and mp.input is not None and mp.input >= Decimal(0): - resolved += 1 - assert resolved >= 1, "expected at least one well-known OpenRouter model to resolve" - - -def test_bedrock_live_table_builds_and_resolves() -> None: - region = "us-east-1" - table = HttpPricingFetcher(timeout=30).fetch_bedrock(region) - # The parser should extract at least some priced models from the live offer. - assert table, "AWS Bedrock offer parsed to an empty table — schema may have changed" - priced = [mp for mp in table.values() if mp.input is not None or mp.output is not None] - assert priced, "no Bedrock models had input/output token prices" - - # A common Bedrock model should resolve (best-effort; logs the key on miss). - for model in [ - "anthropic.claude-3-5-sonnet-20240620-v1:0", - "anthropic.claude-3-haiku-20240307-v1:0", - ]: - mp = lookup_bedrock(table, model) - if mp is not None and (mp.input or mp.output): - return - pytest.skip(f"no probed Bedrock model matched; {len(table)} keys built — refine matcher if needed") diff --git a/tests/integration/test_live_streaming.py b/tests/integration/test_live_streaming.py deleted file mode 100644 index 0556c8d..0000000 --- a/tests/integration/test_live_streaming.py +++ /dev/null @@ -1,101 +0,0 @@ -"""Live streaming end-to-end against real Bedrock + mock Lago. - -Skipped unless AWS_BEARER_TOKEN_BEDROCK is set. Drives real -`converse_stream` and `invoke_model_with_response_stream` via the bearer -REST surface, reshaped into the same flow our wrapper drains. -""" - -from __future__ import annotations - -import json -import os -import threading -from http.server import BaseHTTPRequestHandler, HTTPServer - -import boto3 -import pytest - -from lago_agent_sdk import LagoSDK - -REGION = "eu-west-1" -PROMPT = "One sentence about dolphins." - -pytestmark = pytest.mark.skipif( - not os.environ.get("AWS_BEARER_TOKEN_BEDROCK"), - reason="AWS_BEARER_TOKEN_BEDROCK not set", -) - - -class _MockLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): # silence - return - - -def _spawn_lago(): - s = HTTPServer(("127.0.0.1", 0), _MockLago) - s.received = [] # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def _fresh_client(sdk: LagoSDK): - return sdk.wrap(boto3.client("bedrock-runtime", region_name=REGION)) - - -def test_live_converse_stream_emits_events(): - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = _fresh_client(sdk) - resp = client.converse_stream( - modelId="eu.amazon.nova-lite-v1:0", - messages=[{"role": "user", "content": [{"text": PROMPT}]}], - inferenceConfig={"maxTokens": 30}, - ) - # Drain — wrapper extracts usage from the metadata event - for _event in resp["stream"]: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes and "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "bedrock_converse" - finally: - server.shutdown() - - -def test_live_invoke_model_stream_emits_events(): - server, url = _spawn_lago() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_int") - client = _fresh_client(sdk) - body = json.dumps( - { - "anthropic_version": "bedrock-2023-05-31", - "max_tokens": 40, - "messages": [{"role": "user", "content": PROMPT}], - } - ) - resp = client.invoke_model_with_response_stream(modelId="eu.anthropic.claude-sonnet-4-6", body=body) - for _event in resp["body"]: - pass - assert sdk.flush(timeout=10.0) - sdk.shutdown(timeout=2.0) - events = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - codes = {e["code"] for e in events} - assert "llm_input_tokens" in codes and "llm_output_tokens" in codes - for e in events: - assert e["properties"]["api"] == "bedrock_invoke" - finally: - server.shutdown() diff --git a/tests/integration/test_outage_replay.py b/tests/integration/test_outage_replay.py deleted file mode 100644 index 0b6e0af..0000000 --- a/tests/integration/test_outage_replay.py +++ /dev/null @@ -1,96 +0,0 @@ -"""Outage replay — Lago fails for N seconds; events buffer and arrive in order on recovery.""" - -from __future__ import annotations - -import json -import threading -import time -from http.server import BaseHTTPRequestHandler, HTTPServer - -from lago_agent_sdk import CanonicalUsage, LagoSDK - - -class _ToggleableLago(BaseHTTPRequestHandler): - def do_POST(self): # noqa: N802 - n = int(self.headers.get("Content-Length", 0)) - body = self.rfile.read(n) - if self.server.failing: # type: ignore[attr-defined] - self.send_response(503) - self.end_headers() - return - self.server.received.append(json.loads(body)) # type: ignore[attr-defined] - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.end_headers() - self.wfile.write(b'{"ok": true}') - - def log_message(self, *_args, **_kwargs): - return - - -def _spawn(): - s = HTTPServer(("127.0.0.1", 0), _ToggleableLago) - s.received = [] # type: ignore[attr-defined] - s.failing = False # type: ignore[attr-defined] - threading.Thread(target=s.serve_forever, daemon=True).start() - return s, f"http://127.0.0.1:{s.server_port}" - - -def test_outage_replay_preserves_order_and_count(): - server, url = _spawn() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_test") - # Cap backoff low so the test doesn't take a minute - sdk._queue._max_retry_seconds = 1.0 # type: ignore[attr-defined] - - # 1. Lago is down — push 200 events - server.failing = True # type: ignore[attr-defined] - for i in range(200): - sdk.emit( - CanonicalUsage(input=1, model=f"m{i:03d}", provider="p", api="bedrock_invoke"), - ) - - # Give the queue worker a few attempts during the outage - time.sleep(2.0) - - # 2. Lago comes back - server.failing = False # type: ignore[attr-defined] - assert sdk.flush(timeout=15.0), "queue did not drain after recovery" - sdk.shutdown(timeout=2.0) - finally: - server.shutdown() - - flat = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - assert len(flat) == 200, f"expected 200 events, got {len(flat)}" - - # Order preserved — model field is m000, m001, ..., m199 - models = [e["properties"]["model"] for e in flat] - assert models == [f"m{i:03d}" for i in range(200)] - - -def test_long_outage_at_buffer_cap_drops_oldest_then_drains(): - """Outage long enough to overflow the (small) buffer — oldest dropped, rest drain.""" - server, url = _spawn() - try: - sdk = LagoSDK(api_key="x", api_url=url, default_subscription_id="sub_test") - # Tiny buffer + tiny backoff so the test runs quickly - sdk._queue._max_buffer_size = 30 # type: ignore[attr-defined] - sdk._queue._max_retry_seconds = 0.5 # type: ignore[attr-defined] - - server.failing = True # type: ignore[attr-defined] - # Push 50 — buffer caps at 30, so 20 oldest get dropped (model='m00'..'m19') - for i in range(50): - sdk.emit(CanonicalUsage(input=1, model=f"m{i:02d}", provider="p", api="bedrock_invoke")) - time.sleep(0.5) - - server.failing = False # type: ignore[attr-defined] - assert sdk.flush(timeout=15.0) - sdk.shutdown(timeout=2.0) - finally: - server.shutdown() - - flat = [e for p in server.received for e in p["events"]] # type: ignore[attr-defined] - # Expect exactly 30 events, the most recent ones (m20..m49) - assert len(flat) == 30 - models = sorted({e["properties"]["model"] for e in flat}) - assert models == [f"m{i:02d}" for i in range(20, 50)]