smoke_test_proxy.py

"""Smoke test: run the ACO proxy with a mock upstream, verify all features.

Usage: python smoke_test_proxy.py
"""
import json, time, threading, os, sys, socket

try:
    import fastapi, uvicorn, httpx, openai
except ImportError:
    os.system(f"{sys.executable} -m pip install fastapi uvicorn httpx openai -q")
    import fastapi, uvicorn, httpx, openai

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

# ── Mock upstream ──
mock_app = FastAPI(title="Mock OpenAI")
MOCK_COUNTS = {"calls": 0}

@mock_app.post("/v1/chat/completions")
async def mock_completions(request: Request):
    body = await request.json()
    model = body.get("model", "unknown")
    messages = body.get("messages", [])
    user_text = ""
    for m in messages:
        if m.get("role") == "user":
            user_text = str(m.get("content", ""))
            break

    MOCK_COUNTS["calls"] += 1
    ut = user_text.lower()

    if "bad-model" in model:
        return JSONResponse({"error": {"message": "not found"}}, status_code=404)

    if any(w in ut for w in ['capital', 'define', 'explain', 'who is']):
        content = "Paris is the capital of France."
    elif any(w in ut for w in ['fizzbuzz', 'reverse', 'lru', 'def ']):
        content = "```python\ndef fizzbuzz(n): ...\n```"
    elif 'search' in ut or 'current version' in ut:
        content = 'I need to search. <function=web_search><parameter=query>Python version</parameter></function>'
    elif 'compare' in ut or 'lora' in ut:
        content = "LoRA and QLoRA differ in quantization."
    elif 'email' in ut or 'client' in ut:
        content = "Dear Client,\n\nI must inform you..."
    else:
        content = "I'll help you with that. Could you provide more details?"

    prompt_tokens = sum(len(str(m.get("content", ""))) // 4 for m in messages)
    completion_tokens = len(content) // 4

    return {
        "id": f"mock-{MOCK_COUNTS['calls']}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": content},
                      "finish_reason": "stop"}],
        "usage": {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens,
                   "total_tokens": prompt_tokens + completion_tokens}
    }

def find_free_port():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(('', 0))
        return s.getsockname()[1]

def wait_for_url(url: str, timeout: float = 10.0) -> bool:
    t0 = time.time()
    while time.time() - t0 < timeout:
        try:
            r = httpx.get(url, timeout=2.0)
            if r.status_code < 500:
                return True
        except:
            pass
        time.sleep(0.3)
    return False

