PolyGateway/tests/unit/test_embedding.py

"""Embedding 类型/端口/transport 测试(M2 设计 §7;T8)。

蓝本审计: GovDoc retrieval/embedding.py(分批/index 排序/维度校验)与
VT adapters/embedding.py(归一化);库裁决见设计 §7.3 表。
"""

import dataclasses
import json

import httpx
import pytest

from polygateway.errors import (
    RequestRejectedError,
    ResultInvalidError,
    SourceDeadError,
    TransientError,
)
from polygateway.ports import EmbeddingTransport
from polygateway.transports.openai_compat import OpenAICompatTransport
from polygateway.types import EmbeddingResponse, EmbeddingTransportResult, SourceConfig


def _src(**overrides):
    base = {
        "name": "e1",
        "provider": "openai",
        "base_url": "https://gw.example/v1",
        "api_key": "sk",
        "model": "embed-1",
        "timeout_s": 10.0,
        "est_tokens": 7,
    }
    base.update(overrides)
    return SourceConfig(**base)


class TestTypes:
    def test_embedding_response_frozen_with_defaults(self):
        resp = EmbeddingResponse(
            vectors=[[0.1, 0.2]],
            dim=2,
            model="m",
            provider="p",
            prompt_tokens=3,
            usage_source="measured",
            latency_ms=10,
            call_id="c",
            source_name="e1",
        )
        assert resp.cost is None
        with pytest.raises(dataclasses.FrozenInstanceError):
            resp.dim = 3

    def test_transport_result_frozen(self):
        r = EmbeddingTransportResult(
            vectors=[[1.0]], dim=1, prompt_tokens=1, usage_source="measured", raw={}
        )
        with pytest.raises(dataclasses.FrozenInstanceError):
            r.dim = 2


class _DummyEmbedTransport:
    async def embed(self, *, texts, source, call_id):
        raise NotImplementedError


def test_embedding_transport_protocol_runtime_checkable():
    assert isinstance(_DummyEmbedTransport(), EmbeddingTransport)
    assert isinstance(OpenAICompatTransport(), EmbeddingTransport)


def _transport_with(handler):
    return OpenAICompatTransport(
        client_factory=lambda source: httpx.AsyncClient(transport=httpx.MockTransport(handler))
    )


def _ok_body(vectors, *, usage=None, shuffle=False):
    data = [{"index": i, "embedding": v} for i, v in enumerate(vectors)]
    if shuffle:
        data = list(reversed(data))
    body = {"data": data}
    if usage is not None:
        body["usage"] = usage
    return body


class TestEmbedTransport:
    async def test_sorts_by_index_and_measures_usage(self):
        def handler(request):
            assert request.url.path.endswith("/embeddings")
            payload = json.loads(request.content)
            assert payload == {"model": "embed-1", "input": ["a", "b"]}
            return httpx.Response(
                200,
                json=_ok_body(
                    [[1.0, 0.0], [0.0, 1.0]], usage={"prompt_tokens": 5}, shuffle=True
                ),
            )

        result = await _transport_with(handler).embed(
            texts=["a", "b"], source=_src(), call_id="c"
        )
        assert result.vectors == [[1.0, 0.0], [0.0, 1.0]]  # 乱序响应按 index 重排
        assert result.dim == 2
        assert result.prompt_tokens == 5 and result.usage_source == "measured"

    async def test_missing_usage_falls_back_estimated(self):
        def handler(request):
            return httpx.Response(200, json=_ok_body([[1.0]]))

        result = await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")
        assert result.prompt_tokens == 7 and result.usage_source == "estimated"  # est_tokens

    @pytest.mark.parametrize(
        ("status", "exc_type"),
        [(401, SourceDeadError), (400, RequestRejectedError), (500, TransientError)],
    )
    async def test_http_errors_translate(self, status, exc_type):
        def handler(request):
            return httpx.Response(status, text="boom")

        with pytest.raises(exc_type):
            await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")

    async def test_network_error_is_transient(self):
        def handler(request):
            raise httpx.ConnectError("refused")

        with pytest.raises(TransientError):
            await _transport_with(handler).embed(texts=["a"], source=_src(), call_id="c")

    @pytest.mark.parametrize(
        "body",
        [
            {"data": []},  # 空 data
            {"data": [{"index": 0, "embedding": [1.0]}]},  # 数量与输入不符(输入 2 条)
            {
                "data": [
                    {"index": 0, "embedding": [1.0, 2.0]},
                    {"index": 1, "embedding": [1.0]},  # 维度不一致
                ]
            },
            {"nope": True},  # 缺 data
        ],
    )
    async def test_malformed_payload_is_result_invalid(self, body):
        def handler(request):
            return httpx.Response(200, json=body)

        with pytest.raises(ResultInvalidError):
            await _transport_with(handler).embed(texts=["a", "b"], source=_src(), call_id="c")

    async def test_empty_texts_rejected(self):
        with pytest.raises(ValueError):
            await _transport_with(lambda r: None).embed(texts=[], source=_src(), call_id="c")