313 lines
13 KiB
Python
313 lines
13 KiB
Python
import json
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import aclosing
|
|
from datetime import datetime, timezone
|
|
|
|
import httpx
|
|
|
|
from app.contracts import ModelEvent, ModelEventType, ModelRequest
|
|
from app.providers.base import ProviderError, ProviderTurn
|
|
from app.providers.tool_names import prepare_tool_names
|
|
|
|
|
|
class TurnStreamingMixin:
|
|
async def complete(self, request: ModelRequest) -> ProviderTurn:
|
|
raise NotImplementedError
|
|
|
|
async def stream(self, request: ModelRequest) -> AsyncIterator[ModelEvent]:
|
|
try:
|
|
turn = await self.complete(request)
|
|
sequence = 0
|
|
if turn.text:
|
|
yield ModelEvent(
|
|
event=ModelEventType.text_delta,
|
|
sequence=sequence,
|
|
data={"text": turn.text},
|
|
timestamp=datetime.now(timezone.utc),
|
|
)
|
|
sequence += 1
|
|
for call in turn.tool_calls:
|
|
yield ModelEvent(
|
|
event=ModelEventType.tool_call_start,
|
|
sequence=sequence,
|
|
data={
|
|
"tool_call_id": call.tool_call_id,
|
|
"name": call.name,
|
|
"arguments": call.arguments,
|
|
},
|
|
timestamp=datetime.now(timezone.utc),
|
|
)
|
|
sequence += 1
|
|
yield ModelEvent(
|
|
event=ModelEventType.tool_call_end,
|
|
sequence=sequence,
|
|
data={"tool_call_id": call.tool_call_id},
|
|
timestamp=datetime.now(timezone.utc),
|
|
)
|
|
sequence += 1
|
|
yield ModelEvent(
|
|
event=ModelEventType.usage,
|
|
sequence=sequence,
|
|
data={
|
|
"input_tokens": turn.input_tokens,
|
|
"output_tokens": turn.output_tokens,
|
|
},
|
|
timestamp=datetime.now(timezone.utc),
|
|
)
|
|
yield ModelEvent(
|
|
event=ModelEventType.done,
|
|
sequence=sequence + 1,
|
|
timestamp=datetime.now(timezone.utc),
|
|
)
|
|
except ProviderError as exc:
|
|
yield ModelEvent(
|
|
event=ModelEventType.error,
|
|
data={"code": exc.code, "message": exc.message},
|
|
timestamp=datetime.now(timezone.utc),
|
|
)
|
|
yield ModelEvent(
|
|
event=ModelEventType.done,
|
|
sequence=1,
|
|
timestamp=datetime.now(timezone.utc),
|
|
)
|
|
|
|
|
|
def decode_tool_arguments(value: object) -> dict[str, object]:
|
|
if isinstance(value, dict):
|
|
return value
|
|
if not isinstance(value, str):
|
|
raise ProviderError("PROVIDER_INVALID_RESPONSE", "Tool arguments are not JSON.")
|
|
try:
|
|
decoded = json.loads(value)
|
|
except json.JSONDecodeError as exc:
|
|
raise ProviderError("PROVIDER_INVALID_RESPONSE", "Tool arguments are invalid JSON.") from exc
|
|
if not isinstance(decoded, dict):
|
|
raise ProviderError("PROVIDER_INVALID_RESPONSE", "Tool arguments must be an object.")
|
|
return decoded
|
|
|
|
|
|
def invalid_response() -> ProviderError:
|
|
return ProviderError("PROVIDER_INVALID_RESPONSE", "Provider returned an invalid response.")
|
|
|
|
|
|
def truncated_stream() -> ProviderError:
|
|
return ProviderError("PROVIDER_STREAM_TRUNCATED", "Provider stream ended before completion.")
|
|
|
|
|
|
def object_value(value: object) -> dict:
|
|
if not isinstance(value, dict):
|
|
raise invalid_response()
|
|
return value
|
|
|
|
|
|
def list_value(value: object) -> list:
|
|
if not isinstance(value, list):
|
|
raise invalid_response()
|
|
return value
|
|
|
|
|
|
def string_value(value: object, *, nonempty: bool = False) -> str:
|
|
if not isinstance(value, str) or (nonempty and not value):
|
|
raise invalid_response()
|
|
return value
|
|
|
|
|
|
def token_count(value: object) -> int:
|
|
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
|
raise invalid_response()
|
|
return value
|
|
|
|
|
|
def remote_error(value: object) -> ProviderError:
|
|
# Never reflect upstream messages, URLs, request bodies or credentials.
|
|
error = value if isinstance(value, dict) else {}
|
|
code = error.get("code") or error.get("type")
|
|
mapping = {
|
|
"authentication_error": "PROVIDER_AUTH_FAILED",
|
|
"invalid_api_key": "PROVIDER_AUTH_FAILED",
|
|
"permission_error": "PROVIDER_AUTH_FAILED",
|
|
"rate_limit_error": "PROVIDER_RATE_LIMITED",
|
|
"rate_limit_exceeded": "PROVIDER_RATE_LIMITED",
|
|
"insufficient_quota": "PROVIDER_RATE_LIMITED",
|
|
"not_found_error": "MODEL_NOT_FOUND",
|
|
"model_not_found": "MODEL_NOT_FOUND",
|
|
"invalid_request_error": "PROVIDER_INVALID_REQUEST",
|
|
"context_length_exceeded": "PROVIDER_INVALID_REQUEST",
|
|
}
|
|
mapped = mapping.get(code, "PROVIDER_UNAVAILABLE") if isinstance(code, str) else "PROVIDER_UNAVAILABLE"
|
|
return ProviderError(mapped, "Provider could not complete the request.")
|
|
|
|
|
|
def check_error(data: dict) -> None:
|
|
if data.get("error") is not None or data.get("type") == "error":
|
|
raise remote_error(data.get("error") or data)
|
|
|
|
|
|
class UsageTracker:
|
|
"""Merge cumulative snapshots, including partial usage updates."""
|
|
|
|
def __init__(self, input_key: str = "input_tokens", output_key: str = "output_tokens",
|
|
*, cache_tokens: bool = False) -> None:
|
|
self.input_key = input_key
|
|
self.output_key = output_key
|
|
self.cache_tokens = cache_tokens
|
|
self.counts: dict[str, int] = {}
|
|
|
|
def update(self, value: object) -> dict[str, int]:
|
|
usage = object_value(value)
|
|
keys = [self.input_key, self.output_key]
|
|
if self.cache_tokens:
|
|
keys += ["cache_creation_input_tokens", "cache_read_input_tokens"]
|
|
for key in keys:
|
|
if key in usage:
|
|
self.counts[key] = max(self.counts.get(key, 0), token_count(usage[key]))
|
|
inputs = self.counts.get(self.input_key, 0)
|
|
if self.cache_tokens:
|
|
inputs += sum(self.counts.get(key, 0) for key in keys[2:])
|
|
return {"input_tokens": inputs, "output_tokens": self.counts.get(self.output_key, 0)}
|
|
|
|
|
|
class EventStreamingMixin:
|
|
async def stream(self, request: ModelRequest) -> AsyncIterator[ModelEvent]:
|
|
sequence = 0
|
|
status = "completed"
|
|
try:
|
|
request, originals = prepare_tool_names(request)
|
|
# Closing the public iterator must synchronously close every nested iterator.
|
|
async with aclosing(self._events(request)) as events:
|
|
async for kind, data in events:
|
|
if kind == ModelEventType.tool_call_start and "name" in data:
|
|
data = {**data, "name": originals.get(data["name"], data["name"])}
|
|
if kind == ModelEventType.usage:
|
|
data = {**data, "total_tokens": data["input_tokens"] + data["output_tokens"]}
|
|
yield ModelEvent(event=kind, data=data, sequence=sequence,
|
|
timestamp=datetime.now(timezone.utc))
|
|
sequence += 1
|
|
except ProviderError as exc:
|
|
status = "failed"
|
|
yield ModelEvent(event=ModelEventType.error, sequence=sequence,
|
|
data={"code": exc.code, "message": exc.message},
|
|
timestamp=datetime.now(timezone.utc))
|
|
sequence += 1
|
|
except (ValueError, TypeError, KeyError, IndexError, AttributeError, OverflowError):
|
|
status = "failed"
|
|
error = invalid_response()
|
|
yield ModelEvent(event=ModelEventType.error, sequence=sequence,
|
|
data={"code": error.code, "message": error.message},
|
|
timestamp=datetime.now(timezone.utc))
|
|
sequence += 1
|
|
# CancelledError and GeneratorExit deliberately propagate without a Done event.
|
|
yield ModelEvent(event=ModelEventType.done, sequence=sequence,
|
|
data={"status": status},
|
|
timestamp=datetime.now(timezone.utc))
|
|
|
|
|
|
async def sse_objects(response: httpx.Response) -> AsyncIterator[dict]:
|
|
"""Read SSE frames, accepting the adjacent data lines used by some gateways."""
|
|
parts: list[str] = []
|
|
event_name = ""
|
|
|
|
def decode() -> dict:
|
|
value = "\n".join(parts)
|
|
if value.strip() == "[DONE]":
|
|
return {"type": "[DONE]"}
|
|
try:
|
|
data = object_value(json.loads(value))
|
|
except (ValueError, TypeError) as exc:
|
|
raise invalid_response() from exc
|
|
if event_name and "type" not in data:
|
|
data["type"] = event_name
|
|
check_error(data)
|
|
return data
|
|
|
|
async for line in response.aiter_lines():
|
|
if not line:
|
|
if parts:
|
|
yield decode()
|
|
parts = []
|
|
event_name = ""
|
|
elif line.startswith(":"):
|
|
continue
|
|
elif line.startswith("event:"):
|
|
if parts:
|
|
yield decode()
|
|
parts = []
|
|
event_name = line[6:].strip()
|
|
elif line.startswith("data:"):
|
|
if parts:
|
|
# Legacy compatible endpoints sometimes omit blank separators.
|
|
try:
|
|
json.loads("\n".join(parts))
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
yield decode()
|
|
parts = []
|
|
event_name = ""
|
|
parts.append(line[5:].removeprefix(" "))
|
|
if parts:
|
|
yield decode()
|
|
|
|
|
|
class HTTPProviderMixin:
|
|
stream_path = "/chat/completions"
|
|
stream_format = "sse"
|
|
|
|
def _headers(self) -> dict[str, str]:
|
|
return {"Content-Type": "application/json"}
|
|
|
|
@staticmethod
|
|
def _status_error(exc: httpx.HTTPStatusError) -> ProviderError:
|
|
status = exc.response.status_code
|
|
code = {400: "PROVIDER_INVALID_REQUEST", 401: "PROVIDER_AUTH_FAILED",
|
|
403: "PROVIDER_AUTH_FAILED", 404: "MODEL_NOT_FOUND",
|
|
408: "PROVIDER_TIMEOUT", 413: "PROVIDER_INVALID_REQUEST",
|
|
422: "PROVIDER_INVALID_REQUEST", 429: "PROVIDER_RATE_LIMITED"}.get(
|
|
status, "PROVIDER_UNAVAILABLE")
|
|
return ProviderError(code, f"Provider returned HTTP {status}.")
|
|
|
|
async def _request(self, method: str, path: str, **kwargs) -> dict:
|
|
headers = self._headers()
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout_seconds, transport=self.transport) as client:
|
|
response = await client.request(method, f"{self.base_url}{path}", headers=headers, **kwargs)
|
|
response.raise_for_status()
|
|
data = object_value(response.json())
|
|
check_error(data)
|
|
return data
|
|
except httpx.TimeoutException as exc:
|
|
raise ProviderError("PROVIDER_TIMEOUT", "Provider request timed out.") from exc
|
|
except httpx.HTTPStatusError as exc:
|
|
raise self._status_error(exc) from exc
|
|
except httpx.HTTPError as exc:
|
|
raise ProviderError("PROVIDER_UNAVAILABLE", "Provider is unavailable.") from exc
|
|
except (ValueError, TypeError) as exc:
|
|
raise invalid_response() from exc
|
|
|
|
async def _stream_json(self, payload: dict[str, object]) -> AsyncIterator[dict]:
|
|
headers = self._headers()
|
|
headers["Accept"] = "text/event-stream" if self.stream_format == "sse" else "application/x-ndjson"
|
|
try:
|
|
async with httpx.AsyncClient(timeout=self.timeout_seconds, transport=self.transport) as client:
|
|
async with client.stream("POST", f"{self.base_url}{self.stream_path}",
|
|
headers=headers, json=payload) as response:
|
|
response.raise_for_status()
|
|
if self.stream_format == "sse":
|
|
async with aclosing(sse_objects(response)) as objects:
|
|
async for data in objects:
|
|
yield data
|
|
else:
|
|
async for line in response.aiter_lines():
|
|
if line.strip():
|
|
data = object_value(json.loads(line))
|
|
check_error(data)
|
|
yield data
|
|
except httpx.TimeoutException as exc:
|
|
raise ProviderError("PROVIDER_TIMEOUT", "Provider request timed out.") from exc
|
|
except httpx.HTTPStatusError as exc:
|
|
raise self._status_error(exc) from exc
|
|
except httpx.HTTPError as exc:
|
|
raise ProviderError("PROVIDER_UNAVAILABLE", "Provider is unavailable.") from exc
|
|
except (ValueError, TypeError) as exc:
|
|
raise invalid_response() from exc
|