From e151aa812396f509a2fdb7e04bc7a5f0514c1bb1 Mon Sep 17 00:00:00 2001 From: Von Date: Wed, 11 Mar 2026 00:45:25 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9external?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 1 + src/external/__init__.py | 6 +- src/external/llm_provider.py | 173 ++----------------------------- src/external/response_adapter.py | 167 +++++++++++++++++++++++++++++ src/external/stream_probe.py | 60 +++++++++++ src/services/proxy_service.py | 73 ++----------- 6 files changed, 248 insertions(+), 232 deletions(-) create mode 100644 src/external/response_adapter.py create mode 100644 src/external/stream_probe.py diff --git a/.gitignore b/.gitignore index ff220e2..b2ab7d1 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ logs .history config.yaml .record + diff --git a/src/external/__init__.py b/src/external/__init__.py index 36cf9c1..8972224 100644 --- a/src/external/__init__.py +++ b/src/external/__init__.py @@ -1,7 +1,9 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -"""外部服务层导出。""" +"""external 集成层导出。""" from .llm_provider import LLMProvider +from .response_adapter import build_proxy_response +from .stream_probe import probe_stream_response -__all__ = ['LLMProvider'] +__all__ = ["LLMProvider", "build_proxy_response", "probe_stream_response"] diff --git a/src/external/llm_provider.py b/src/external/llm_provider.py index eb9651f..71f8626 100644 --- a/src/external/llm_provider.py +++ b/src/external/llm_provider.py @@ -1,18 +1,14 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -"""外部 LLM Provider 抽象。""" +"""LLM Provider 配置与请求阶段钩子封装。""" -import json -from typing import Any, Callable, Dict, Iterator, Optional +from typing import Any, Dict, Optional -from flask import Response - -from ..application.app_context import Logger from ..hooks import HookContext, HookModule class LLMProvider: - """封装 provider 配置与钩子处理流程。""" + """封装 provider 配置与请求阶段钩子调用。""" def __init__( self, @@ -35,168 +31,11 @@ class LLMProvider: self.hook = hook def apply_header_hook(self, ctx: HookContext, headers: Dict[str, str]) -> Dict[str, str]: - if self.hook and hasattr(self.hook, 'header_hook'): + if self.hook and hasattr(self.hook, "header_hook"): return self.hook.header_hook(ctx, headers) return headers def apply_input_body_hook(self, ctx: HookContext, body: Dict[str, Any]) -> Dict[str, Any]: - if self.hook and hasattr(self.hook, 'input_body_hook'): + if self.hook and hasattr(self.hook, "input_body_hook"): return self.hook.input_body_hook(ctx, body) - return body - - def apply_output_body_hook( - self, - ctx: HookContext, - response: Any, - is_stream: bool, - filter_headers_func: Any, - stream_context_func: Any, - logger: Logger, - on_complete: Optional[Callable[[Dict[str, Any]], None]] = None, - ) -> Response: - hook_func = getattr(self.hook, 'output_body_hook', None) if self.hook else None - meta: Dict[str, Any] = { - 'response_model': None, - 'total_tokens': 0, - 'prompt_tokens': 0, - 'completion_tokens': 0, - } - completed = False - - def update_meta_from_payload(payload: Dict[str, Any]) -> None: - model = payload.get('model') - if model: - meta['response_model'] = model - usage = payload.get('usage') - if isinstance(usage, dict): - if 'total_tokens' in usage and usage.get('total_tokens') is not None: - meta['total_tokens'] = int(usage.get('total_tokens') or 0) - if 'prompt_tokens' in usage and usage.get('prompt_tokens') is not None: - meta['prompt_tokens'] = int(usage.get('prompt_tokens') or 0) - if 'completion_tokens' in usage and usage.get('completion_tokens') is not None: - meta['completion_tokens'] = int(usage.get('completion_tokens') or 0) - - def safe_on_complete() -> None: - nonlocal completed - if completed: - return - completed = True - if on_complete: - try: - on_complete(meta) - except Exception as exc: - logger.error(f'Error in on_complete callback: {exc}') - - def parse_sse_line(line: str) -> Optional[str]: - """解析单行 SSE data 内容。""" - stripped = line.strip() - if not stripped or not stripped.startswith('data:'): - return None - return stripped[5:].strip() - - def process_sse_event(event_text: str) -> Iterator[bytes]: - """处理单个 SSE 事件块并执行输出钩子。""" - data_lines = [] - for raw_line in event_text.splitlines(): - if raw_line.startswith('data:'): - data_lines.append(raw_line[5:].strip()) - - if not data_lines: - if event_text.strip(): - yield (event_text + '\n\n').encode('utf-8') - return - - data_str = '\n'.join(data_lines).strip() - if not data_str: - return - - if data_str == '[DONE]': - yield b'data: [DONE]\n\n' - return - - try: - data = json.loads(data_str) - except json.JSONDecodeError: - yield f'data: {data_str}\n\n'.encode('utf-8') - return - - if isinstance(data, dict): - update_meta_from_payload(data) - - if hook_func: - modified = hook_func(ctx, data) - if modified is not None: - yield f"data: {json.dumps(modified)}\n\n".encode('utf-8') - else: - yield f'data: {data_str}\n\n'.encode('utf-8') - else: - yield f'data: {data_str}\n\n'.encode('utf-8') - - if not is_stream: - body = response.content - response_body = body - try: - payload = json.loads(body.decode('utf-8')) - if isinstance(payload, dict): - update_meta_from_payload(payload) - if hook_func: - modified = hook_func(ctx, payload) - response_body = json.dumps(modified).encode('utf-8') - except json.JSONDecodeError as exc: - logger.warning( - 'Non-stream response is not valid JSON, skip output_body_hook: status=%s content_type=%s error=%s', - response.status_code, - response.headers.get('Content-Type'), - exc, - ) - finally: - safe_on_complete() - - return Response( - response_body, - status=response.status_code, - headers=filter_headers_func(response.headers), - mimetype=response.headers.get('Content-Type', 'application/json'), - ) - - def generate() -> Iterator[bytes]: - buffer = '' - try: - for chunk in response.iter_content(chunk_size=None): - if not chunk: - continue - - try: - buffer += chunk.decode('utf-8') - except UnicodeDecodeError: - yield chunk - continue - - while '\n\n' in buffer: - event_text, buffer = buffer.split('\n\n', 1) - yield from process_sse_event(event_text) - - if buffer: - # 对不完整流做兜底：尽量解析已完整接收的 data 行。 - for line in buffer.splitlines(): - data_str = parse_sse_line(line) - if data_str is None: - continue - if data_str == '[DONE]': - yield b'data: [DONE]\n\n' - continue - try: - data = json.loads(data_str) - if isinstance(data, dict): - update_meta_from_payload(data) - except Exception: - continue - finally: - safe_on_complete() - - return Response( - stream_context_func(generate()), - status=response.status_code, - headers=filter_headers_func(response.headers), - mimetype=response.headers.get('Content-Type', 'text/event-stream'), - ) + return body diff --git a/src/external/response_adapter.py b/src/external/response_adapter.py new file mode 100644 index 0000000..82d7163 --- /dev/null +++ b/src/external/response_adapter.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""上游响应适配：执行 provider 输出钩子并构造 Flask 响应。""" + +import json +from typing import Any, Callable, Dict, Iterator, Optional + +from flask import Response + +from ..application.app_context import Logger +from ..hooks import HookContext, HookModule + + +def build_proxy_response( + hook: Optional[HookModule], + ctx: HookContext, + response: Any, + is_stream: bool, + filter_headers_func: Any, + stream_context_func: Any, + logger: Logger, + on_complete: Optional[Callable[[Dict[str, Any]], None]] = None, +) -> Response: + """基于上游响应与输出钩子构建 Flask 响应。""" + hook_func = getattr(hook, "output_body_hook", None) if hook else None + meta: Dict[str, Any] = { + "response_model": None, + "total_tokens": 0, + "prompt_tokens": 0, + "completion_tokens": 0, + } + completed = False + + def update_meta_from_payload(payload: Dict[str, Any]) -> None: + model = payload.get("model") + if model: + meta["response_model"] = model + usage = payload.get("usage") + if isinstance(usage, dict): + if "total_tokens" in usage and usage.get("total_tokens") is not None: + meta["total_tokens"] = int(usage.get("total_tokens") or 0) + if "prompt_tokens" in usage and usage.get("prompt_tokens") is not None: + meta["prompt_tokens"] = int(usage.get("prompt_tokens") or 0) + if "completion_tokens" in usage and usage.get("completion_tokens") is not None: + meta["completion_tokens"] = int(usage.get("completion_tokens") or 0) + + def safe_on_complete() -> None: + nonlocal completed + if completed: + return + completed = True + if on_complete: + try: + on_complete(meta) + except Exception as exc: + logger.error(f"Error in on_complete callback: {exc}") + + def parse_sse_line(line: str) -> Optional[str]: + stripped = line.strip() + if not stripped or not stripped.startswith("data:"): + return None + return stripped[5:].strip() + + def process_sse_event(event_text: str) -> Iterator[bytes]: + data_lines = [] + for raw_line in event_text.splitlines(): + if raw_line.startswith("data:"): + data_lines.append(raw_line[5:].strip()) + + if not data_lines: + if event_text.strip(): + yield (event_text + "\n\n").encode("utf-8") + return + + data_str = "\n".join(data_lines).strip() + if not data_str: + return + + if data_str == "[DONE]": + yield b"data: [DONE]\n\n" + return + + try: + data = json.loads(data_str) + except json.JSONDecodeError: + yield f"data: {data_str}\n\n".encode("utf-8") + return + + if isinstance(data, dict): + update_meta_from_payload(data) + + if hook_func: + modified = hook_func(ctx, data) + if modified is not None: + yield f"data: {json.dumps(modified)}\n\n".encode("utf-8") + else: + yield f"data: {data_str}\n\n".encode("utf-8") + else: + yield f"data: {data_str}\n\n".encode("utf-8") + + if not is_stream: + body = response.content + response_body = body + try: + payload = json.loads(body.decode("utf-8")) + if isinstance(payload, dict): + update_meta_from_payload(payload) + if hook_func: + modified = hook_func(ctx, payload) + response_body = json.dumps(modified).encode("utf-8") + except json.JSONDecodeError as exc: + logger.warning( + "Non-stream response is not valid JSON, skip output_body_hook: status=%s content_type=%s error=%s", + response.status_code, + response.headers.get("Content-Type"), + exc, + ) + finally: + safe_on_complete() + + return Response( + response_body, + status=response.status_code, + headers=filter_headers_func(response.headers), + mimetype=response.headers.get("Content-Type", "application/json"), + ) + + def generate() -> Iterator[bytes]: + buffer = "" + try: + for chunk in response.iter_content(chunk_size=None): + if not chunk: + continue + + try: + buffer += chunk.decode("utf-8") + except UnicodeDecodeError: + yield chunk + continue + + while "\n\n" in buffer: + event_text, buffer = buffer.split("\n\n", 1) + yield from process_sse_event(event_text) + + if buffer: + for line in buffer.splitlines(): + data_str = parse_sse_line(line) + if data_str is None: + continue + if data_str == "[DONE]": + yield b"data: [DONE]\n\n" + continue + try: + data = json.loads(data_str) + if isinstance(data, dict): + update_meta_from_payload(data) + except Exception: + continue + finally: + safe_on_complete() + + return Response( + stream_context_func(generate()), + status=response.status_code, + headers=filter_headers_func(response.headers), + mimetype=response.headers.get("Content-Type", "text/event-stream"), + ) diff --git a/src/external/stream_probe.py b/src/external/stream_probe.py new file mode 100644 index 0000000..b693af4 --- /dev/null +++ b/src/external/stream_probe.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""上游流式响应探测工具。""" + +from typing import Any, Iterator, Optional, Tuple + +import requests + + +class PrefetchedStreamResponse: + def __init__(self, response: requests.Response, first_chunk: bytes): + self._response = response + self._first_chunk = first_chunk + self.status_code = response.status_code + self.headers = response.headers + + def iter_content(self, chunk_size: Optional[int] = None) -> Iterator[bytes]: + if self._first_chunk: + yield self._first_chunk + self._first_chunk = b"" + yield from self._response.iter_content(chunk_size=chunk_size) + + def close(self) -> None: + self._response.close() + + +class BufferedUpstreamResponse: + def __init__(self, response: requests.Response, body: bytes): + self._response = response + self.content = body + self.status_code = response.status_code + self.headers = response.headers + + def close(self) -> None: + self._response.close() + + +def looks_like_sse_chunk(chunk: bytes) -> bool: + if not chunk: + return False + text = chunk.decode("utf-8", errors="ignore").lstrip() + return text.startswith("data:") or text.startswith("event:") or text.startswith(":") + + +def probe_stream_response(upstream_response: requests.Response) -> Tuple[Any, bool]: + chunk_iter = upstream_response.iter_content(chunk_size=None) + first_chunk = b"" + for chunk in chunk_iter: + if chunk: + first_chunk = chunk + break + + if not first_chunk: + return BufferedUpstreamResponse(upstream_response, b""), False + + if looks_like_sse_chunk(first_chunk): + return PrefetchedStreamResponse(upstream_response, first_chunk), True + + remaining = b"".join(chunk_iter) + return BufferedUpstreamResponse(upstream_response, first_chunk + remaining), False diff --git a/src/services/proxy_service.py b/src/services/proxy_service.py index a18c242..2fd1cb4 100644 --- a/src/services/proxy_service.py +++ b/src/services/proxy_service.py @@ -1,48 +1,20 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -"""代理服务。""" +"""上游 LLM 代理服务。""" import threading -from typing import Any, Callable, Dict, Iterator, Optional, Tuple +from typing import Any, Callable, Dict, Optional, Tuple import requests -from requests.adapters import HTTPAdapter 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 +from ..external import LLMProvider, build_proxy_response, probe_stream_response from ..hooks import HookContext -class _PrefetchedStreamResponse: - def __init__(self, response: requests.Response, first_chunk: bytes): - self._response = response - self._first_chunk = first_chunk - self.status_code = response.status_code - self.headers = response.headers - - def iter_content(self, chunk_size: Optional[int] = None) -> Iterator[bytes]: - if self._first_chunk: - yield self._first_chunk - self._first_chunk = b"" - yield from self._response.iter_content(chunk_size=chunk_size) - - def close(self) -> None: - self._response.close() - - -class _BufferedUpstreamResponse: - def __init__(self, response: requests.Response, body: bytes): - self._response = response - self.content = body - self.status_code = response.status_code - self.headers = response.headers - - def close(self) -> None: - self._response.close() - - class ProxyService: """处理上游 LLM 代理请求。""" @@ -58,7 +30,7 @@ class ProxyService: request_headers: Dict[str, str], on_complete: Optional[Callable[[Dict[str, Any]], None]] = None, ) -> Tuple[Optional[Response], int]: - """代理请求到指定 provider，并执行重试与输出钩子。""" + """代理请求到目标 provider，并处理重试与输出钩子。""" target_url = provider.api model_name = request_data["model"] @@ -67,7 +39,6 @@ class ProxyService: verify_ssl = provider.verify_ssl def build_request(attempt: int) -> Tuple[Dict[str, str], Dict[str, Any], HookContext]: - """构建重试请求的 headers/body/context。""" headers = dict(request_headers) headers["content-type"] = "application/json" @@ -85,7 +56,6 @@ class ProxyService: for attempt in range(max_retries): headers, body, request_ctx = build_request(attempt) requested_stream = bool(body.get("stream", False)) - http_session: Optional[requests.Session] = None self._logger.info( "Proxying upstream request: provider=%s model=%s attempt=%s/%s stream=%s", provider.name, @@ -123,7 +93,7 @@ class ProxyService: 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 = self._probe_stream_response(upstream_response) + 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, @@ -141,7 +111,8 @@ class ProxyService: upstream_stream, ) - response = provider.apply_output_body_hook( + response = build_proxy_response( + provider.hook, request_ctx, response_for_hook, is_stream, @@ -157,6 +128,7 @@ class ProxyService: is_stream, ) if is_stream: + def _close_upstream_resources() -> None: response_for_hook.close() @@ -200,33 +172,8 @@ class ProxyService: session.mount("https://", adapter) return session - @staticmethod - def _looks_like_sse_chunk(chunk: bytes) -> bool: - if not chunk: - return False - text = chunk.decode("utf-8", errors="ignore").lstrip() - return text.startswith("data:") or text.startswith("event:") or text.startswith(":") - - def _probe_stream_response(self, upstream_response: requests.Response) -> Tuple[Any, bool]: - chunk_iter = upstream_response.iter_content(chunk_size=None) - first_chunk = b"" - for chunk in chunk_iter: - if chunk: - first_chunk = chunk - break - - if not first_chunk: - return _BufferedUpstreamResponse(upstream_response, b""), False - - if self._looks_like_sse_chunk(first_chunk): - return _PrefetchedStreamResponse(upstream_response, first_chunk), True - - remaining = b"".join(chunk_iter) - return _BufferedUpstreamResponse(upstream_response, first_chunk + remaining), False - @staticmethod def _filter_response_headers(headers: Any) -> Dict[str, str]: - """过滤不应透传给客户端的 hop-by-hop 响应头。""" excluded = { "transfer-encoding", "connection", @@ -238,4 +185,4 @@ class ProxyService: "upgrade", "set-cookie", } - return {k: v for k, v in headers.items() if k.lower() not in excluded} + return {k: v for k, v in headers.items() if k.lower() not in excluded}