mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
26ede849e2
- cli_handler_base.py 拆分为 7 个 mixin (event/monitor/prefetch/request/sse_helpers/stream/sync) - usage/service.py 拆分为 6 个子模块 (types/active_requests/cache_analysis/lifecycle/pricing/query/recording) - models/database 拆分为独立模型文件 (auth/misc/model/provider/stats/usage/user) - DUMMY_THOUGHT_SIGNATURE 常量提升到 core/api_format/conversion/constants 统一管理 - task/service.py 内联导入提升为顶层导入 - 流处理函数签名移除冗余的 http_client 参数
124 lines
4.3 KiB
Python
124 lines
4.3 KiB
Python
import json
|
|
from typing import Any, AsyncIterator
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from src.api.handlers.base.response_parser import (
|
|
ParsedChunk,
|
|
ParsedResponse,
|
|
ResponseParser,
|
|
StreamStats,
|
|
)
|
|
from src.api.handlers.base.stream_context import StreamContext
|
|
from src.api.handlers.base.stream_processor import StreamProcessor
|
|
from src.utils.sse_parser import SSEEventParser
|
|
|
|
|
|
class DummyParser(ResponseParser):
|
|
def parse_sse_line(self, line: str, stats: StreamStats) -> ParsedChunk | None:
|
|
return None
|
|
|
|
def parse_response(self, response: dict[str, Any], status_code: int) -> ParsedResponse:
|
|
return ParsedResponse(raw_response=response, status_code=status_code)
|
|
|
|
def extract_usage_from_response(self, response: dict[str, Any]) -> dict[str, int]:
|
|
return {}
|
|
|
|
def extract_text_content(self, response: dict[str, Any]) -> str:
|
|
return ""
|
|
|
|
|
|
def test_process_line_strips_newlines_and_finalizes_event() -> None:
|
|
ctx = StreamContext(model="test-model", api_format="openai:chat")
|
|
processor = StreamProcessor(request_id="test-request", default_parser=DummyParser())
|
|
sse_parser = SSEEventParser()
|
|
|
|
processor._process_line(ctx, sse_parser, 'data: {"type":"response.completed"}\n')
|
|
processor._process_line(ctx, sse_parser, "\n")
|
|
|
|
assert ctx.has_completion is True
|
|
|
|
|
|
def test_process_line_updates_openai_usage_from_usage_only_chunk() -> None:
|
|
ctx = StreamContext(model="test-model", api_format="openai:chat")
|
|
ctx.provider_api_format = "openai:chat"
|
|
processor = StreamProcessor(request_id="test-request", default_parser=DummyParser())
|
|
sse_parser = SSEEventParser()
|
|
|
|
usage_chunk = {
|
|
"id": "chatcmpl_test",
|
|
"object": "chat.completion.chunk",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
|
}
|
|
|
|
processor._process_line(ctx, sse_parser, f"data: {json.dumps(usage_chunk)}\n")
|
|
processor._process_line(ctx, sse_parser, "\n")
|
|
|
|
assert ctx.input_tokens == 10
|
|
assert ctx.output_tokens == 5
|
|
|
|
|
|
def test_process_line_handles_openai_usage_chunk_followed_by_done_without_blank_line() -> None:
|
|
ctx = StreamContext(model="test-model", api_format="openai:chat")
|
|
ctx.provider_api_format = "openai:chat"
|
|
processor = StreamProcessor(request_id="test-request", default_parser=DummyParser())
|
|
sse_parser = SSEEventParser()
|
|
|
|
usage_chunk = {
|
|
"id": "chatcmpl_test",
|
|
"object": "chat.completion.chunk",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10},
|
|
}
|
|
|
|
# Some SSE implementations may emit consecutive data lines without an empty separator.
|
|
processor._process_line(ctx, sse_parser, f"data: {json.dumps(usage_chunk)}\n")
|
|
processor._process_line(ctx, sse_parser, "data: [DONE]\n")
|
|
processor._process_line(ctx, sse_parser, "\n")
|
|
|
|
assert ctx.input_tokens == 7
|
|
assert ctx.output_tokens == 3
|
|
assert ctx.has_completion is True
|
|
|
|
|
|
class _DummyResponseCtx:
|
|
async def __aexit__(self, exc_type: type | None, exc: BaseException | None, tb: object) -> None:
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_response_stream_flushes_usage_on_remote_protocol_error() -> None:
|
|
ctx = StreamContext(model="test-model", api_format="openai:chat")
|
|
ctx.provider_api_format = "openai:chat"
|
|
processor = StreamProcessor(request_id="test-request", default_parser=DummyParser())
|
|
|
|
usage_chunk = {
|
|
"id": "chatcmpl_test",
|
|
"object": "chat.completion.chunk",
|
|
"choices": [],
|
|
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
|
|
}
|
|
|
|
async def _iter_bytes_then_remote_protocol_error() -> AsyncIterator[bytes]:
|
|
yield f"data: {json.dumps(usage_chunk)}\n".encode("utf-8")
|
|
raise httpx.RemoteProtocolError("boom")
|
|
|
|
out = b""
|
|
async for b in processor.create_response_stream(
|
|
ctx=ctx,
|
|
byte_iterator=_iter_bytes_then_remote_protocol_error(),
|
|
response_ctx=_DummyResponseCtx(),
|
|
prefetched_chunks=[],
|
|
start_time=None,
|
|
):
|
|
out += b
|
|
|
|
# Stream ends gracefully (no exception), but usage is best-effort captured and request is marked failed.
|
|
assert b"data:" in out
|
|
assert ctx.input_tokens == 11
|
|
assert ctx.output_tokens == 4
|
|
assert ctx.status_code == 502
|
|
assert (ctx.error_message or "").startswith("upstream_stream_error:")
|