From 55d9fbca88054e9e345da8ea2fc95f2d162e98e5 Mon Sep 17 00:00:00 2001 From: Von Date: Thu, 19 Mar 2026 19:48:57 +0800 Subject: [PATCH] =?UTF-8?q?update=20ui,=E5=A2=9E=E5=8A=A0websocket?= =?UTF-8?q?=E7=9A=84=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 29 ++- config.sample.yaml | 9 + config.yaml | 20 +- docs/architecture-4plus1.md | 35 +-- src/application/__init__.py | 11 +- src/config/__init__.py | 9 +- src/config/provider_config.py | 33 +++ src/config/provider_manager.py | 1 + src/external/__init__.py | 14 +- src/external/llm_provider.py | 1 + src/external/upstream_websocket.py | 203 +++++++++++++++++ src/presentation/static/css/admin-base.css | 28 +++ src/presentation/static/css/index.css | 17 +- src/presentation/static/css/providers.css | 129 +++++++++-- src/presentation/templates/base_page.html | 2 +- src/presentation/templates/index.html | 9 +- src/presentation/templates/providers.html | 132 ++++++++--- src/services/model_discovery_service.py | 8 +- src/services/provider_service.py | 10 +- src/services/proxy_service.py | 247 +++++++++++++++++---- src/utils/net.py | 32 ++- tests/test_proxy_transport.py | 173 +++++++++++++++ 22 files changed, 1022 insertions(+), 130 deletions(-) create mode 100644 src/external/upstream_websocket.py create mode 100644 tests/test_proxy_transport.py diff --git a/README.md b/README.md index 7797d40..7598a14 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,9 @@ - 提供 `POST /v1/chat/completions` 和 `GET /v1/models` - 多 Provider / 多模型路由 - 通过 `provider/model` 形式精确路由到不同上游 +- 双上游传输 + - `Provider` 运行时显式区分 `http` 与 `websocket` + - 对下游仍保持统一的 `/v1/chat/completions` 入口 - Hook 扩展机制 - 可在请求头、请求体、响应体三个阶段做定制化改写 - 非标准上游适配 @@ -146,7 +149,7 @@ Hook 适配能力的设计目标是支持合法授权前提下的协议兼容和 ### 1. 安装依赖 ```bash -pip install flask gevent requests pyyaml urllib3 +pip install flask gevent requests pyyaml urllib3 websocket-client ``` ### 2. 准备配置文件 @@ -157,6 +160,7 @@ pip install flask gevent requests pyyaml urllib3 - `server.port` - `providers[].name` - `providers[].api` +- `providers[].transport` - `providers[].model_list` 可选字段： @@ -179,6 +183,9 @@ pip install flask gevent requests pyyaml urllib3 - `chat.whitelist_enabled` - 为 `true` 时启用白名单访问控制，仅允许已录入且 `whitelist_access_enabled` 为启用状态的用户 IP 调用 `/v1/chat/completions` - 为 `false` 时关闭白名单限制，代理按普通开放模式工作 +- `providers[].transport` + - 支持 `http` 与 `websocket` + - 后台表单使用显式选择；如果你直接编辑配置文件并省略该字段，运行时会按 `api` 的 scheme 自动推断 参考样例见： @@ -221,6 +228,7 @@ curl http://127.0.0.1:8080/v1/chat/completions ^ providers: - name: openai api: https://api.openai.com/v1/chat/completions + transport: http api_key: ${OPENAI_API_KEY} model_list: - gpt-4.1 @@ -233,12 +241,31 @@ providers: providers: - name: anthropic api: https://example.com/v1/chat/completions + transport: http api_key: model_list: - claude-sonnet hook: anthropic_compat.py ``` +### WebSocket 上游 Provider + +```yaml +providers: + - name: codex + api: wss://example.com/v1/chat/completions + transport: websocket + api_key: ${CODEX_TOKEN} + model_list: + - gpt-4.1 +``` + +说明： + +- 对下游客户端仍然使用 `POST /v1/chat/completions` +- 代理层会把请求体作为 websocket 文本消息发给上游，并把上游消息桥接成 OpenAI 风格流式响应 +- 后台“拉取模型”会在 `api` 为 `ws://` 或 `wss://` 时，自动尝试对应的 `http://` 或 `https://` 地址去探测 `/v1/models` / `/models` + 这类 Hook 常见用途： - 把 OpenAI 风格 `messages` 转成上游要求的结构 diff --git a/config.sample.yaml b/config.sample.yaml index 5b8d1ac..2f0b62e 100644 --- a/config.sample.yaml +++ b/config.sample.yaml @@ -12,6 +12,7 @@ admin: providers: - name: volcengine api: https://ark.cn-beijing.volces.com/api/coding/v3/chat/completions + transport: http api_key: proxy: timeout_seconds: 300 @@ -23,6 +24,7 @@ providers: - name: volc api: https://ark.cn-beijing.volces.com/api/v3/chat/completions + transport: http api_key: model_list: - glm-4-7-251222 @@ -30,6 +32,13 @@ providers: - MiniMax-M2.5 - Kimi-K2.5 + # - name: codex + # api: wss://example.com/v1/chat/completions + # transport: websocket + # api_key: + # model_list: + # - gpt-4.1 + # logging: # 可选 # path: # level: diff --git a/config.yaml b/config.yaml index 0ce39de..4969e3b 100644 --- a/config.yaml +++ b/config.yaml @@ -1,24 +1,14 @@ server: port: 8080 host: 0.0.0.0 - chat: whitelist_enabled: true - admin: username: admin password: admin - providers: - - name: volcengine - api: https://ark.cn-beijing.volces.com/api/coding/v3/chat/completions - api_key: - model_list: - - kimi-k2.5 - -# logging: # 可选 -# path: -# level: - -# database: # 可选 -# path: +- name: volcengine + api: https://ark.cn-beijing.volces.com/api/coding/v3/chat/completions + transport: http + model_list: + - kimi-k2.5 diff --git a/docs/architecture-4plus1.md b/docs/architecture-4plus1.md index 41b8439..0912bb2 100644 --- a/docs/architecture-4plus1.md +++ b/docs/architecture-4plus1.md @@ -16,6 +16,7 @@ 这个项目当前最有区分度的能力不是“普通 API 代理”，而是： - 通过 Hook 机制在请求头、请求体、响应体三个阶段做协议适配 +- Provider 运行时显式携带 `transport`，把 HTTP/SSE 与 WebSocket 上游统一收敛到同一条代理链路 - 既能接标准 OpenAI 风格上游，也能接非标准协议上游 - 适合处理 Claude / Anthropic 一类需要格式转换的接入场景 - 也适合处理经过授权的会话型集成，例如额外 Cookie、session token、自定义 header 或私有参数 @@ -29,6 +30,7 @@ - 配置与 Provider 运行时：[config_manager.py](../src/config/config_manager.py) [provider_manager.py](../src/config/provider_manager.py) - Hook 协议：[contracts.py](../src/hooks/contracts.py) - 代理主链路：[proxy_controller.py](../src/presentation/proxy_controller.py) [proxy_service.py](../src/services/proxy_service.py) +- WebSocket 上游桥接：[upstream_websocket.py](../src/external/upstream_websocket.py) ## 1. Logical View @@ -60,7 +62,7 @@ flowchart LR ConfigMgr["ConfigManager"] ProviderCfg["provider_config"] ProviderMgr["ProviderManager"] - LLMProvider["LLMProvider"] + LLMProvider["LLMProvider\ntransport=http|websocket"] end subgraph Persistence["Persistence"] @@ -73,6 +75,7 @@ flowchart LR Hooks["Hook Modules"] Adapter["response_adapter"] Probe["stream_probe"] + WSBridge["upstream_websocket"] Upstream["Upstream LLM Providers"] end end @@ -102,6 +105,7 @@ flowchart LR ProxySvc --> LLMProvider ProxySvc --> Adapter ProxySvc --> Probe + ProxySvc --> WSBridge ProviderMgr --> ProviderCfg ProviderMgr --> LLMProvider ProviderMgr --> Hooks @@ -118,9 +122,9 @@ flowchart LR - `Presentation` 负责 HTTP 路由、鉴权入口、页面渲染和请求/响应封装 - `Services` 负责用例编排，不直接共享全局可变配置 -- `Config / Provider Runtime` 负责 YAML 配置快照、显式 `ProviderConfigSchema` / `RuntimeProviderSpec` 工厂、只读运行时视图和模型到 Provider 的运行时映射 +- `Config / Provider Runtime` 负责 YAML 配置快照、显式 `ProviderConfigSchema` / `RuntimeProviderSpec` 工厂、只读运行时视图、模型到 Provider 的运行时映射，以及 provider 传输类型解析 - `Persistence` 负责 SQLite 读写 -- `External Integration` 负责上游协议适配、流式探测、Hook 扩展和上游 Provider 集成 +- `External Integration` 负责上游协议适配、流式探测、WebSocket 消息桥接、Hook 扩展和上游 Provider 集成 其中最重要的扩展边界是 Hook： @@ -196,7 +200,7 @@ flowchart TB UserCache["UserService\nIP -> User 缓存"] ConfigCache["ConfigManager\n配置快照缓存"] ProviderRegistry["ProviderManager\nmodel -> provider 映射\nhook cache"] - HttpLocal["ProxyService\nthread-local requests.Session"] + TransportClients["ProxyService\nthread-local requests.Session\nper-request websocket connection"] end end @@ -216,13 +220,13 @@ flowchart TB FlaskApp --> UserCache FlaskApp --> ConfigCache FlaskApp --> ProviderRegistry - FlaskApp --> HttpLocal + FlaskApp --> TransportClients ConfigCache <--> ConfigFile UserCache <--> DB ProviderRegistry <--> HooksDir FlaskApp --> LogFiles - HttpLocal <--> Upstream + TransportClients <--> Upstream ``` 进程视图要点： @@ -231,7 +235,7 @@ flowchart TB - `AuthenticationService` 的 Session 存在进程内存中，重启后失效 - `UserService` 有 IP 维度缓存 - `ProviderManager` 持有运行时模型路由表、只读 `ProviderRuntimeView` 注册表和 Hook 缓存 -- `ProxyService` 维护 thread-local `requests.Session` +- `ProxyService` 维护 thread-local `requests.Session`，并在需要时建立短生命周期 websocket 上游连接 - Hook 模块以运行时动态加载方式参与请求路径 - 如果未来引入多实例部署，这些内存态要重新设计 @@ -271,16 +275,16 @@ flowchart LR App --> HooksDir App --> Logs - App <-->|HTTPS| P1 - App <-->|HTTPS| P2 - App <-->|HTTPS| PN + App <-->|HTTPS / WSS| P1 + App <-->|HTTPS / WSS| P2 + App <-->|HTTPS / WSS| PN ``` 物理视图要点： - 当前部署结构很简单，单机即可运行 - 本地状态包括配置文件、SQLite、日志文件、Hook 文件 -- 外部依赖主要是上游模型 Provider 的 HTTP API +- 外部依赖主要是上游模型 Provider 的 HTTP API 或 WebSocket API ## 5. Scenario View @@ -351,6 +355,7 @@ sequenceDiagram Admin->>ProviderCtl: GET /api/providers/fetch-models?api=... ProviderCtl->>DiscoverySvc: fetch_models_preview(...) + Note over DiscoverySvc: 如果 provider.api 为 ws:// 或 wss://\n先映射为对应的 http:// 或 https:// 再探测 DiscoverySvc->>Upstream: GET /v1/models or /models Upstream-->>DiscoverySvc: JSON models payload DiscoverySvc-->>ProviderCtl: fetched_models @@ -382,11 +387,11 @@ sequenceDiagram end ProxyCtl->>ProviderMgr: get_provider_for_model(model) - ProviderMgr-->>ProxyCtl: LLMProvider + ProviderMgr-->>ProxyCtl: LLMProvider(transport) ProxyCtl->>ProxySvc: proxy_request(provider, body, headers, on_complete) ProxySvc->>Hook: header_hook / input_body_hook - ProxySvc->>Upstream: POST provider.api - Upstream-->>ProxySvc: response or SSE stream + ProxySvc->>Upstream: HTTP POST or WebSocket connect + send JSON + Upstream-->>ProxySvc: JSON / SSE / WebSocket message stream ProxySvc->>Adapter: build_proxy_response(...) Adapter->>Hook: output_body_hook Adapter-->>Client: Flask Response @@ -431,6 +436,7 @@ sequenceDiagram - 控制平面和数据平面在职责上已经分离出清晰轮廓 - Hook 扩展点足够轻量，但表达力很强 - 兼容层能力使项目不局限于“标准 API Key 型接入” +- 通过显式 `transport` 与外部桥接层，HTTP/SSE 和 WebSocket 上游不再混杂在同一段响应解析逻辑里 当前架构边界： @@ -438,6 +444,7 @@ sequenceDiagram - `presentation` 层仍承担部分接口协议细节和错误映射责任 - `config` 子系统已经收敛为配置快照、显式 schema/factory、Provider 注册三类职责，但仍集中在同一 package 中 - Hook 动态加载带来了灵活性，也带来了调试、观测和测试复杂性 +- 当前 WebSocket 上游连接仍是“按请求建立、按请求关闭”，尚未像 CLIProxyAPI 那样引入长会话复用与增量输入状态机 ## 7. Evolution Suggestions diff --git a/src/application/__init__.py b/src/application/__init__.py index cac078d..57e8d65 100644 --- a/src/application/__init__.py +++ b/src/application/__init__.py @@ -2,7 +2,14 @@ # -*- coding: utf-8 -*- """应用层导出。""" -from .application import Application from .app_context import AppContext, Logger -__all__ = ['Application', 'AppContext', 'Logger'] +__all__ = ["Application", "AppContext", "Logger"] + + +def __getattr__(name: str): + if name == "Application": + from .application import Application + + return Application + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/config/__init__.py b/src/config/__init__.py index c6612e8..a334365 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -9,7 +9,6 @@ from .provider_config import ( RuntimeProviderSpec, build_provider_schemas, ) -from .provider_manager import ProviderManager __all__ = [ "ConfigManager", @@ -19,3 +18,11 @@ __all__ = [ "build_provider_schemas", "ProviderManager", ] + + +def __getattr__(name: str): + if name == "ProviderManager": + from .provider_manager import ProviderManager + + return ProviderManager + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/config/provider_config.py b/src/config/provider_config.py index 07f5084..def5bcb 100644 --- a/src/config/provider_config.py +++ b/src/config/provider_config.py @@ -6,12 +6,16 @@ from __future__ import annotations from dataclasses import dataclass from typing import Any, Dict, List, Mapping, Optional, Sequence +from urllib.parse import urlparse from ..utils.net import normalize_proxy_url DEFAULT_PROVIDER_TIMEOUT_SECONDS = 300 DEFAULT_PROVIDER_MAX_RETRIES = 3 DEFAULT_PROVIDER_VERIFY_SSL = False +DEFAULT_PROVIDER_TRANSPORT = "http" +SUPPORTED_PROVIDER_TRANSPORTS = {"http", "websocket"} +SUPPORTED_PROVIDER_API_SCHEMES = {"http", "https", "ws", "wss"} def clean_optional_string(value: Any) -> Optional[str]: @@ -91,12 +95,35 @@ def normalize_model_list(value: Any) -> List[str]: return models +def resolve_provider_transport(api: str, transport: Any = None) -> str: + """解析 provider 上游传输类型。""" + scheme = urlparse(api).scheme.lower() + if scheme not in SUPPORTED_PROVIDER_API_SCHEMES: + supported_schemes = ", ".join(sorted(SUPPORTED_PROVIDER_API_SCHEMES)) + raise ValueError(f"Provider api must use one of: {supported_schemes}") + + normalized_transport = clean_optional_string(transport) + if normalized_transport is not None: + lowered = normalized_transport.lower() + if lowered not in SUPPORTED_PROVIDER_TRANSPORTS: + supported = ", ".join(sorted(SUPPORTED_PROVIDER_TRANSPORTS)) + raise ValueError(f"Provider transport must be one of: {supported}") + if lowered == "http" and scheme in {"ws", "wss"}: + raise ValueError("Provider transport 'http' requires api to use http:// or https://") + return lowered + + if scheme in {"ws", "wss"}: + return "websocket" + return DEFAULT_PROVIDER_TRANSPORT + + @dataclass(frozen=True, slots=True) class ProviderConfigSchema: """显式表示配置文件中的单个 provider。""" name: str api: str + transport: str = DEFAULT_PROVIDER_TRANSPORT api_key: Optional[str] = None proxy: Optional[str] = None timeout_seconds: Optional[int] = None @@ -130,6 +157,7 @@ class ProviderConfigSchema: return cls( name=name, api=api, + transport=resolve_provider_transport(api, config.get("transport")), api_key=clean_optional_string(config.get("api_key")), proxy=normalize_proxy_url(config.get("proxy")), timeout_seconds=parse_optional_positive_int(config.get("timeout_seconds")), @@ -144,6 +172,7 @@ class ProviderConfigSchema: config: Dict[str, Any] = { "name": self.name, "api": self.api, + "transport": self.transport, } if self.api_key is not None: @@ -170,6 +199,7 @@ class RuntimeProviderSpec: name: str api: str + transport: str api_key: str model_list: tuple[str, ...] proxy: Optional[str] @@ -187,6 +217,7 @@ class RuntimeProviderSpec: return cls( name=config.name, api=config.api, + transport=config.transport, api_key=config.api_key or "", model_list=config.model_list, proxy=config.proxy, @@ -207,6 +238,7 @@ class ProviderRuntimeView: name: str api: str + transport: str model_list: tuple[str, ...] proxy: Optional[str] timeout_seconds: int @@ -219,6 +251,7 @@ class ProviderRuntimeView: return cls( name=spec.name, api=spec.api, + transport=spec.transport, model_list=spec.model_list, proxy=spec.proxy, timeout_seconds=spec.timeout_seconds, diff --git a/src/config/provider_manager.py b/src/config/provider_manager.py index 377b438..60ec190 100644 --- a/src/config/provider_manager.py +++ b/src/config/provider_manager.py @@ -72,6 +72,7 @@ class ProviderManager: provider = LLMProvider( name=spec.name, api=spec.api, + transport=spec.transport, api_key=spec.api_key, model_list=spec.model_list, proxy=spec.proxy, diff --git a/src/external/__init__.py b/src/external/__init__.py index 3ce9119..1474dea 100644 --- a/src/external/__init__.py +++ b/src/external/__init__.py @@ -5,5 +5,17 @@ from .llm_provider import LLMProvider from .response_adapter import build_proxy_response from .stream_probe import probe_stream_response +from .upstream_websocket import ( + StaticUpstreamResponse, + WebSocketUpstreamResponse, + collect_websocket_response_body, +) -__all__ = ["LLMProvider", "build_proxy_response", "probe_stream_response"] +__all__ = [ + "LLMProvider", + "StaticUpstreamResponse", + "WebSocketUpstreamResponse", + "build_proxy_response", + "collect_websocket_response_body", + "probe_stream_response", +] diff --git a/src/external/llm_provider.py b/src/external/llm_provider.py index 78bc5a6..663d29b 100644 --- a/src/external/llm_provider.py +++ b/src/external/llm_provider.py @@ -14,6 +14,7 @@ class LLMProvider: name: str api: str + transport: str = "http" api_key: Optional[str] = None model_list: tuple[str, ...] = () proxy: Optional[str] = None diff --git a/src/external/upstream_websocket.py b/src/external/upstream_websocket.py new file mode 100644 index 0000000..2f7a5e2 --- /dev/null +++ b/src/external/upstream_websocket.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""WebSocket 上游响应桥接。""" + +from __future__ import annotations + +import json +from typing import Any, Callable, Dict, Iterator, Optional + +from websocket import ABNF + + +class StaticUpstreamResponse: + """为非 requests 响应提供统一元数据与关闭能力。""" + + def __init__( + self, + status_code: int = 200, + headers: Optional[Dict[str, str]] = None, + on_close: Optional[Callable[[], None]] = None, + ): + self.status_code = status_code + self.headers = headers or {} + self._on_close = on_close + + def close(self) -> None: + if self._on_close is None: + return + on_close = self._on_close + self._on_close = None + on_close() + + +class WebSocketUpstreamResponse: + """把 websocket 消息流包装成 SSE 风格字节流。""" + + def __init__( + self, + connection: Any, + *, + status_code: int = 200, + headers: Optional[Dict[str, str]] = None, + ): + self._connection = connection + self.status_code = status_code + self.headers = headers or {"Content-Type": "text/event-stream"} + + def iter_content(self, chunk_size: Optional[int] = None) -> Iterator[bytes]: + del chunk_size + while True: + opcode, payload = self._connection.recv_data(control_frame=True) + if opcode == ABNF.OPCODE_CLOSE: + return + if opcode == ABNF.OPCODE_PING: + self._connection.pong(_coerce_bytes(payload)) + continue + if opcode not in {ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY}: + continue + + normalized = normalize_websocket_message(payload) + if normalized is not None: + yield normalized + + def close(self) -> None: + self._connection.close() + + +def collect_websocket_response_body(connection: Any, logger: Any = None) -> bytes: + """读取非流式 websocket 响应体。""" + payloads: list[bytes] = [] + + while True: + opcode, payload = connection.recv_data(control_frame=True) + if opcode == ABNF.OPCODE_CLOSE: + break + if opcode == ABNF.OPCODE_PING: + connection.pong(_coerce_bytes(payload)) + continue + if opcode not in {ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY}: + continue + + normalized = extract_websocket_payload(payload) + if normalized is None: + continue + + payloads.append(normalized) + if is_terminal_websocket_payload(normalized): + break + + if not payloads: + return b"" + + if len(payloads) > 1 and logger is not None: + logger.warning( + "WebSocket upstream returned %s messages for non-stream request, using the last payload", + len(payloads), + ) + return payloads[-1] + + +def normalize_websocket_message(payload: Any) -> Optional[bytes]: + """把单个 websocket 消息归一化成 SSE 事件。""" + text = _coerce_text(payload) + if text is None: + return None + + stripped = text.strip() + if not stripped: + return None + + if stripped.startswith(("data:", "event:", ":")): + if stripped.endswith("\n\n"): + return stripped.encode("utf-8") + return f"{stripped}\n\n".encode("utf-8") + + return f"data: {stripped}\n\n".encode("utf-8") + + +def extract_websocket_payload(payload: Any) -> Optional[bytes]: + """从 websocket 消息中提取最终 JSON/文本载荷。""" + text = _coerce_text(payload) + if text is None: + return None + + stripped = text.strip() + if not stripped: + return None + + sse_payload = _extract_last_sse_data(stripped) + if sse_payload is not None: + return sse_payload + + if stripped == "[DONE]": + return None + return stripped.encode("utf-8") + + +def is_terminal_websocket_payload(payload: bytes) -> bool: + """判断 websocket 载荷是否已到终止帧。""" + text = payload.decode("utf-8", errors="ignore").strip() + if not text: + return False + if text == "[DONE]": + return True + + try: + data = json.loads(text) + except json.JSONDecodeError: + return False + + if not isinstance(data, dict): + return False + if data.get("done") is True: + return True + + event_type = str(data.get("type") or "").strip().lower() + if event_type in {"response.completed", "response.done"}: + return True + + usage = data.get("usage") + choices = data.get("choices") + if isinstance(usage, dict) and (choices is None or choices == []): + return True + + if isinstance(choices, list): + for choice in choices: + if not isinstance(choice, dict): + continue + finish_reason = choice.get("finish_reason") + if finish_reason not in (None, ""): + return True + return False + + +def _extract_last_sse_data(text: str) -> Optional[bytes]: + data_lines: list[str] = [] + for raw_line in text.replace("\r\n", "\n").replace("\r", "\n").split("\n"): + if not raw_line.startswith("data:"): + continue + data = raw_line[5:].strip() + if not data or data == "[DONE]": + continue + data_lines.append(data) + + if not data_lines: + return None + return data_lines[-1].encode("utf-8") + + +def _coerce_text(payload: Any) -> Optional[str]: + if payload is None: + return None + if isinstance(payload, bytes): + return payload.decode("utf-8", errors="ignore") + return str(payload) + + +def _coerce_bytes(payload: Any) -> bytes: + if payload is None: + return b"" + if isinstance(payload, bytes): + return payload + return str(payload).encode("utf-8") diff --git a/src/presentation/static/css/admin-base.css b/src/presentation/static/css/admin-base.css index 2f5b1bf..fa637dc 100644 --- a/src/presentation/static/css/admin-base.css +++ b/src/presentation/static/css/admin-base.css @@ -232,6 +232,34 @@ body.app-page { background: var(--danger-hover); } +#ui-toast-container { + width: min(420px, calc(100vw - 24px)); + max-width: calc(100vw - 24px); + pointer-events: none; +} + +#ui-toast-container .toast { + width: 100%; + max-width: 100%; + pointer-events: auto; +} + +#ui-toast-container .toast .d-flex { + min-width: 0; + align-items: flex-start; +} + +#ui-toast-container .toast-body { + min-width: 0; + white-space: pre-wrap; + overflow-wrap: anywhere; + word-break: break-word; +} + +#ui-toast-container .btn-close { + flex-shrink: 0; +} + .app-page .btn, .app-page .header-nav-btn { min-height: 40px; diff --git a/src/presentation/static/css/index.css b/src/presentation/static/css/index.css index c31db6a..3054787 100644 --- a/src/presentation/static/css/index.css +++ b/src/presentation/static/css/index.css @@ -390,21 +390,28 @@ color: #c7e0ff; } -.dashboard-page .loading, -.dashboard-page .empty-state, -.dashboard-page .error { +.dashboard-page td.loading, +.dashboard-page td.empty-state, +.dashboard-page td.error { text-align: center; padding: 32px 16px; font-size: 14px; color: var(--text-secondary); } -.dashboard-page .error { +.dashboard-page td.empty-state p, +.dashboard-page td.loading p, +.dashboard-page td.error p { + margin: 0; + color: inherit; +} + +.dashboard-page td.error { color: #9b2c2c; background: #fef2f2; } -:root[data-theme="dark"] .dashboard-page .error { +:root[data-theme="dark"] .dashboard-page td.error { color: #ffb8bf; background: #3a1f24; } diff --git a/src/presentation/static/css/providers.css b/src/presentation/static/css/providers.css index 97b6222..c249049 100644 --- a/src/presentation/static/css/providers.css +++ b/src/presentation/static/css/providers.css @@ -233,11 +233,47 @@ } .providers-page .provider-name-cell { - font-weight: 700; color: var(--text-primary); white-space: nowrap; } +.providers-page .provider-name-line { + display: inline-flex; + align-items: center; + gap: 8px; + flex-wrap: wrap; +} + +.providers-page .provider-name-text { + font-weight: 700; +} + +.providers-page .provider-transport-badge { + display: inline-flex; + align-items: center; + min-height: 24px; + padding: 0 8px; + border: 1px solid transparent; + border-radius: 999px; + font-size: 11px; + font-weight: 800; + letter-spacing: 0.02em; + text-transform: uppercase; + white-space: nowrap; +} + +.providers-page .provider-transport-badge.is-http { + border-color: #bfd6ee; + background: linear-gradient(180deg, #f4f9ff 0%, #eaf3ff 100%); + color: #29507c; +} + +.providers-page .provider-transport-badge.is-websocket { + border-color: #cfe7d6; + background: linear-gradient(180deg, #f2fbf5 0%, #e8f6ee 100%); + color: #1d6b46; +} + .providers-page .provider-api-cell { width: 520px; max-width: 520px; @@ -289,6 +325,18 @@ box-shadow: inset 0 1px 0 rgba(125, 180, 255, 0.08); } +:root[data-theme="dark"] .providers-page .provider-transport-badge.is-http { + border-color: #2d5a8a; + background: linear-gradient(180deg, rgba(19, 35, 56, 0.96) 0%, rgba(14, 28, 45, 0.96) 100%); + color: #cfe5ff; +} + +:root[data-theme="dark"] .providers-page .provider-transport-badge.is-websocket { + border-color: #2a6b46; + background: linear-gradient(180deg, rgba(18, 48, 33, 0.96) 0%, rgba(14, 39, 27, 0.96) 100%); + color: #c8f0d7; +} + .providers-page .muted-text, .providers-page .loading, .providers-page .empty-state, @@ -673,6 +721,7 @@ display: inline-flex; align-items: center; gap: 10px; + min-height: 24px; min-width: 0; } @@ -684,6 +733,8 @@ .providers-page .field-header .form-label { margin-bottom: 0; + align-self: center; + line-height: 1; } .providers-page .field-meta-badge { @@ -697,6 +748,7 @@ color: var(--text-secondary); font-size: 12px; font-weight: 600; + line-height: 1; white-space: nowrap; } @@ -825,7 +877,6 @@ .providers-page .fetch-model-toolbar { display: flex; align-items: center; - gap: 12px; margin-bottom: 12px; } @@ -833,13 +884,6 @@ flex: 1 1 auto; } -.providers-page .fetch-model-toolbar-actions { - display: inline-flex; - align-items: center; - gap: 12px; - flex-shrink: 0; -} - .providers-page .fetch-model-summary, .providers-page .fetch-model-note { color: var(--text-secondary); @@ -858,13 +902,53 @@ .providers-page .fetch-model-list { height: min(52vh, 480px); overflow-y: auto; - padding: 6px; border: 1px solid var(--line); border-radius: var(--radius-md); background: rgba(17, 24, 39, 0.02); } -.providers-page .fetch-model-list > .empty-state { +.providers-page .fetch-model-list-header { + position: sticky; + top: 0; + z-index: 1; + display: flex; + align-items: center; + justify-content: space-between; + gap: 12px; + padding: 10px 12px; + border-bottom: 1px solid var(--line); + background: rgba(248, 250, 252, 0.98); + backdrop-filter: blur(12px); +} + +.providers-page .fetch-model-list-header-main { + min-width: 0; + display: inline-flex; + align-items: center; + gap: 10px; + margin: 0; + color: var(--text-primary); + font-size: 13px; + font-weight: 700; + cursor: pointer; +} + +.providers-page .fetch-model-list-header-title { + line-height: 1; +} + +.providers-page .fetch-model-list-header-hint { + color: var(--text-secondary); + font-size: 12px; + font-weight: 600; + white-space: nowrap; +} + +.providers-page .fetch-model-list-body { + padding: 6px; +} + +.providers-page .fetch-model-list-body > .empty-state { min-height: 100%; margin: 0; display: flex; @@ -1065,6 +1149,19 @@ background: rgba(255, 255, 255, 0.02); } +:root[data-theme="dark"] .providers-page .fetch-model-list-header { + border-bottom-color: #30363d; + background: rgba(19, 24, 32, 0.96); +} + +:root[data-theme="dark"] .providers-page .fetch-model-list-header-main { + color: #e5edf6; +} + +:root[data-theme="dark"] .providers-page .fetch-model-list-header-hint { + color: #9aa8ba; +} + :root[data-theme="dark"] .providers-page .fetch-model-item:hover { border-color: #30363d; background: rgba(255, 255, 255, 0.03); @@ -1166,8 +1263,14 @@ align-items: stretch; } - .providers-page .fetch-model-toolbar-actions { - justify-content: flex-end; + .providers-page .fetch-model-list-header { + align-items: flex-start; + flex-direction: column; + } + + .providers-page .fetch-model-list-header-hint { + padding-left: 26px; + white-space: normal; } .providers-page .fetch-model-item { diff --git a/src/presentation/templates/base_page.html b/src/presentation/templates/base_page.html index 38c2062..a175bff 100644 --- a/src/presentation/templates/base_page.html +++ b/src/presentation/templates/base_page.html @@ -6,7 +6,7 @@ {% include "_theme_bootstrap.html" %} {% block page_title %}LLM Proxy{% endblock %} - + {% block page_css %}{% endblock %} diff --git a/src/presentation/templates/index.html b/src/presentation/templates/index.html index ed6ca6f..5f65005 100644 --- a/src/presentation/templates/index.html +++ b/src/presentation/templates/index.html @@ -1,7 +1,7 @@ {% extends "base_admin.html" %} {% block page_title %}LLM使用统计{% endblock %} -{% block page_css %}{% endblock %} +{% block page_css %}{% endblock %} {% block body_class %}dashboard-page{% endblock %} {% block header_title %}LLM使用统计{% endblock %} {% block header_subtitle %}实时监控 API 请求和 Token 消耗{% endblock %} @@ -525,7 +525,12 @@ updateSortIndicators('stats', statsSortState, ['ip_address', 'username', 'request_model', 'response_model', 'prompt_tokens', 'completion_tokens', 'total_tokens', 'request_count']); if (statsData.length === 0) { - statsTable.innerHTML = '暂无统计数据'; + statsTable.innerHTML = ` + + + +

