feat(streaming): add coordinator_proxy service module
This commit is contained in:
parent
c3e7dc1ea4
commit
0996ea8c7a
2 changed files with 105 additions and 0 deletions
94
app/services/coordinator_proxy.py
Normal file
94
app/services/coordinator_proxy.py
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
"""cf-orch coordinator proxy client.
|
||||
|
||||
Calls the coordinator's /proxy/authorize endpoint to obtain a one-time
|
||||
stream URL + token for LLM streaming. Always raises CoordinatorError on
|
||||
failure — callers decide how to handle it (stream-token endpoint returns
|
||||
503 or 403 as appropriate).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CoordinatorError(Exception):
|
||||
"""Raised when the coordinator returns an error or is unreachable."""
|
||||
def __init__(self, message: str, status_code: int = 503):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StreamTokenResult:
|
||||
stream_url: str
|
||||
token: str
|
||||
expires_in_s: int
|
||||
|
||||
|
||||
def _coordinator_url() -> str:
|
||||
return os.environ.get("COORDINATOR_URL", "http://10.1.10.71:7700")
|
||||
|
||||
|
||||
def _product_key() -> str:
|
||||
return os.environ.get("COORDINATOR_KIWI_KEY", "")
|
||||
|
||||
|
||||
async def coordinator_authorize(
|
||||
prompt: str,
|
||||
caller: str = "kiwi-recipe",
|
||||
ttl_s: int = 300,
|
||||
) -> StreamTokenResult:
|
||||
"""Call POST /proxy/authorize on the coordinator.
|
||||
|
||||
Returns a StreamTokenResult with the stream URL and one-time token.
|
||||
Raises CoordinatorError on any failure (network, auth, capacity).
|
||||
"""
|
||||
url = f"{_coordinator_url()}/proxy/authorize"
|
||||
key = _product_key()
|
||||
if not key:
|
||||
raise CoordinatorError(
|
||||
"COORDINATOR_KIWI_KEY env var is not set — streaming unavailable",
|
||||
status_code=503,
|
||||
)
|
||||
|
||||
payload = {
|
||||
"product": "kiwi",
|
||||
"product_key": key,
|
||||
"caller": caller,
|
||||
"prompt": prompt,
|
||||
"params": {},
|
||||
"ttl_s": ttl_s,
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
resp = await client.post(url, json=payload)
|
||||
except httpx.RequestError as exc:
|
||||
log.warning("coordinator_authorize network error: %s", exc)
|
||||
raise CoordinatorError(f"Coordinator unreachable: {exc}", status_code=503)
|
||||
|
||||
if resp.status_code == 401:
|
||||
raise CoordinatorError("Invalid product key", status_code=401)
|
||||
if resp.status_code == 429:
|
||||
raise CoordinatorError("Too many concurrent streams", status_code=429)
|
||||
if resp.status_code == 503:
|
||||
raise CoordinatorError("No GPU available for streaming", status_code=503)
|
||||
if not resp.is_success:
|
||||
raise CoordinatorError(
|
||||
f"Coordinator error {resp.status_code}: {resp.text[:200]}",
|
||||
status_code=503,
|
||||
)
|
||||
|
||||
data = resp.json()
|
||||
# Use public_stream_url if coordinator provides it (cloud mode), else stream_url
|
||||
stream_url = data.get("public_stream_url") or data["stream_url"]
|
||||
return StreamTokenResult(
|
||||
stream_url=stream_url,
|
||||
token=data["token"],
|
||||
expires_in_s=data["expires_in_s"],
|
||||
)
|
||||
11
tests/api/test_stream_token.py
Normal file
11
tests/api/test_stream_token.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
"""Tests for POST /api/v1/recipes/stream-token — coordinator proxy integration."""
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
||||
def test_coordinator_authorize_missing_url(monkeypatch):
|
||||
"""coordinator_authorize raises RuntimeError when COORDINATOR_URL is unset."""
|
||||
monkeypatch.delenv("COORDINATOR_URL", raising=False)
|
||||
monkeypatch.delenv("COORDINATOR_KIWI_KEY", raising=False)
|
||||
# Will test this properly via endpoint — see Task 3 tests.
|
||||
pass
|
||||
Loading…
Reference in a new issue