okay

app/api/v1/chat.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import json
 import logging
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.responses import StreamingResponse
 
 from app.api.deps import require_auth
-from app.services.chat_service import chat_completion
+from app.services.chat_service import chat_completion, _stream_chat_completion
 
 logger = logging.getLogger(__name__)
 
@@ -17,12 +19,50 @@ VALID_PROVIDERS = ["openprovider", "meganova", "aionlabs"]
 router = APIRouter()
 
 
+async def _stream_events(
+    body: Dict[str, Any],
+    request: Request,
+) -> AsyncGenerator[str, None]:
+    messages = body["messages"]
+    model = body.get("model", "agentdeck-1.0")
+    provider = body.get("provider")
+    max_tokens = body.get("max_tokens", 1024)
+    temperature = body.get("temperature", 0.7)
+    top_p = body.get("top_p", 0.9)
+    response_format = body.get("response_format")
+    redis = getattr(request.app.state, "redis", None)
+    scripts = getattr(request.app.state, "scripts", None)
+
+    try:
+        async for chunk in _stream_chat_completion(
+            messages=messages,
+            model=model,
+            provider=provider,
+            response_format=response_format,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            redis=redis,
+            scripts=scripts,
+        ):
+            yield f"data: {json.dumps(chunk)}\n\n"
+        yield "data: [DONE]\n\n"
+    except RuntimeError as e:
+        error_detail = str(e)
+        if "All AI providers exhausted" in error_detail:
+            error_chunk = {"error": {"message": "All API keys are currently locked. Retry after a few minutes.", "code": 503}}
+        else:
+            error_chunk = {"error": {"message": error_detail, "code": 502}}
+        yield f"data: {json.dumps(error_chunk)}\n\n"
+        yield "data: [DONE]\n\n"
+
+
 @router.post("/chat/completions")
 async def create_chat_completion(
     body: Dict[str, Any],
     request: Request,
     token: str = Depends(require_auth),
-) -> Dict[str, Any]:
+):
     start = time.monotonic()
     req_id = hex(int(time.time() * 1_000_000))[-8:]
 
