feat(llm): support incremental Responses WebSocket transport
This commit is contained in:
chenmingxuan committed
1 parent
73adcf25a7
commit
2a3f43fb60
22 files changed
+1176
-12
No files matched your search
@@ -0,0 +1,449 @@
|
||||
"""Responses WebSocket transport beneath the existing SDK and safety guards.
|
||||
|
||||
The SDK still serializes complete, authoritative requests. Only this wire boundary
|
||||
omits a verified prefix. No model messages, tool results, or screenshots are lost
|
||||
from local history, so any new connection can start from the full request.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from contextlib import suppress
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from uuid import uuid4
|
||||
|
||||
import httpx
|
||||
from websockets.asyncio.client import connect
|
||||
from websockets.exceptions import ConnectionClosed, InvalidStatus, SecurityError
|
||||
from websockets.protocol import State
|
||||
|
||||
from artemis.llm.diagnostics import ACTIVE_ATTEMPT
|
||||
from artemis.llm.http import model_http_options
|
||||
from artemis.llm.responses_diagnostics import identifier
|
||||
|
||||
IDLE_SECONDS = 60
|
||||
MAX_SESSIONS = 32
|
||||
MAX_CONNECTION_SECONDS = 55 * 60
|
||||
|
||||
|
||||
def _json(value):
|
||||
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
|
||||
|
||||
|
||||
def _fingerprint(item):
|
||||
# Assistant message IDs/status are transport metadata; store=false strips
|
||||
# message IDs on replay. Keep every semantic field, including call_id/phase.
|
||||
if isinstance(item, dict):
|
||||
item = dict(item)
|
||||
item.pop("status", None)
|
||||
if item.get("type") == "message" or item.get("role") == "assistant":
|
||||
item.pop("id", None)
|
||||
item.setdefault("type", "message")
|
||||
# Completion events may include token probabilities; LangChain
|
||||
# correctly omits those response-only statistics from replay input.
|
||||
if isinstance(item.get("content"), list):
|
||||
item["content"] = [
|
||||
{k: v for k, v in block.items() if k != "logprobs"}
|
||||
if isinstance(block, dict) and block.get("type") == "output_text"
|
||||
else block
|
||||
for block in item["content"]
|
||||
]
|
||||
if item.get("type") == "function_call" and isinstance(item.get("arguments"), str):
|
||||
with suppress(ValueError):
|
||||
item["arguments"] = json.loads(item["arguments"])
|
||||
return hashlib.sha256(_json(item).encode()).digest()
|
||||
|
||||
|
||||
class _Connect(connect):
|
||||
def process_redirect(self, exc):
|
||||
# Never forward the API credential to a redirected endpoint.
|
||||
if not isinstance(exc, InvalidStatus) or exc.response.status_code not in {
|
||||
301,
|
||||
302,
|
||||
303,
|
||||
307,
|
||||
308,
|
||||
}:
|
||||
return exc
|
||||
return SecurityError("Responses WebSocket redirects are disabled")
|
||||
|
||||
|
||||
class _Session:
|
||||
def __init__(self):
|
||||
self.lock = asyncio.Lock()
|
||||
self.users = 0
|
||||
self.connection = None
|
||||
self.opened_at = 0.0
|
||||
self.previous_id = None
|
||||
self.history = ()
|
||||
self.options = None
|
||||
self.timer = None
|
||||
self.handshake = {}
|
||||
|
||||
async def close(self):
|
||||
if self.timer:
|
||||
self.timer.cancel()
|
||||
self.timer = None
|
||||
self.previous_id = None
|
||||
self.history = ()
|
||||
self.options = None
|
||||
connection, self.connection = self.connection, None
|
||||
if connection is not None:
|
||||
with suppress(ConnectionClosed, OSError, TimeoutError):
|
||||
async with asyncio.timeout(1):
|
||||
await connection.close()
|
||||
|
||||
|
||||
class ResponsesWebSocketTransport(httpx.AsyncBaseTransport):
|
||||
"""One serial socket per conversation/role, independent sockets for concurrency."""
|
||||
|
||||
def __init__(self, *, timeout, use_proxy):
|
||||
self.timeout = timeout
|
||||
self.use_proxy = use_proxy
|
||||
self.http = httpx.AsyncClient(timeout=timeout, **model_http_options())
|
||||
self.sessions = OrderedDict()
|
||||
self.cleanup_tasks = set()
|
||||
|
||||
async def aclose(self):
|
||||
sessions = list(self.sessions.values())
|
||||
self.sessions.clear()
|
||||
await asyncio.gather(*(session.close() for session in sessions))
|
||||
if self.cleanup_tasks:
|
||||
await asyncio.gather(*tuple(self.cleanup_tasks), return_exceptions=True)
|
||||
await self.http.aclose()
|
||||
|
||||
def _close_background(self, session):
|
||||
task = asyncio.create_task(session.close())
|
||||
self.cleanup_tasks.add(task)
|
||||
task.add_done_callback(self.cleanup_tasks.discard)
|
||||
|
||||
async def _session(self, request, body):
|
||||
key = body.get("prompt_cache_key")
|
||||
# Header identity prevents credentials, routing, and different gateways
|
||||
# from sharing a socket. Per-request correlation headers aren't affinity.
|
||||
headers = {
|
||||
k: v
|
||||
for k, v in request.headers.items()
|
||||
if k not in {"content-length", "x-client-request-id", "x-client-attempt-id"}
|
||||
}
|
||||
identity = hashlib.sha256(_json([str(request.url), key, headers]).encode()).digest()
|
||||
reusable = isinstance(key, str) and bool(key)
|
||||
if reusable and identity in self.sessions:
|
||||
session = self.sessions.pop(identity)
|
||||
self.sessions[identity] = session
|
||||
else:
|
||||
if len(self.sessions) >= MAX_SESSIONS:
|
||||
for old_key, old in list(self.sessions.items()):
|
||||
if not old.users:
|
||||
del self.sessions[old_key]
|
||||
# Reserve the replacement before yielding: concurrent
|
||||
# requests for its scope must find the same session.
|
||||
self._close_background(old)
|
||||
break
|
||||
session = _Session()
|
||||
reusable = reusable and len(self.sessions) < MAX_SESSIONS
|
||||
if reusable:
|
||||
self.sessions[identity] = session
|
||||
acquired = False
|
||||
session.users += 1
|
||||
try:
|
||||
await session.lock.acquire()
|
||||
acquired = True
|
||||
finally:
|
||||
if not acquired:
|
||||
session.users -= 1
|
||||
self._idle(identity, session, reusable)
|
||||
if session.timer:
|
||||
session.timer.cancel()
|
||||
session.timer = None
|
||||
return identity, session, reusable
|
||||
|
||||
def release(self, identity, session, reusable):
|
||||
session.lock.release()
|
||||
session.users -= 1
|
||||
self._idle(identity, session, reusable)
|
||||
|
||||
def _idle(self, identity, session, reusable):
|
||||
if not reusable or session.users or self.sessions.get(identity) is not session:
|
||||
return
|
||||
|
||||
def expire():
|
||||
if session.users or self.sessions.get(identity) is not session:
|
||||
return
|
||||
del self.sessions[identity]
|
||||
self._close_background(session)
|
||||
|
||||
session.timer = asyncio.get_running_loop().call_later(IDLE_SECONDS, expire)
|
||||
|
||||
async def handle_async_request(self, request):
|
||||
if request.method != "POST" or not request.url.path.rstrip("/").endswith("/responses"):
|
||||
return await self.http.send(request, stream=True)
|
||||
body = json.loads(await request.aread())
|
||||
identity, session, reusable = await self._session(request, body)
|
||||
exchange = _Exchange(self, request, body, identity, session, reusable)
|
||||
stream_handed_off = False
|
||||
try:
|
||||
await exchange.start()
|
||||
first = await exchange.read()
|
||||
if first.get("type") == "error":
|
||||
status = first.get("status", 400)
|
||||
status = status if type(status) is int and 400 <= status <= 599 else 400
|
||||
return httpx.Response(status, json={"error": first.get("error", {})})
|
||||
if not body.get("stream"):
|
||||
event = first
|
||||
while not exchange.done:
|
||||
event = await exchange.read()
|
||||
if event.get("type") == "error":
|
||||
raise httpx.ReadError("Responses WebSocket request failed", request=request)
|
||||
return httpx.Response(
|
||||
200, json=event["response"], extensions={"artemis_websocket": True}
|
||||
)
|
||||
response = httpx.Response(
|
||||
200,
|
||||
headers={"content-type": "text/event-stream"},
|
||||
stream=_EventStream(exchange, first),
|
||||
extensions={"artemis_websocket": True},
|
||||
)
|
||||
stream_handed_off = True
|
||||
return response
|
||||
except InvalidStatus as error:
|
||||
return httpx.Response(
|
||||
error.response.status_code,
|
||||
json={
|
||||
"error": {
|
||||
"message": "Responses WebSocket 升级失败，请检查网关 WebSocket 支持或选择 HTTP。"
|
||||
}
|
||||
},
|
||||
headers={
|
||||
k: v
|
||||
for k, v in error.response.headers.raw_items()
|
||||
if k.lower() in {"retry-after", "x-request-id"}
|
||||
},
|
||||
)
|
||||
finally:
|
||||
if not stream_handed_off:
|
||||
await exchange.aclose()
|
||||
|
||||
|
||||
class _Exchange:
|
||||
def __init__(self, transport, request, body, identity, session, reusable):
|
||||
self.transport = transport
|
||||
self.request = request
|
||||
self.body = body
|
||||
self.identity = identity
|
||||
self.session = session
|
||||
self.reusable = reusable
|
||||
self.done = False
|
||||
self.closed = False
|
||||
self.recovered = False
|
||||
self.output_seen = False
|
||||
self.input = body.get("input")
|
||||
self.input_hashes = (
|
||||
tuple(map(_fingerprint, self.input)) if isinstance(self.input, list) else ()
|
||||
)
|
||||
self.options = hashlib.sha256(
|
||||
_json(
|
||||
{
|
||||
k: v
|
||||
for k, v in body.items()
|
||||
if k not in {"input", "stream", "background", "previous_response_id"}
|
||||
}
|
||||
).encode()
|
||||
).digest()
|
||||
self.explicit_state = bool(body.get("previous_response_id") or body.get("conversation"))
|
||||
self.attempt = ACTIVE_ATTEMPT.get()
|
||||
self.stats = {
|
||||
"mode": "websocket",
|
||||
"full_request_bytes": len(request.content),
|
||||
"sent_payload_bytes": 0,
|
||||
"received_payload_bytes": 0,
|
||||
"create_events": 0,
|
||||
"incremental": False,
|
||||
"reused_input_items": 0,
|
||||
"connection_reused": session.connection is not None,
|
||||
}
|
||||
if self.attempt is not None:
|
||||
self.attempt.data["transport"] = self.stats
|
||||
|
||||
async def start(self):
|
||||
session = self.session
|
||||
if session.connection is not None and (
|
||||
session.connection.state is not State.OPEN
|
||||
or time.monotonic() - session.opened_at >= MAX_CONNECTION_SECONDS
|
||||
):
|
||||
await session.close()
|
||||
self.stats["connection_reused"] = False
|
||||
if session.connection is None:
|
||||
url = self.request.url.copy_with(
|
||||
scheme="wss" if self.request.url.scheme == "https" else "ws"
|
||||
)
|
||||
headers = {
|
||||
k: v
|
||||
for k, v in self.request.headers.items()
|
||||
if k
|
||||
not in {"host", "content-length", "content-type", "connection", "accept-encoding"}
|
||||
}
|
||||
options = {"ssl": model_http_options()["verify"]} if url.scheme == "wss" else {}
|
||||
session.connection = await _Connect(
|
||||
str(url),
|
||||
additional_headers=headers,
|
||||
user_agent_header=None,
|
||||
proxy=True if self.transport.use_proxy else None,
|
||||
open_timeout=self.transport.timeout,
|
||||
close_timeout=1,
|
||||
max_size=64 * 1024 * 1024,
|
||||
**options,
|
||||
)
|
||||
session.opened_at = time.monotonic()
|
||||
session.handshake = {
|
||||
"client_request_id": identifier(headers.get("x-client-request-id")),
|
||||
"client_attempt_id": identifier(headers.get("x-client-attempt-id")),
|
||||
}
|
||||
if self.attempt:
|
||||
self.attempt.data.update(session.handshake, request_header_scope="websocket_handshake")
|
||||
for field, local in (
|
||||
("client_request_id", "request_id"),
|
||||
("client_attempt_id", "attempt_id"),
|
||||
):
|
||||
self.attempt.data[field + "_matches"] = session.handshake.get(
|
||||
field
|
||||
) == self.attempt.data.get(local)
|
||||
self.stats["handshake_http_status"] = 101
|
||||
await self.send(incremental=True)
|
||||
|
||||
async def send(self, *, incremental):
|
||||
session = self.session
|
||||
payload = {k: v for k, v in self.body.items() if k not in {"stream", "background"}}
|
||||
count = len(session.history)
|
||||
can_continue = (
|
||||
incremental
|
||||
and not self.explicit_state
|
||||
and session.previous_id
|
||||
and session.options == self.options
|
||||
and count
|
||||
and len(self.input_hashes) > count
|
||||
and self.input_hashes[:count] == session.history
|
||||
)
|
||||
if can_continue:
|
||||
payload["previous_response_id"] = session.previous_id
|
||||
payload["input"] = self.input[count:]
|
||||
self.stats.update(incremental=True, reused_input_items=count)
|
||||
else:
|
||||
self.stats.update(incremental=False, reused_input_items=0)
|
||||
self.stats["full_context_reason"] = (
|
||||
"explicit_state"
|
||||
if self.explicit_state
|
||||
else "history_or_options_changed"
|
||||
if session.previous_id
|
||||
else "new_connection_or_chain"
|
||||
)
|
||||
payload["type"] = "response.create"
|
||||
encoded = _json(payload)
|
||||
await session.connection.send(encoded)
|
||||
self.stats["sent_payload_bytes"] += len(encoded.encode())
|
||||
self.stats["create_events"] += 1
|
||||
self.stats["payload_bytes_saved"] = max(
|
||||
0, self.stats["full_request_bytes"] - self.stats["sent_payload_bytes"]
|
||||
)
|
||||
if self.attempt:
|
||||
self.attempt.data["request_sent_at"] = time.time()
|
||||
self.attempt.publish_if_due()
|
||||
|
||||
async def read(self):
|
||||
try:
|
||||
async with asyncio.timeout(self.transport.timeout):
|
||||
raw = await self.session.connection.recv()
|
||||
self.stats["received_payload_bytes"] += len(
|
||||
raw.encode() if isinstance(raw, str) else raw
|
||||
)
|
||||
event = json.loads(raw)
|
||||
if not isinstance(event, dict):
|
||||
raise ValueError("Responses WebSocket event is not an object")
|
||||
except ConnectionClosed as error:
|
||||
raise httpx.ReadError(
|
||||
"Responses WebSocket connection closed before completion", request=self.request
|
||||
) from error
|
||||
except (ValueError, UnicodeError) as error:
|
||||
raise httpx.ReadError(
|
||||
"Invalid Responses WebSocket event", request=self.request
|
||||
) from error
|
||||
kind = event.get("type", "")
|
||||
error = event.get("error") or {}
|
||||
if (
|
||||
kind == "error"
|
||||
and error.get("code") == "previous_response_not_found"
|
||||
and self.stats["incremental"]
|
||||
and not self.recovered
|
||||
and not self.output_seen
|
||||
):
|
||||
# An explicit rejection means no model output was generated. Recover
|
||||
# once from the authoritative input; never replay an ambiguous drop.
|
||||
self.recovered = True
|
||||
self.session.previous_id = None
|
||||
self.stats["chain_recovered"] = True
|
||||
await self.send(incremental=False)
|
||||
return await self.read()
|
||||
if kind not in {"error", "response.created", "response.in_progress", "response.queued"}:
|
||||
self.output_seen = True
|
||||
if kind in {
|
||||
"response.completed",
|
||||
"response.failed",
|
||||
"response.incomplete",
|
||||
"response.cancelled",
|
||||
}:
|
||||
self.done = True
|
||||
response = event.get("response") or {}
|
||||
self.session.previous_id = None
|
||||
if (
|
||||
kind == "response.completed"
|
||||
and response.get("status") == "completed"
|
||||
and isinstance(response.get("id"), str)
|
||||
and response["id"]
|
||||
and isinstance(response.get("output"), list)
|
||||
and self.input_hashes
|
||||
and not self.explicit_state
|
||||
):
|
||||
output = [
|
||||
item
|
||||
for item in response["output"]
|
||||
if not (
|
||||
item.get("type") == "reasoning"
|
||||
and not item.get("encrypted_content")
|
||||
and self.body.get("store") is False
|
||||
)
|
||||
]
|
||||
self.session.previous_id = response["id"]
|
||||
self.session.history = self.input_hashes + tuple(map(_fingerprint, output))
|
||||
self.session.options = self.options
|
||||
return event
|
||||
|
||||
async def aclose(self):
|
||||
if self.closed:
|
||||
return
|
||||
self.closed = True
|
||||
try:
|
||||
if not self.done or not self.reusable:
|
||||
await self.session.close()
|
||||
finally:
|
||||
self.transport.release(self.identity, self.session, self.reusable)
|
||||
|
||||
|
||||
class _EventStream(httpx.AsyncByteStream):
|
||||
def __init__(self, exchange, first):
|
||||
self.exchange = exchange
|
||||
self.first = first
|
||||
|
||||
async def __aiter__(self):
|
||||
try:
|
||||
event = self.first
|
||||
while True:
|
||||
yield b"data: " + _json(event).encode() + b"\n\n"
|
||||
if self.exchange.done or event.get("type") == "error":
|
||||
return
|
||||
event = await self.exchange.read()
|
||||
finally:
|
||||
await self.aclose()
|
||||
|
||||
async def aclose(self):
|
||||
await self.exchange.aclose()
|
||||
Reference in new issue
Block a user