diff --git a/router/agy_proxy.py b/router/agy_proxy.py index 7a4d6f0..5343f2b 100644 --- a/router/agy_proxy.py +++ b/router/agy_proxy.py @@ -26,6 +26,7 @@ import time import httpx from typing import Optional, Protocol, runtime_checkable +from dataclasses import dataclass @runtime_checkable class CooldownPersistence(Protocol): @@ -181,30 +182,35 @@ def _wrap_response(text: str, model_name: str, prompt: str) -> dict: }, } -async def try_agy_proxy(prompt: str, messages: list = None, - session_id: str = None, - total_timeout: float = AGY_TOTAL_TIMEOUT_SECS, - stream: bool = False, - target_tier: str = "agent-advanced-core", - client: Optional[httpx.AsyncClient] = None, - cooldown_persistence: Optional[CooldownPersistence] = None) -> Optional[dict]: +@dataclass +class AgyProxyRequest: + prompt: str + messages: Optional[list] = None + session_id: Optional[str] = None + total_timeout: float = AGY_TOTAL_TIMEOUT_SECS + stream: bool = False + target_tier: str = "agent-advanced-core" + client: Optional[httpx.AsyncClient] = None + cooldown_persistence: Optional[CooldownPersistence] = None + +async def try_agy_proxy(request: AgyProxyRequest) -> Optional[dict]: """ Attempt agy proxy with session-aware tier fallback. Args: - prompt: Current user prompt - messages: Full message history for context - session_id: Router session identifier for conversation continuity - total_timeout: Max total time across all tiers - stream: If True, returns a dict with {"stream": async_generator, "model": model_name} - target_tier: Classified tier — "agent-reasoning-core" uses gemini-3.5-flash (low thinking), - "agent-advanced-core" uses full 2-tier chain (gemini-3.5-flash → claude-opus-4.6) - client: Shared HTTP client instance from caller - cooldown_persistence: Valkey synchronization callback interface + request: AgyProxyRequest containing all parameters Returns: OpenAI-compatible response dict, streaming dict, or None if all tiers failed. """ + prompt = request.prompt + messages = request.messages + session_id = request.session_id + total_timeout = request.total_timeout + stream = request.stream + target_tier = request.target_tier + client = request.client + cooldown_persistence = request.cooldown_persistence # Select model chain based on target tier # Reasoning: single tier, gemini-3.5-flash with low thinking # Advanced: full 2-tier chain (gemini-3.5-flash → claude-opus-4.6) diff --git a/router/main.py b/router/main.py index 2487b4b..9c371e0 100644 --- a/router/main.py +++ b/router/main.py @@ -2197,7 +2197,7 @@ async def chat_completions(request: Request): if should_try_agy: agy_span_obj = None try: - from agy_proxy import try_agy_proxy + from agy_proxy import try_agy_proxy, AgyProxyRequest last_prompt = "" for msg in reversed(messages): @@ -2239,7 +2239,7 @@ async def chat_completions(request: Request): pass is_stream_requested = body.get("stream", False) - agy_response = await try_agy_proxy( + agy_request = AgyProxyRequest( prompt=last_prompt, messages=messages, session_id=session_id, @@ -2249,6 +2249,7 @@ async def chat_completions(request: Request): client=get_http_client(), cooldown_persistence=ValkeyCooldownPersistence(), ) + agy_response = await try_agy_proxy(agy_request) if agy_response: model_name = agy_response.get("model", "gemini-3.5-flash (via agy)") diff --git a/tests/test_a2_verify.py b/tests/test_a2_verify.py index 60edee8..d1435b5 100644 --- a/tests/test_a2_verify.py +++ b/tests/test_a2_verify.py @@ -2,13 +2,13 @@ """Verify circuit breaker integration into agy_proxy.py""" try: from router.circuit_breaker import get_breaker - from router.agy_proxy import try_agy_proxy + from router.agy_proxy import try_agy_proxy, AgyProxyRequest except ImportError: import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from router.circuit_breaker import get_breaker - from router.agy_proxy import try_agy_proxy + from router.agy_proxy import try_agy_proxy, AgyProxyRequest import asyncio, time b = get_breaker() @@ -17,6 +17,6 @@ sub.cooldown_until = time.time() + 18000 sub.probe_granted = False -result = asyncio.run(try_agy_proxy('test prompt')) +result = asyncio.run(try_agy_proxy(AgyProxyRequest(prompt='test prompt'))) assert result is None, f'Breaker should return None when blocked, got: {result}' print('Integration verified: blocked breaker returns None from try_agy_proxy')