def test_proxy():
    upstream_port = find_free_port()
    proxy_port = find_free_port()
    print(f"Mock upstream:  http://localhost:{upstream_port}")
    print(f"Proxy:          http://localhost:{proxy_port}")

    # ── Start mock upstream ──
    upstream_cfg = uvicorn.Config(mock_app, host="127.0.0.1", port=upstream_port, log_level="error")
    upstream_srv = uvicorn.Server(upstream_cfg)
    threading.Thread(target=upstream_srv.run, daemon=True).start()
    if not wait_for_url(f"http://127.0.0.1:{upstream_port}", 5.0):
        print("ERROR: Mock upstream failed to start")
        return False
    print("  Mock upstream ready")

    # ── Download proxy.py from Hub ──
    proxy_url = "https://huggingface.co/narcolepticchicken/agent-cost-optimizer/resolve/main/aco/proxy.py"
    proxy_path = "/tmp/proxy.py"
    if not os.path.exists(proxy_path):
        import urllib.request
        print("  Downloading proxy.py...")
        urllib.request.urlretrieve(proxy_url, proxy_path)

    # ── Set env: ALL providers point to mock ──
    mock_base = f"http://127.0.0.1:{upstream_port}/v1"
    os.environ["OPENAI_BASE_URL"] = mock_base
    os.environ["ANTHROPIC_BASE_URL"] = mock_base
    os.environ["GOOGLE_BASE_URL"] = mock_base
    os.environ["DEEPSEEK_BASE_URL"] = mock_base
    os.environ["OPENAI_API_KEY"] = "mock-key"
    os.environ["ANTHROPIC_API_KEY"] = "mock-key"
    os.environ["GOOGLE_API_KEY"] = "mock-key"
    os.environ["DEEPSEEK_API_KEY"] = "mock-key"

    # ── Import proxy ──
    sys.path.insert(0, "/tmp")
    import importlib.util
    spec = importlib.util.spec_from_file_location("proxy", proxy_path)
    proxy_mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(proxy_mod)
    sys.modules["proxy"] = proxy_mod

    # ── Start proxy ──
    proxy_cfg = uvicorn.Config(proxy_mod.app, host="127.0.0.1", port=proxy_port, log_level="error")
    proxy_srv = uvicorn.Server(proxy_cfg)
    threading.Thread(target=proxy_srv.run, daemon=True).start()
    if not wait_for_url(f"http://127.0.0.1:{proxy_port}", 5.0):
        print("ERROR: Proxy failed to start")
        return False
    print("  Proxy ready")

    client = openai.OpenAI(base_url=f"http://127.0.0.1:{proxy_port}/v1", api_key="mock")
    passed = 0

    # Test 1: Health
    print("\n─── Test 1: Health ───")
    r = httpx.get(f"http://127.0.0.1:{proxy_port}/health", timeout=5)
    print(f"  Status: {r.status_code}")
    assert r.status_code == 200
    passed += 1

    # Test 2: Simple QA
    print("\n─── Test 2: Simple QA ───")
    resp = client.chat.completions.create(
        model="gpt-5-mini",
        messages=[{"role": "user", "content": "What is the capital of France?"}],
        max_tokens=100,
    )
    content = resp.choices[0].message.content
    aco_model = getattr(resp.usage, "aco_model", "N/A") if hasattr(resp, "usage") else "N/A"
    aco_cost = getattr(resp.usage, "aco_cost_usd", 0) if hasattr(resp, "usage") else 0
    print(f"  ACO model: {aco_model} | Cost: ${aco_cost:.6f}")
    print(f"  Content: {content[:60]}...")
    assert "Paris" in content, f"Expected Paris, got: {content[:80]}"
    passed += 1

    # Test 3: Coding floor (tier 2 min)
    print("\n─── Test 3: Coding floor ───")
    resp = client.chat.completions.create(
        model="gpt-5-mini",
        messages=[{"role": "user", "content": "Write a Python function to reverse a string."}],
        max_tokens=100,
    )
    aco_model = getattr(resp.usage, "aco_model", "N/A")
    print(f"  ACO model: {aco_model}")
    assert "deepseek-v4-flash" != aco_model, "Coding should not downgrade to tier 1"
    passed += 1

    # Test 4: Tier 3+ downgrade for short text
    print("\n─── Test 4: Tier 3 downgrade ───")
    resp = client.chat.completions.create(
        model="gemini-2.5-pro",  # tier 3
        messages=[{"role": "user", "content": "What is 2+2?"}],
        max_tokens=100,
    )
    aco_model = getattr(resp.usage, "aco_model", "N/A")
    aco_tier = getattr(resp.usage, "aco_tier", 0)
    print(f"  Requested: gemini-2.5-pro (tier 3) → ACO: {aco_model} (tier {aco_tier})")
    assert aco_tier == 1, f"Expected tier 1 for short text, got tier {aco_tier}"
    passed += 1

    # Test 5: Tool gating
    print("\n─── Test 5: Tool gating ───")
    resp = client.chat.completions.create(
        model="gpt-5-mini",
        messages=[{"role": "user", "content": "What is the capital of France?"}],
        tools=[{"type": "function", "function": {"name": "search", "description": "Search web",
                  "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}}}],
        max_tokens=100,
    )
    tool_gated = getattr(resp.usage, "aco_tool_gated", False) if hasattr(resp, "usage") else False
    print(f"  Tool gated: {tool_gated}")
    passed += 1

    # Test 6: Dashboard
    print("\n─── Test 6: Dashboard ───")
    r = httpx.get(f"http://127.0.0.1:{proxy_port}/dashboard", timeout=5)
    assert r.status_code == 200 and "ACO Proxy" in r.text
    print(f"  ✓ Dashboard ({len(r.text)} chars)")
    passed += 1

    # Test 7: Telemetry JSON
    print("\n─── Test 7: Telemetry ───")
    r = httpx.get(f"http://127.0.0.1:{proxy_port}/telemetry", timeout=5)
    data = r.json()
    print(f"  Calls: {data['total_calls']}, Cost: ${data['total_cost']:.6f}")
    assert data["total_calls"] >= 4
    passed += 1

    # Test 8: Context compression
    print("\n─── Test 8: Context compression ───")
    long_log = "Traceback (most recent call last):\n  File \"app.py\", line 42\n" * 30
    long_log += "\nERROR: Connection to database timed out\n" * 5
    resp = client.chat.completions.create(
        model="gpt-5-mini",
        messages=[{"role": "system", "content": "You are a log analyzer."},
                  {"role": "user", "content": long_log}],
        max_tokens=100,
    )
    ratio = getattr(resp.usage, "aco_compression_ratio", 1.0) if hasattr(resp, "usage") else 1.0
    print(f"  Compression ratio: {ratio:.2f}")
    assert ratio < 1.0, f"Expected compression, got {ratio}"
    passed += 1

    # Test 9: Telemetry reset
    print("\n─── Test 9: Telemetry reset ───")
    r = httpx.get(f"http://127.0.0.1:{proxy_port}/telemetry/reset", timeout=5)
    assert r.json()["status"] == "ok"
    r2 = httpx.get(f"http://127.0.0.1:{proxy_port}/telemetry", timeout=5)
    assert r2.json()["total_calls"] == 0
    print("  ✓ Reset works")
    passed += 1

    # ── Results ──
    print(f"\n{'='*60}")
    print(f"  ALL {passed} TESTS PASSED ✓")
    print(f"  Mock upstream calls: {MOCK_COUNTS['calls']}")
    print(f"{'='*60}")

    upstream_srv.should_exit = True
    proxy_srv.should_exit = True
    return True


if __name__ == "__main__":
    test_proxy()