Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
@@ -0,0 +1,371 @@
|
||||
"""HTTP/2-first client with an aiohttp fallback for dependency-light installs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, AsyncIterator, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError: # Manual plugin copies may not install new requirements.
|
||||
httpx = None
|
||||
|
||||
|
||||
def _verbose_http_logging_enabled() -> bool:
|
||||
value = os.environ.get("O1KEY_VERBOSE_LOG", "")
|
||||
return value.strip().lower() not in ("", "0", "false", "no", "off")
|
||||
|
||||
|
||||
# httpx emits one INFO line for every polling request. Those access logs drown
|
||||
# out the batch lifecycle and add no information on successful requests. Keep
|
||||
# warnings/errors, and allow the existing verbose switch to restore raw logs.
|
||||
if not _verbose_http_logging_enabled():
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
logging.getLogger("httpcore").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
HTTPX_AVAILABLE = httpx is not None
|
||||
HTTP_CLIENT_ERRORS = (aiohttp.ClientError,)
|
||||
HTTP_STREAM_ERRORS = (aiohttp.ClientPayloadError,)
|
||||
if HTTPX_AVAILABLE:
|
||||
HTTP_CLIENT_ERRORS += (httpx.HTTPError,)
|
||||
HTTP_STREAM_ERRORS += (httpx.StreamError,)
|
||||
|
||||
|
||||
class ResponseBodyIntegrityError(OSError):
|
||||
"""The response body ended early or disagreed with Content-Length."""
|
||||
|
||||
|
||||
class ResponseTaskIdMismatchError(OSError):
|
||||
"""A task query returned a different task ID than the one requested."""
|
||||
|
||||
|
||||
def _response_header(response: Any, name: str) -> str:
|
||||
headers = getattr(response, "headers", None) or {}
|
||||
value = headers.get(name)
|
||||
if value is None:
|
||||
value = headers.get(name.lower())
|
||||
if value is None:
|
||||
for key, candidate in headers.items():
|
||||
if str(key).lower() == name.lower():
|
||||
value = candidate
|
||||
break
|
||||
return "" if value is None else str(value).strip()
|
||||
|
||||
|
||||
def response_body_diagnostics(response: Any, received_bytes: int) -> dict[str, Any]:
|
||||
"""Describe whether a fully-read response agrees with its declared length."""
|
||||
declared_text = _response_header(response, "Content-Length")
|
||||
content_encoding = _response_header(response, "Content-Encoding") or "identity"
|
||||
transfer_encoding = _response_header(response, "Transfer-Encoding") or "<none>"
|
||||
declared_bytes = None
|
||||
length_check = "not-declared"
|
||||
|
||||
if declared_text:
|
||||
try:
|
||||
declared_bytes = int(declared_text)
|
||||
except ValueError:
|
||||
length_check = "invalid-header"
|
||||
else:
|
||||
if content_encoding.lower() not in ("", "identity"):
|
||||
# aiohttp/httpx expose decoded bytes while Content-Length can
|
||||
# describe the compressed wire representation.
|
||||
length_check = "skipped-compressed"
|
||||
elif declared_bytes == received_bytes:
|
||||
length_check = "match"
|
||||
else:
|
||||
length_check = "mismatch"
|
||||
|
||||
return {
|
||||
"http_version": str(getattr(response, "http_version", "") or "HTTP"),
|
||||
"status": int(getattr(response, "status", 0) or 0),
|
||||
"content_length": declared_text or "<none>",
|
||||
"declared_bytes": declared_bytes,
|
||||
"received_bytes": int(received_bytes),
|
||||
"content_encoding": content_encoding,
|
||||
"transfer_encoding": transfer_encoding,
|
||||
"length_check": length_check,
|
||||
}
|
||||
|
||||
|
||||
def format_response_body_diagnostics(diagnostics: dict[str, Any]) -> str:
|
||||
return (
|
||||
f"{diagnostics['http_version']} {diagnostics['status']}"
|
||||
f" | Content-Length={diagnostics['content_length']}"
|
||||
f" | received={diagnostics['received_bytes']}B"
|
||||
f" | Content-Encoding={diagnostics['content_encoding']}"
|
||||
f" | Transfer-Encoding={diagnostics['transfer_encoding']}"
|
||||
f" | length_check={diagnostics['length_check']}"
|
||||
)
|
||||
|
||||
|
||||
async def read_response_body_with_diagnostics(
|
||||
response: Any,
|
||||
*,
|
||||
chunk_size: int = 64 * 1024,
|
||||
) -> tuple[bytes, dict[str, Any]]:
|
||||
"""Read a body while retaining the received byte count if the stream fails."""
|
||||
buffer = bytearray()
|
||||
try:
|
||||
content = getattr(response, "content", None)
|
||||
iter_chunked = getattr(content, "iter_chunked", None)
|
||||
if callable(iter_chunked):
|
||||
async for chunk in iter_chunked(chunk_size):
|
||||
buffer.extend(chunk)
|
||||
else:
|
||||
read = getattr(response, "read", None)
|
||||
if callable(read):
|
||||
buffer.extend(await read())
|
||||
else:
|
||||
text = getattr(response, "text", None)
|
||||
if not callable(text):
|
||||
raise TypeError("response does not expose a readable body")
|
||||
buffer.extend((await text()).encode("utf-8"))
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
diagnostics = response_body_diagnostics(response, len(buffer))
|
||||
raise ResponseBodyIntegrityError(
|
||||
"响应体读取提前中断 | "
|
||||
f"{format_response_body_diagnostics(diagnostics)}"
|
||||
f" | cause={type(exc).__name__}: {exc}"
|
||||
) from exc
|
||||
|
||||
body = bytes(buffer)
|
||||
diagnostics = response_body_diagnostics(response, len(body))
|
||||
if diagnostics["length_check"] == "mismatch":
|
||||
raise ResponseBodyIntegrityError(
|
||||
"响应体长度与服务端声明不一致，可能在传输中提前中断 | "
|
||||
f"{format_response_body_diagnostics(diagnostics)}"
|
||||
)
|
||||
return body, diagnostics
|
||||
|
||||
|
||||
def response_task_id(payload: Any) -> Optional[str]:
|
||||
"""Return an explicit task_id/taskId without mistaking result item IDs for it."""
|
||||
if not isinstance(payload, dict):
|
||||
return None
|
||||
sources = [payload]
|
||||
for key in ("data", "result", "task"):
|
||||
nested = payload.get(key)
|
||||
if isinstance(nested, dict):
|
||||
sources.append(nested)
|
||||
for source in sources:
|
||||
for key in ("task_id", "taskId"):
|
||||
value = source.get(key)
|
||||
if value is not None and str(value).strip():
|
||||
return str(value).strip()
|
||||
return None
|
||||
|
||||
|
||||
def validate_response_task_id(payload: Any, expected_task_id: str) -> Optional[str]:
|
||||
actual_task_id = response_task_id(payload)
|
||||
if actual_task_id is not None and actual_task_id != str(expected_task_id):
|
||||
raise ResponseTaskIdMismatchError(
|
||||
f"任务查询响应 ID 不匹配 | requested_task_id={expected_task_id}"
|
||||
f" | response_task_id={actual_task_id}"
|
||||
)
|
||||
return actual_task_id
|
||||
|
||||
|
||||
def http2_runtime_available() -> bool:
|
||||
return HTTPX_AVAILABLE and importlib.util.find_spec("h2") is not None
|
||||
|
||||
|
||||
def create_timeout(
|
||||
total: float,
|
||||
*,
|
||||
connect: float,
|
||||
read: float,
|
||||
write: float,
|
||||
pool: float,
|
||||
):
|
||||
"""Return a timeout object understood by the selected HTTP backend."""
|
||||
if HTTPX_AVAILABLE:
|
||||
return httpx.Timeout(total, connect=connect, read=read, write=write, pool=pool)
|
||||
return aiohttp.ClientTimeout(
|
||||
total=total,
|
||||
connect=connect,
|
||||
sock_connect=connect,
|
||||
sock_read=read,
|
||||
)
|
||||
|
||||
|
||||
class _ResponseContent:
|
||||
def __init__(self, response: httpx.Response):
|
||||
self._response = response
|
||||
|
||||
async def iter_chunked(self, chunk_size: int) -> AsyncIterator[bytes]:
|
||||
async for chunk in self._response.aiter_bytes(chunk_size):
|
||||
yield chunk
|
||||
|
||||
|
||||
class HttpResponse:
|
||||
def __init__(self, response: httpx.Response):
|
||||
self._response = response
|
||||
self.content = _ResponseContent(response)
|
||||
|
||||
@property
|
||||
def status(self) -> int:
|
||||
return self._response.status_code
|
||||
|
||||
@property
|
||||
def headers(self):
|
||||
return self._response.headers
|
||||
|
||||
@property
|
||||
def raw_headers(self):
|
||||
return tuple(self._response.headers.raw)
|
||||
|
||||
@property
|
||||
def http_version(self) -> str:
|
||||
return self._response.http_version
|
||||
|
||||
async def text(self) -> str:
|
||||
await self._response.aread()
|
||||
return self._response.text
|
||||
|
||||
|
||||
class _RequestContext:
|
||||
def __init__(self, context):
|
||||
self._context = context
|
||||
|
||||
async def __aenter__(self) -> HttpResponse:
|
||||
response = await self._context.__aenter__()
|
||||
return HttpResponse(response)
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return await self._context.__aexit__(exc_type, exc, tb)
|
||||
|
||||
|
||||
class O1keyAsyncHttpClient:
|
||||
"""HTTP/2-first client; falls back to httpx HTTP/1.1, then aiohttp."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
max_connections: int = 32,
|
||||
max_keepalive_connections: int = 16,
|
||||
keepalive_expiry: float = 30.0,
|
||||
http2: bool = True,
|
||||
):
|
||||
self.http2_enabled = bool(http2 and http2_runtime_available())
|
||||
self.backend = "httpx" if HTTPX_AVAILABLE else "aiohttp"
|
||||
self._max_connections = max_connections
|
||||
self._keepalive_expiry = keepalive_expiry
|
||||
self._client = None
|
||||
if HTTPX_AVAILABLE:
|
||||
self._client = httpx.AsyncClient(
|
||||
http2=self.http2_enabled,
|
||||
verify=True,
|
||||
limits=httpx.Limits(
|
||||
max_connections=max_connections,
|
||||
max_keepalive_connections=max_keepalive_connections,
|
||||
keepalive_expiry=keepalive_expiry,
|
||||
),
|
||||
follow_redirects=False,
|
||||
)
|
||||
|
||||
async def __aenter__(self):
|
||||
if self.backend == "aiohttp":
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=self._max_connections,
|
||||
limit_per_host=self._max_connections,
|
||||
keepalive_timeout=self._keepalive_expiry,
|
||||
)
|
||||
self._client = aiohttp.ClientSession(connector=connector, trust_env=True)
|
||||
await self._client.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return await self._client.__aexit__(exc_type, exc, tb)
|
||||
|
||||
def post(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
data: Any = None,
|
||||
files: Any = None,
|
||||
timeout: Any = None,
|
||||
) -> _RequestContext:
|
||||
if self._client is None:
|
||||
raise RuntimeError("HTTP client must be entered before use")
|
||||
if self.backend == "httpx":
|
||||
return _RequestContext(self._client.stream(
|
||||
"POST",
|
||||
url,
|
||||
headers=headers,
|
||||
content=data if files is None else None,
|
||||
files=files,
|
||||
timeout=timeout,
|
||||
))
|
||||
|
||||
payload = data
|
||||
if files:
|
||||
payload = aiohttp.FormData()
|
||||
for field_name, (filename, source, content_type) in files.items():
|
||||
payload.add_field(
|
||||
field_name,
|
||||
source,
|
||||
filename=filename,
|
||||
content_type=content_type,
|
||||
)
|
||||
return self._client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
data=payload,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def get(
|
||||
self,
|
||||
url: str,
|
||||
*,
|
||||
headers: Optional[dict[str, str]] = None,
|
||||
allow_redirects: bool = False,
|
||||
timeout: Any = None,
|
||||
) -> _RequestContext:
|
||||
if self._client is None:
|
||||
raise RuntimeError("HTTP client must be entered before use")
|
||||
if self.backend == "httpx":
|
||||
return _RequestContext(self._client.stream(
|
||||
"GET",
|
||||
url,
|
||||
headers=headers,
|
||||
follow_redirects=allow_redirects,
|
||||
timeout=timeout,
|
||||
))
|
||||
return self._client.get(
|
||||
url,
|
||||
headers=headers,
|
||||
allow_redirects=allow_redirects,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
|
||||
def create_http_client(**kwargs) -> O1keyAsyncHttpClient:
|
||||
return O1keyAsyncHttpClient(**kwargs)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HTTPX_AVAILABLE",
|
||||
"HTTP_CLIENT_ERRORS",
|
||||
"HTTP_STREAM_ERRORS",
|
||||
"O1keyAsyncHttpClient",
|
||||
"ResponseBodyIntegrityError",
|
||||
"ResponseTaskIdMismatchError",
|
||||
"create_http_client",
|
||||
"create_timeout",
|
||||
"format_response_body_diagnostics",
|
||||
"http2_runtime_available",
|
||||
"read_response_body_with_diagnostics",
|
||||
"response_body_diagnostics",
|
||||
"response_task_id",
|
||||
"validate_response_task_id",
|
||||
]
|
||||
Reference in New Issue
Block a user