from __future__ import annotations import json from collections.abc import AsyncIterator import httpx import pytest from app.adapters.model_gateway import ModelGatewayAdapter from app.ports.model_providers import ChatMessage, ModelProviderError, ProviderErrorKind def _usage() -> dict[str, int]: return {"input_tokens": 4, "output_tokens": 0, "total_tokens": 4} def _vector() -> list[float]: return [1.0] + [0.0] * 1023 @pytest.mark.asyncio async def test_query_embedding_uses_fixed_origin_auth_and_query_type() -> None: seen: list[httpx.Request] = [] def handler(request: httpx.Request) -> httpx.Response: seen.append(request) return httpx.Response( 200, json={ "vectors": [_vector()], "model": "text-embedding-v4", "request_id": "req-1", "usage": _usage(), "elapsed_ms": 8.5, }, ) client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) try: result = await adapter.embed_query("斑岩铜矿") finally: await client.aclose() assert len(result.vectors[0]) == 1024 assert seen[0].url == "http://model-gateway:8000/internal/v1/embeddings" assert seen[0].headers["authorization"] == "Bearer internal-token" assert seen[0].headers["x-rag-caller"] == "api" assert json.loads(seen[0].content) == {"texts": ["斑岩铜矿"], "input_type": "query"} @pytest.mark.asyncio async def test_document_embedding_identifies_worker_caller() -> None: def handler(request: httpx.Request) -> httpx.Response: assert request.headers["x-rag-caller"] == "worker" assert json.loads(request.content)["input_type"] == "document" return httpx.Response( 200, json={ "vectors": [_vector()], "model": "text-embedding-v4", "request_id": None, "usage": _usage(), "elapsed_ms": 1, }, ) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="worker-token", caller="worker", http_client=client) await adapter.embed_documents(["合成文档"]) @pytest.mark.asyncio async def test_gateway_http_error_does_not_leak_response_or_token() -> None: def handler(_: httpx.Request) -> httpx.Response: return httpx.Response(401, json={"detail": "internal-token private upstream body"}) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) with pytest.raises(ModelProviderError) as caught: await adapter.embed_query("secret query") assert caught.value.kind is ProviderErrorKind.AUTHENTICATION assert "internal-token" not in str(caught.value) assert "secret query" not in str(caught.value) @pytest.mark.asyncio async def test_gateway_preserves_only_fixed_provider_error_category() -> None: def handler(_: httpx.Request) -> httpx.Response: return httpx.Response( 502, json={ "error": { "kind": "authentication", "retryable": False, "request_id": "req-safe", "untrusted_extra": "must-not-cross-boundary", } }, ) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) with pytest.raises(ModelProviderError) as caught: await adapter.embed_query("secret query") assert caught.value.kind is ProviderErrorKind.AUTHENTICATION assert caught.value.status_code == 502 assert caught.value.request_id == "req-safe" assert "untrusted_extra" not in str(caught.value) @pytest.mark.asyncio async def test_rerank_rejects_document_mismatch_from_gateway() -> None: def handler(_: httpx.Request) -> httpx.Response: return httpx.Response( 200, json={ "items": [{"index": 0, "relevance_score": 0.9, "document": "tampered"}], "model": "qwen3-rerank", "request_id": "req-2", "usage": _usage(), "elapsed_ms": 3, }, ) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) with pytest.raises(ModelProviderError) as caught: await adapter.rerank("铜矿", ["候选"], top_n=1) assert caught.value.kind is ProviderErrorKind.INVALID_RESPONSE @pytest.mark.asyncio async def test_chat_completion_maps_gateway_response() -> None: def handler(_: httpx.Request) -> httpx.Response: return httpx.Response( 200, json={ "content": "仅依据证据回答。[S1]", "finish_reason": "stop", "model": "deepseek-v4-flash", "request_id": "req-chat", "usage": {"input_tokens": 5, "output_tokens": 3, "total_tokens": 8}, "elapsed_ms": 9, }, ) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) result = await adapter.complete([ChatMessage(role="user", content="问题")], max_tokens=20) assert result.content.endswith("[S1]") assert result.usage.total_tokens == 8 class _SseStream(httpx.AsyncByteStream): async def __aiter__(self) -> AsyncIterator[bytes]: yield ( b"event: delta\n" b'data: {"delta":"answer","model":"deepseek-v4-flash",' b'"request_id":"req"}\n\n' ) yield ( b"event: complete\n" b'data: {"delta":"","finish_reason":"stop",' b'"model":"deepseek-v4-flash","request_id":"req",' b'"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3},' b'"elapsed_ms":5}\n\n' ) class _IncompleteSseStream(httpx.AsyncByteStream): async def __aiter__(self) -> AsyncIterator[bytes]: yield ( b"event: delta\n" b'data: {"delta":"partial","model":"deepseek-v4-flash",' b'"request_id":"req"}\n\n' ) class _EventAfterCompleteStream(httpx.AsyncByteStream): async def __aiter__(self) -> AsyncIterator[bytes]: yield ( b"event: complete\n" b'data: {"finish_reason":"stop","model":"deepseek-v4-flash",' b'"request_id":"req","usage":{"input_tokens":1,"output_tokens":1,' b'"total_tokens":2},"elapsed_ms":5}\n\n' ) yield ( b"event: delta\n" b'data: {"delta":"late","model":"deepseek-v4-flash",' b'"request_id":"req"}\n\n' ) @pytest.mark.asyncio async def test_chat_stream_parses_typed_events() -> None: def handler(_: httpx.Request) -> httpx.Response: return httpx.Response( 200, headers={"content-type": "text/event-stream"}, stream=_SseStream(), ) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) events = [ event async for event in adapter.stream( [ChatMessage(role="user", content="question")], max_tokens=20 ) ] assert [event.delta for event in events] == ["answer", ""] assert events[-1].finish_reason == "stop" assert events[-1].usage.total_tokens == 3 @pytest.mark.asyncio @pytest.mark.parametrize("stream", [_IncompleteSseStream(), _EventAfterCompleteStream()]) async def test_chat_stream_requires_exactly_one_terminal_complete( stream: httpx.AsyncByteStream, ) -> None: def handler(_: httpx.Request) -> httpx.Response: return httpx.Response( 200, headers={"content-type": "text/event-stream"}, stream=stream, ) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) with pytest.raises(ModelProviderError) as caught: _ = [ event async for event in adapter.stream( [ChatMessage(role="user", content="question")], max_tokens=20 ) ] assert caught.value.kind is ProviderErrorKind.INVALID_RESPONSE @pytest.mark.asyncio async def test_chat_stream_rejects_nonfinite_terminal_metadata() -> None: class InvalidTerminal(httpx.AsyncByteStream): async def __aiter__(self) -> AsyncIterator[bytes]: yield ( b"event: complete\n" b'data: {"finish_reason":"stop","model":"deepseek-v4-flash",' b'"request_id":"req","usage":{},"elapsed_ms":-1}\n\n' ) def handler(_: httpx.Request) -> httpx.Response: return httpx.Response(200, stream=InvalidTerminal()) async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: adapter = ModelGatewayAdapter(token="internal-token", caller="api", http_client=client) with pytest.raises(ModelProviderError) as caught: _ = [ event async for event in adapter.stream( [ChatMessage(role="user", content="question")], max_tokens=20 ) ] assert caught.value.kind is ProviderErrorKind.INVALID_RESPONSE def test_client_rejects_any_nonfixed_gateway_origin() -> None: with pytest.raises(ModelProviderError): ModelGatewayAdapter(token="token", caller="api", base_url="http://127.0.0.1:8000")