@@ -78,6 +118,17 @@ async def create_chat_completion(
     redis = getattr(request.app.state, "redis", None)
     scripts = getattr(request.app.state, "scripts", None)
 
+    if stream:
+        return StreamingResponse(
+            _stream_events(body, request),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
+
     try:
         result = await chat_completion(
             messages=messages,

app/services/chat_service.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, AsyncGenerator, Dict, List, Optional
 
 import aiohttp
 from redis.asyncio import Redis
@@ -602,3 +602,336 @@ async def call_aion_labs_no_redis(
             continue
 
     return None
+
+
+async def _sse_chunks(resp: aiohttp.ClientResponse) -> AsyncGenerator[Dict[str, Any], None]:
+    buffer = b""
+    async for raw in resp.content.iter_any():
+        buffer += raw
+        while b"\n" in buffer:
+            line, buffer = buffer.split(b"\n", 1)
+            line_str = line.decode("utf-8").strip()
+            if line_str.startswith("data: "):
+                payload = line_str[6:]
+                if payload.strip() == "[DONE]":
+                    return
+                try:
+                    yield json.loads(payload)
+                except json.JSONDecodeError:
+                    continue
+
+
+async def _stream_meganova(
+    redis: Redis,
+    scripts: Dict[str, str],
+    messages: List[Dict[str, str]],
+    response_format: Optional[Dict[str, str]],
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    temperature: float = DEFAULT_TEMPERATURE,
+    top_p: float = DEFAULT_TOP_P,
+    target_model: Optional[str] = None,
+) -> AsyncGenerator[Dict[str, Any], None]:
+    if target_model and target_model not in MODELS:
+        return
+
+    total_unique_slots = len(MODELS) * len(KEY_IDS)
+    tried: set = set()
+    hard_cap = total_unique_slots * 2 + 2
+    loops = 0
+    max_tries = len(KEY_IDS) if target_model else total_unique_slots
+
+    while len(tried) < max_tries and loops < hard_cap:
+        loops += 1
+        slot = (
+            await _acquire_slot_for_model(redis, target_model)
+            if target_model
+            else await _acquire_slot(redis, scripts)
+        )
+        if not slot:
+            return
+
+        combo_key = f"{slot['keyId']}:{slot['modelIndex']}"
+        if combo_key in tried:
+            continue
+        tried.add(combo_key)
+
+        prepared = prepare_messages(messages, response_format)
+        api_key = KEY_MAP.get(slot["keyId"])
+        if not api_key:
+            continue
+
+        payload = {
+            "messages": prepared,
+            "model": slot["model"],
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "stream": True,
+        }
+
+        try:
+            timeout = aiohttp.ClientTimeout(total=_settings.request_timeout_ms / 1000)
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(
+                    f"{MEGANOVA_BASE_URL}{MEGANOVA_CHAT_PATH}",
+                    json=payload,
+                    headers={
+                        "Authorization": f"Bearer {api_key}",
+                        "Content-Type": "application/json",
+                    },
+                ) as resp:
+                    if resp.status >= 400:
+                        logger.warning(
+                            "MegaNova stream HTTP %d for keyId=%s model=%s",
+                            resp.status, slot["keyId"], slot["model"],
+                        )
+                        await _mark_failure(redis, scripts, slot["keyId"], slot["modelIndex"])
+                        continue
+                    async for chunk in _sse_chunks(resp):
+                        yield chunk
+                    return
+        except Exception as exc:
+            logger.warning("MegaNova stream failed: %s", exc)
+            continue
+
+
+async def _stream_meganova_no_redis(
+    messages: List[Dict[str, str]],
+    response_format: Optional[Dict[str, str]],
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    temperature: float = DEFAULT_TEMPERATURE,
+    top_p: float = DEFAULT_TOP_P,
+    target_model: Optional[str] = None,
+) -> AsyncGenerator[Dict[str, Any], None]:
+    api_key = _get_meganova_key()
+    if not api_key:
+        logger.info("No MegaNova keys configured, skipping stream")
+        return
+
+    prepared = prepare_messages(messages, response_format)
+    payload = {
+        "model": target_model or MODELS[1],
+        "messages": prepared,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "stream": True,
+    }
+
+    try:
+        timeout = aiohttp.ClientTimeout(total=_settings.request_timeout_ms / 1000)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.post(
+                f"{MEGANOVA_BASE_URL}{MEGANOVA_CHAT_PATH}",
+                json=payload,
+                headers={
+                    "Authorization": f"Bearer {api_key}",
+                    "Content-Type": "application/json",
+                },
+            ) as resp:
+                if resp.status != 200:
+                    logger.warning("MegaNova stream HTTP %d", resp.status)
+                    return
+                async for chunk in _sse_chunks(resp):
+                    yield chunk
+    except Exception as exc:
+        logger.warning("MegaNova stream failed: %s", exc)
+
+
+async def _stream_aion_labs(
+    redis: Redis,
+    messages: List[Dict[str, str]],
+    response_format: Optional[Dict[str, str]],
+    model: str = AION_LABS_DEFAULT_MODEL,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    temperature: float = DEFAULT_TEMPERATURE,
+    top_p: float = DEFAULT_TOP_P,
+) -> AsyncGenerator[Dict[str, Any], None]:
+    keys_json = await redis.get(f"{AION_PREFIX}:keys")
+    if not keys_json:
+        logger.info("[aion] No keys in Redis, skipping stream")
+        return
+    keys: List[str] = json.loads(keys_json)
+    if not keys:
+        return
+
+    total_tries = min(len(keys), 3)
+    for attempt in range(total_tries):
+        key = await _get_next_aion_key(redis)
+        if not key:
+            return
+
+        prepared = prepare_messages(messages, response_format)
+        payload = {
+            "model": model,
+            "messages": prepared,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "stream": True,
+        }
+
+        try:
+            timeout = aiohttp.ClientTimeout(total=_settings.request_timeout_ms / 1000)
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(
+                    f"{AION_LABS_BASE_URL}{AION_LABS_CHAT_PATH}",
+                    json=payload,
+                    headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
+                ) as resp:
+                    if resp.status != 200:
+                        logger.warning("[aion] Stream attempt %s HTTP %s", attempt + 1, resp.status)
+                        continue
+                    async for chunk in _sse_chunks(resp):
+                        yield chunk
+                    return
+        except Exception as exc:
+            logger.warning("[aion] Stream attempt %s failed: %s", attempt + 1, exc)
+            continue
+
+
+async def _stream_aion_labs_no_redis(
+    messages: List[Dict[str, str]],
+    response_format: Optional[Dict[str, str]],
+    model: str = AION_LABS_DEFAULT_MODEL,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    temperature: float = DEFAULT_TEMPERATURE,
+    top_p: float = DEFAULT_TOP_P,
+) -> AsyncGenerator[Dict[str, Any], None]:
+    keys_str = _settings.aion_lab_keys
+    if not keys_str:
+        logger.info("No AION keys configured, skipping stream")
+        return
+    keys = [k.strip() for k in keys_str.split(",") if k.strip()]
+    if not keys:
+        return
+
+    for attempt, key in enumerate(keys[:3]):
+        prepared = prepare_messages(messages, response_format)
+        payload = {
+            "model": model,
+            "messages": prepared,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "stream": True,
+        }
+
+        try:
+            timeout = aiohttp.ClientTimeout(total=_settings.request_timeout_ms / 1000)
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.post(
+                    f"{AION_LABS_BASE_URL}{AION_LABS_CHAT_PATH}",
+                    json=payload,
+                    headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
+                ) as resp:
+                    if resp.status != 200:
+                        continue
+                    async for chunk in _sse_chunks(resp):
+                        yield chunk
+                    return
+        except Exception:
+            continue
+
+
+async def _stream_openrouter_mimika(
+    messages: List[Dict[str, str]],
+    response_format: Optional[Dict[str, str]],
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    temperature: float = DEFAULT_TEMPERATURE,
+    top_p: float = DEFAULT_TOP_P,
+) -> AsyncGenerator[Dict[str, Any], None]:
+    api_key = _settings.openrouter_mimika_api_key
+    if not api_key:
+        logger.info("OPENROUTER_MIMIKA_API_KEY not set, skipping stream")
+        return
+
+    prepared = prepare_messages(messages, response_format)
+    payload = {
+        "model": OPENROUTER_MIMIKA_MODEL,
+        "messages": prepared,
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "stream": True,
+    }
+
+    try:
+        timeout = aiohttp.ClientTimeout(total=_settings.request_timeout_ms / 1000)
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            async with session.post(
+                f"{OPENROUTER_MIMIKA_BASE_URL}{OPENROUTER_MIMIKA_CHAT_PATH}",
+                json=payload,
+                headers={
+                    "Authorization": f"Bearer {api_key}",
+                    "Content-Type": "application/json",
+                },
+            ) as resp:
+                if resp.status != 200:
+                    logger.warning("OpenRouter Mimika stream HTTP %d", resp.status)
+                    return
+                async for chunk in _sse_chunks(resp):
+                    yield chunk
+    except Exception as exc:
+        logger.warning("OpenRouter Mimika stream failed: %s", exc)
+
+
+async def _stream_chat_completion(
+    messages: List[Dict[str, str]],
+    model: str = "agentdeck-1.0",
+    response_format: Optional[Dict[str, str]] = None,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+    temperature: float = DEFAULT_TEMPERATURE,
+    top_p: float = DEFAULT_TOP_P,
+    provider: Optional[str] = None,
+    redis: Optional[Redis] = None,
+    scripts: Optional[Dict[str, str]] = None,
+) -> AsyncGenerator[Dict[str, Any], None]:
+    messages = inject_system_identity(messages, model)
+
+    if not provider:
+        provider = "meganova"
+
+    target = model if model in MODELS else None
+
+    if provider == "aionlabs":
+        if redis and scripts:
+            gen = _stream_aion_labs(redis, messages, response_format, model or AION_LABS_DEFAULT_MODEL, max_tokens, temperature, top_p)
+        else:
+            gen = _stream_aion_labs_no_redis(messages, response_format, model or AION_LABS_DEFAULT_MODEL, max_tokens, temperature, top_p)
+        async for chunk in gen:
+            yield chunk
+        return
+
+    if provider == "openprovider":
+        async for chunk in _stream_openrouter_mimika(messages, response_format, max_tokens, temperature, top_p):
+            yield chunk
+        return
+
+    for _ in range(3):
+        if provider == "meganova":
+            if redis and scripts:
+                gen = _stream_meganova(redis, scripts, messages, response_format, max_tokens, temperature, top_p, target)
+            else:
+                gen = _stream_meganova_no_redis(messages, response_format, max_tokens, temperature, top_p, target)
+        elif provider == "aionlabs":
+            if redis and scripts:
+                gen = _stream_aion_labs(redis, messages, response_format, model or AION_LABS_DEFAULT_MODEL, max_tokens, temperature, top_p)
+            else:
+                gen = _stream_aion_labs_no_redis(messages, response_format, model or AION_LABS_DEFAULT_MODEL, max_tokens, temperature, top_p)
+        else:
+            gen = _stream_openrouter_mimika(messages, response_format, max_tokens, temperature, top_p)
+
+        yielded = False
+        async for chunk in gen:
+            yielded = True
+            yield chunk
+        if yielded:
+            return
+
+        if provider == "meganova":
+            provider = "aionlabs"
+        elif provider == "aionlabs":
+            provider = "openprovider"
+        else:
+            break