暂无统计数据

+ `; return; } diff --git a/src/presentation/templates/providers.html b/src/presentation/templates/providers.html index 9eb7540..120dc2a 100644 --- a/src/presentation/templates/providers.html +++ b/src/presentation/templates/providers.html @@ -1,7 +1,7 @@ {% extends "base_admin.html" %} {% block page_title %}Provider 管理{% endblock %} -{% block page_css %}{% endblock %} +{% block page_css %}{% endblock %} {% block body_class %}providers-page{% endblock %} {% block header_title %}Provider 管理{% endblock %} {% block header_subtitle %}维护 providers 配置，修改后自动重载生效。{% endblock %} @@ -57,6 +57,13 @@ +
+ + +
@@ -78,9 +85,8 @@
- + @@ -100,7 +106,7 @@ class="field-help-button" data-bs-toggle="tooltip" data-bs-placement="top" - data-bs-title="从 /v1/models 或 /models 接口拉取模型，勾选后再加入当前列表。" + data-bs-title="从 /v1/models 或 /models 接口拉取模型；如果当前 API 是 ws/wss，会自动尝试对应的 HTTP/HTTPS 地址。" aria-label="查看拉取模型说明" >?
@@ -122,21 +128,29 @@
选择拉取模型
-
已根据当前 Model List 自动预选重复模型，取消勾选可将它们从当前清单移除。
+
根据当前 Model List 自动预选重复模型，取消勾选可将它们从当前清单移除，未出现在本次拉取结果中的现有模型会保留。
- -
- - -
+
共拉取 0 个模型，已选 0 个
-
未出现在本次拉取结果中的现有模型会保留。
-
+
+
+ +
+
+
@@ -400,9 +414,16 @@ summaryElement.textContent = '共拉取 0 个模型，已选 0 个'; } - const listElement = document.getElementById('fetchModelList'); - if (listElement) { - listElement.innerHTML = ''; + const selectAllCheckbox = document.getElementById('fetchModelSelectAllCheckbox'); + if (selectAllCheckbox) { + selectAllCheckbox.checked = false; + selectAllCheckbox.indeterminate = false; + selectAllCheckbox.disabled = true; + } + + const listBodyElement = document.getElementById('fetchModelListBody'); + if (listBodyElement) { + listBodyElement.innerHTML = ''; } } @@ -425,19 +446,32 @@ summaryElement.textContent = `共拉取 ${totalCount} 个模型${visibleText}，已选 ${selectedCount} 个${existingText}`; } + function syncFetchModelSelectAllState(filteredModels = getFilteredFetchedModels()) { + const selectAllCheckbox = document.getElementById('fetchModelSelectAllCheckbox'); + if (!selectAllCheckbox) return; + + const visibleCount = filteredModels.length; + const selectedVisibleCount = filteredModels.filter(model => fetchedModelSelection.has(model)).length; + + selectAllCheckbox.disabled = visibleCount === 0; + selectAllCheckbox.indeterminate = visibleCount > 0 && selectedVisibleCount > 0 && selectedVisibleCount < visibleCount; + selectAllCheckbox.checked = visibleCount > 0 && selectedVisibleCount === visibleCount; + } + function renderFetchModelPicker() { - const listElement = document.getElementById('fetchModelList'); - if (!listElement) return; + const listBodyElement = document.getElementById('fetchModelListBody'); + if (!listBodyElement) return; const filteredModels = getFilteredFetchedModels(); updateFetchModelSummary(filteredModels.length); + syncFetchModelSelectAllState(filteredModels); if (!filteredModels.length) { - listElement.innerHTML = '
没有匹配的拉取模型
'; + listBodyElement.innerHTML = '
没有匹配的拉取模型
'; return; } - listElement.innerHTML = filteredModels + listBodyElement.innerHTML = filteredModels .map(model => { const encodedModel = encodeURIComponent(model); const isSelected = fetchedModelSelection.has(model); @@ -492,16 +526,19 @@ itemElement.classList.toggle('selected', checked); } - updateFetchModelSummary(getFilteredFetchedModels().length); + const filteredModels = getFilteredFetchedModels(); + updateFetchModelSummary(filteredModels.length); + syncFetchModelSelectAllState(filteredModels); } - function selectAllFetchedModels() { - fetchedModelSelection = new Set(fetchedModelCandidates); - renderFetchModelPicker(); - } - - function clearFetchedModelsSelection() { - fetchedModelSelection = new Set(); + function toggleFilteredFetchedModels(checked) { + getFilteredFetchedModels().forEach(model => { + if (checked) { + fetchedModelSelection.add(model); + } else { + fetchedModelSelection.delete(model); + } + }); renderFetchModelPicker(); } @@ -612,6 +649,7 @@ return { name: document.getElementById('providerName').value.trim(), api: document.getElementById('providerApi').value.trim(), + transport: document.getElementById('providerTransport').value, api_key: document.getElementById('providerApiKey').value.trim(), proxy: document.getElementById('providerProxy').value.trim(), timeout_seconds: document.getElementById('providerTimeout').value.trim(), @@ -623,19 +661,48 @@ } function fillForm(provider) { + const inferredTransport = inferTransportFromApi(provider?.api || ''); + const resolvedTransport = provider?.transport || inferredTransport; document.getElementById('providerName').value = provider?.name || ''; document.getElementById('providerApi').value = provider?.api || ''; + document.getElementById('providerTransport').value = resolvedTransport; document.getElementById('providerApiKey').value = provider?.api_key || ''; document.getElementById('providerProxy').value = provider?.proxy || ''; document.getElementById('providerTimeout').value = provider?.timeout_seconds || ''; document.getElementById('providerRetries').value = provider?.max_retries || ''; - document.getElementById('providerVerifySsl').value = provider?.verify_ssl === true ? 'true' : (provider?.verify_ssl === false ? 'false' : ''); + document.getElementById('providerVerifySsl').value = provider?.verify_ssl === true ? 'true' : 'false'; document.getElementById('providerModelList').value = Array.isArray(provider?.model_list) ? buildModelListValue(provider.model_list) : ''; document.getElementById('providerHook').value = provider?.hook || ''; updateModelListCount(); updateFetchModelsButtonState(); } + function inferTransportFromApi(apiValue) { + const normalized = String(apiValue ?? '').trim().toLowerCase(); + if (normalized.startsWith('ws://') || normalized.startsWith('wss://')) { + return 'websocket'; + } + return 'http'; + } + + function getTransportDisplayMeta(transportValue) { + const normalized = String(transportValue ?? '').trim().toLowerCase(); + if (normalized === 'websocket') { + return { + label: 'WebSocket', + className: 'is-websocket' + }; + } + return { + label: 'HTTP', + className: 'is-http' + }; + } + + function getEffectiveTransport() { + return document.getElementById('providerTransport')?.value || inferTransportFromApi(document.getElementById('providerApi')?.value || ''); + } + function initTooltips() { if (!window.bootstrap || !window.bootstrap.Tooltip) return; document.querySelectorAll('[data-bs-toggle="tooltip"]').forEach(element => { @@ -717,6 +784,7 @@ const models = Array.isArray(provider.model_list) ? provider.model_list : []; const preview = getModelPreview(models); const isDeletePending = pendingDeleteProviderName === provider.name; + const transportMeta = getTransportDisplayMeta(provider.transport || inferTransportFromApi(provider.api || '')); const modelTags = preview.visible.length ? preview.visible.map(model => `${escapeHtml(model)}`).join('') : '未配置'; @@ -732,7 +800,12 @@ html += ` - ${escapeHtml(provider.name)} + +
+ ${escapeHtml(provider.name)} + ${transportMeta.label} +
+ ${escapeHtml(provider.api)}
${modelTags}
@@ -933,6 +1006,7 @@ }); document.getElementById('providerName').addEventListener('input', updateFetchModelsButtonState); document.getElementById('providerApi').addEventListener('input', updateFetchModelsButtonState); + document.getElementById('providerTransport').addEventListener('change', updateFetchModelsButtonState); document.getElementById('providerModelList').addEventListener('input', updateModelListCount); document.getElementById('providerModelList').addEventListener('blur', function() { applyNormalizedModelListValue(); diff --git a/src/services/model_discovery_service.py b/src/services/model_discovery_service.py index cccdf9a..8716817 100644 --- a/src/services/model_discovery_service.py +++ b/src/services/model_discovery_service.py @@ -99,7 +99,13 @@ class ModelDiscoveryService: if not parsed.scheme or not parsed.netloc: raise ValueError('Provider api must be a valid absolute URL') - root = f'{parsed.scheme}://{parsed.netloc}' + normalized_scheme = parsed.scheme.lower() + if normalized_scheme == 'ws': + normalized_scheme = 'http' + elif normalized_scheme == 'wss': + normalized_scheme = 'https' + + root = f'{normalized_scheme}://{parsed.netloc}' path = parsed.path.rstrip('/') path_prefixes = [''] diff --git a/src/services/provider_service.py b/src/services/provider_service.py index 56dd024..eee3cd3 100644 --- a/src/services/provider_service.py +++ b/src/services/provider_service.py @@ -23,12 +23,18 @@ class ProviderService: def list_providers(self) -> List[Dict[str, Any]]: config = self._config_manager.get_raw_config() - return self._extract_providers(config) + return [ + ProviderConfigSchema.from_mapping(provider).to_mapping() + for provider in self._extract_providers(config) + ] def get_provider(self, name: str) -> Optional[Dict[str, Any]]: config = self._config_manager.get_raw_config() providers = self._extract_providers(config) - return self._find_provider(providers, name) + provider = self._find_provider(providers, name) + if provider is None: + return None + return ProviderConfigSchema.from_mapping(provider).to_mapping() def create_provider(self, payload: Dict[str, Any]) -> Dict[str, Any]: config = self._config_manager.get_raw_config() diff --git a/src/services/proxy_service.py b/src/services/proxy_service.py index 654c7ac..e4b5b94 100644 --- a/src/services/proxy_service.py +++ b/src/services/proxy_service.py @@ -2,18 +2,29 @@ # -*- coding: utf-8 -*- """上游 LLM 代理服务。""" +import json import threading from typing import Any, Callable, Dict, Optional, Tuple +from urllib.parse import urlparse import requests +import websocket from flask import Response, stream_with_context from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry from ..application.app_context import AppContext -from ..external import LLMProvider, build_proxy_response, probe_stream_response +from ..external import ( + LLMProvider, + StaticUpstreamResponse, + WebSocketUpstreamResponse, + build_proxy_response, + collect_websocket_response_body, + probe_stream_response, +) +from ..external.stream_probe import BufferedUpstreamResponse from ..hooks import HookContext -from ..utils.net import build_requests_proxies +from ..utils.net import build_requests_proxies, build_websocket_connect_options class ProxyService: @@ -60,8 +71,9 @@ class ProxyService: headers, body, request_ctx = build_request(attempt) requested_stream = bool(body.get("stream", False)) self._logger.info( - "Proxying upstream request: provider=%s model=%s attempt=%s/%s stream=%s", + "Proxying upstream request: provider=%s transport=%s model=%s attempt=%s/%s stream=%s", provider.name, + provider.transport, body.get("model"), attempt + 1, max_retries, @@ -69,50 +81,36 @@ class ProxyService: ) try: - http_session = self._get_http_session() - upstream_response = http_session.post( + response_for_hook, is_stream, status_code = self._open_upstream_response( + provider, + headers, + body, + requested_stream, target_url, - headers=headers, - json=body, - stream=requested_stream, - proxies=request_proxies, - verify=verify_ssl, - timeout=timeout_seconds, + request_proxies, + timeout_seconds, + verify_ssl, ) - status_code = upstream_response.status_code should_retry_status = self._should_retry_status_code(status_code) if should_retry_status and attempt < max_retries - 1: self._logger.warning( - "Retryable upstream status (attempt %s/%s): provider=%s status=%s, retrying", + "Retryable upstream status (attempt %s/%s): provider=%s transport=%s status=%s, retrying", attempt + 1, max_retries, provider.name, + provider.transport, status_code, ) - upstream_response.close() + response_for_hook.close() continue - content_type = (upstream_response.headers.get("Content-Type") or "").lower() - upstream_stream = "text/event-stream" in content_type - response_for_hook: Any = upstream_response - if requested_stream and not upstream_stream: - response_for_hook, is_stream = probe_stream_response(upstream_response) + if requested_stream != is_stream: self._logger.warning( - "Stream mode probed for mismatch: requested_stream=%s upstream_stream=%s probed_stream=%s content_type=%s", + "Stream mode mismatch: requested_stream=%s upstream_stream=%s transport=%s", requested_stream, - upstream_stream, is_stream, - content_type, - ) - else: - is_stream = upstream_stream - - if requested_stream != upstream_stream: - self._logger.warning( - "Stream mode mismatch: requested_stream=%s, upstream_stream=%s", - requested_stream, - upstream_stream, + provider.transport, ) response = build_proxy_response( @@ -127,23 +125,29 @@ class ProxyService: forward_stream_usage=forward_stream_usage, ) self._logger.info( - "Upstream request completed: provider=%s status=%s stream=%s", + "Upstream request completed: provider=%s transport=%s status=%s stream=%s", provider.name, + provider.transport, status_code, is_stream, ) if is_stream: - - def _close_upstream_resources() -> None: - response_for_hook.close() - - response.call_on_close(_close_upstream_resources) + response.call_on_close(response_for_hook.close) else: response_for_hook.close() return response, status_code except requests.exceptions.RequestException as exc: self._logger.error( - "Request error (attempt %s/%s): %s", + "HTTP upstream request error (attempt %s/%s): %s", + attempt + 1, + max_retries, + exc, + ) + if attempt < max_retries - 1: + continue + except (websocket.WebSocketException, OSError) as exc: + self._logger.error( + "WebSocket upstream request error (attempt %s/%s): %s", attempt + 1, max_retries, exc, @@ -153,10 +157,153 @@ class ProxyService: return None, 502 + def _open_upstream_response( + self, + provider: LLMProvider, + headers: Dict[str, str], + body: Dict[str, Any], + requested_stream: bool, + target_url: str, + request_proxies: Optional[Dict[str, str]], + timeout_seconds: int, + verify_ssl: bool, + ) -> Tuple[Any, bool, int]: + if provider.transport == "websocket": + return self._open_websocket_upstream_response( + provider, + headers, + body, + requested_stream, + timeout_seconds, + verify_ssl, + ) + + return self._open_http_upstream_response( + provider, + headers, + body, + requested_stream, + target_url, + request_proxies, + timeout_seconds, + verify_ssl, + ) + + def _open_http_upstream_response( + self, + provider: LLMProvider, + headers: Dict[str, str], + body: Dict[str, Any], + requested_stream: bool, + target_url: str, + request_proxies: Optional[Dict[str, str]], + timeout_seconds: int, + verify_ssl: bool, + ) -> Tuple[Any, bool, int]: + http_session = self._get_http_session() + upstream_response = http_session.post( + target_url, + headers=headers, + json=body, + stream=requested_stream, + proxies=request_proxies, + verify=verify_ssl, + timeout=timeout_seconds, + ) + + status_code = upstream_response.status_code + content_type = (upstream_response.headers.get("Content-Type") or "").lower() + upstream_stream = "text/event-stream" in content_type + response_for_hook: Any = upstream_response + if requested_stream and not upstream_stream: + response_for_hook, is_stream = probe_stream_response(upstream_response) + self._logger.warning( + "Stream mode probed for mismatch: requested_stream=%s upstream_stream=%s probed_stream=%s content_type=%s", + requested_stream, + upstream_stream, + is_stream, + content_type, + ) + else: + is_stream = upstream_stream + + return response_for_hook, is_stream, status_code + + def _open_websocket_upstream_response( + self, + provider: LLMProvider, + headers: Dict[str, str], + body: Dict[str, Any], + requested_stream: bool, + timeout_seconds: int, + verify_ssl: bool, + ) -> Tuple[Any, bool, int]: + websocket_url = self._get_upstream_websocket_url(provider.api) + handshake_headers = self._build_websocket_handshake_headers(headers) + + try: + connection = websocket.create_connection( + websocket_url, + timeout=timeout_seconds, + header=handshake_headers, + **build_websocket_connect_options(provider.proxy, verify_ssl), + ) + except websocket.WebSocketBadStatusException as exc: + status_code = int(getattr(exc, "status_code", 502) or 502) + body_bytes = self._build_websocket_error_body(exc) + response = BufferedUpstreamResponse( + StaticUpstreamResponse( + status_code=status_code, + headers={"Content-Type": "application/json"}, + ), + body_bytes, + ) + return response, False, status_code + + connection.send(json.dumps(body, ensure_ascii=False)) + + if requested_stream: + response = WebSocketUpstreamResponse(connection) + return response, True, response.status_code + + response_body = collect_websocket_response_body(connection, self._logger) + response = BufferedUpstreamResponse( + StaticUpstreamResponse( + status_code=200, + headers={"Content-Type": "application/json"}, + on_close=connection.close, + ), + response_body, + ) + return response, False, response.status_code + @staticmethod - def _should_retry_status_code(status_code: int) -> bool: - retryable_status_codes = {408, 409, 425, 429, 500, 502, 503, 504} - return status_code in retryable_status_codes + def _build_websocket_handshake_headers(headers: Dict[str, str]) -> Dict[str, str]: + excluded = { + "accept", + "content-length", + "content-type", + "connection", + "upgrade", + } + return { + key: value + for key, value in headers.items() + if key.lower() not in excluded + } + + @staticmethod + def _build_websocket_error_body(exc: websocket.WebSocketBadStatusException) -> bytes: + response_body = getattr(exc, "resp_body", None) + if isinstance(response_body, bytes) and response_body.strip(): + return response_body + if isinstance(response_body, str) and response_body.strip(): + return response_body.encode("utf-8") + + payload = { + "error": f"WebSocket upstream handshake failed: {exc}", + } + return json.dumps(payload, ensure_ascii=False).encode("utf-8") @staticmethod def _get_upstream_model_name(provider_name: str, requested_model_name: str) -> str: @@ -165,6 +312,22 @@ class ProxyService: return requested_model_name[len(prefix):] return requested_model_name + @staticmethod + def _get_upstream_websocket_url(api: str) -> str: + parsed = urlparse(api.strip()) + if parsed.scheme.lower() in {"ws", "wss"}: + return parsed.geturl() + if parsed.scheme.lower() == "http": + return parsed._replace(scheme="ws").geturl() + if parsed.scheme.lower() == "https": + return parsed._replace(scheme="wss").geturl() + raise ValueError("WebSocket upstream api must use ws://, wss://, http:// or https://") + + @staticmethod + def _should_retry_status_code(status_code: int) -> bool: + retryable_status_codes = {408, 409, 425, 429, 500, 502, 503, 504} + return status_code in retryable_status_codes + def _get_http_session(self) -> requests.Session: session = getattr(self._http_local, "session", None) if session is None: @@ -197,4 +360,4 @@ class ProxyService: "upgrade", "set-cookie", } - return {k: v for k, v in headers.items() if k.lower() not in excluded} + return {key: value for key, value in headers.items() if key.lower() not in excluded} diff --git a/src/utils/net.py b/src/utils/net.py index d3b27af..cac673a 100644 --- a/src/utils/net.py +++ b/src/utils/net.py @@ -3,7 +3,8 @@ """网络相关工具函数。""" import ipaddress -from typing import Dict, Optional +import ssl +from typing import Any, Dict, Optional from urllib.parse import urlparse @@ -59,3 +60,32 @@ def build_requests_proxies(proxy_value: Optional[str]) -> Optional[Dict[str, str "http": normalized, "https": normalized, } + + +def build_websocket_connect_options( + proxy_value: Optional[str], + verify_ssl: bool, +) -> Dict[str, Any]: + """构造 websocket-client 连接参数。""" + options: Dict[str, Any] = { + "enable_multithread": True, + "sslopt": { + "cert_reqs": ssl.CERT_REQUIRED if verify_ssl else ssl.CERT_NONE, + }, + } + + normalized = normalize_proxy_url(proxy_value) + if normalized is None: + return options + + parsed = urlparse(normalized) + scheme = parsed.scheme.lower() + if scheme not in {"http", "https"}: + raise ValueError("WebSocket upstream currently only supports http/https proxy") + + options["http_proxy_host"] = parsed.hostname + options["http_proxy_port"] = parsed.port or (443 if scheme == "https" else 80) + if parsed.username: + options["http_proxy_auth"] = (parsed.username, parsed.password or "") + + return options diff --git a/tests/test_proxy_transport.py b/tests/test_proxy_transport.py new file mode 100644 index 0000000..5641977 --- /dev/null +++ b/tests/test_proxy_transport.py @@ -0,0 +1,173 @@ +import ssl +import sys +import unittest +from pathlib import Path + +from websocket import ABNF + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from src.config.provider_config import ProviderConfigSchema +from src.external.upstream_websocket import ( + WebSocketUpstreamResponse, + collect_websocket_response_body, + normalize_websocket_message, +) +from src.services.model_discovery_service import ModelDiscoveryService +from src.utils.net import build_websocket_connect_options + + +class FakeWebSocketConnection: + def __init__(self, frames): + self._frames = iter(frames) + self.closed = False + self.pongs = [] + + def recv_data(self, control_frame=True): + return next(self._frames) + + def pong(self, payload): + self.pongs.append(payload) + + def close(self): + self.closed = True + + +class ProviderTransportTests(unittest.TestCase): + def test_provider_transport_defaults_to_websocket_for_ws_scheme(self) -> None: + schema = ProviderConfigSchema.from_mapping( + { + "name": "codex", + "api": "wss://example.com/v1/chat/completions", + "model_list": ["gpt-4.1"], + } + ) + + self.assertEqual("websocket", schema.transport) + + def test_provider_transport_allows_explicit_override(self) -> None: + schema = ProviderConfigSchema.from_mapping( + { + "name": "codex", + "api": "https://example.com/v1/chat/completions", + "transport": "websocket", + "model_list": ["gpt-4.1"], + } + ) + + self.assertEqual("websocket", schema.transport) + + def test_provider_transport_rejects_http_transport_with_ws_scheme(self) -> None: + with self.assertRaisesRegex(ValueError, "requires api to use http:// or https://"): + ProviderConfigSchema.from_mapping( + { + "name": "bad-provider", + "api": "wss://example.com/v1/chat/completions", + "transport": "http", + "model_list": ["demo"], + } + ) + + +class WebSocketProxyBridgeTests(unittest.TestCase): + def test_normalize_websocket_message_wraps_json_as_sse(self) -> None: + chunk = normalize_websocket_message(b'{"id":"evt_1"}') + + self.assertEqual(b'data: {"id":"evt_1"}\n\n', chunk) + + def test_stream_response_converts_websocket_messages_to_sse(self) -> None: + connection = FakeWebSocketConnection( + [ + (ABNF.OPCODE_TEXT, b'{"id":"evt_1"}'), + (ABNF.OPCODE_TEXT, b"data: [DONE]\n\n"), + (ABNF.OPCODE_CLOSE, b""), + ] + ) + response = WebSocketUpstreamResponse(connection) + + chunks = list(response.iter_content()) + + self.assertEqual( + [ + b'data: {"id":"evt_1"}\n\n', + b"data: [DONE]\n\n", + ], + chunks, + ) + response.close() + self.assertTrue(connection.closed) + + def test_collect_non_stream_websocket_body_uses_terminal_payload(self) -> None: + connection = FakeWebSocketConnection( + [ + (ABNF.OPCODE_TEXT, b'{"delta":"hello"}'), + ( + ABNF.OPCODE_TEXT, + b'{"choices":[],"usage":{"total_tokens":1},"model":"demo"}', + ), + (ABNF.OPCODE_CLOSE, b""), + ] + ) + + body = collect_websocket_response_body(connection) + + self.assertEqual( + b'{"choices":[],"usage":{"total_tokens":1},"model":"demo"}', + body, + ) + + def test_collect_non_stream_websocket_body_extracts_sse_data(self) -> None: + connection = FakeWebSocketConnection( + [ + (ABNF.OPCODE_TEXT, b"data: {\"id\":\"evt_1\"}\n\ndata: [DONE]\n\n"), + (ABNF.OPCODE_CLOSE, b""), + ] + ) + + body = collect_websocket_response_body(connection) + + self.assertEqual(b'{"id":"evt_1"}', body) + + +class WebSocketConnectOptionsTests(unittest.TestCase): + def test_build_websocket_connect_options_supports_http_proxy(self) -> None: + options = build_websocket_connect_options("http://user:pass@proxy.local:8080", False) + + self.assertEqual("proxy.local", options["http_proxy_host"]) + self.assertEqual(8080, options["http_proxy_port"]) + self.assertEqual(("user", "pass"), options["http_proxy_auth"]) + self.assertEqual(ssl.CERT_NONE, options["sslopt"]["cert_reqs"]) + + +class ModelDiscoveryCandidateTests(unittest.TestCase): + def test_model_discovery_maps_websocket_scheme_to_https(self) -> None: + candidates = ModelDiscoveryService._build_model_endpoint_candidates( + "wss://example.com/v1/chat/completions" + ) + + self.assertIn("https://example.com/v1/models", candidates) + self.assertIn("https://example.com/models", candidates) + self.assertNotIn("wss://example.com/v1/models", candidates) + + +class ProviderTemplateTransportTests(unittest.TestCase): + def test_provider_template_contains_transport_field(self) -> None: + template_path = Path(__file__).resolve().parents[1] / "src" / "presentation" / "templates" / "providers.html" + html = template_path.read_text(encoding="utf-8") + + self.assertIn('id="providerTransport"', html) + self.assertIn('class="field-help-button"', html) + self.assertIn('data-bs-toggle="tooltip"', html) + self.assertIn('id="fetchModelSelectAllCheckbox"', html) + self.assertIn('toggleFilteredFetchedModels(this.checked)', html) + self.assertNotIn('selectAllFetchedModels()', html) + self.assertNotIn('clearFetchedModelsSelection()', html) + self.assertNotIn("自动（按 API 推断）", html) + self.assertIn("证书校验（HTTPS/WSS）", html) + self.assertNotIn('', html) + self.assertIn("WebSocket", html) + self.assertIn("provider-transport-badge", html) + + +if __name__ == "__main__": + unittest.main()