RAG/backend/tests/contract/test_bailian_embedding.py

from __future__ import annotations

import json

import httpx
import pytest

from app.adapters.bailian._base import response_model
from app.adapters.bailian.embedding import BailianEmbeddingAdapter
from app.ports.model_providers import ModelProviderError, ProviderErrorKind

BASE_URL = "https://workspace.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"
API_KEY = "sk-test-embedding-secret"


def _vector(marker: float = 1.0) -> list[float]:
    return [marker, *([0.0] * 1_023)]


def test_embedding_rejects_non_bailian_endpoint_before_request() -> None:
    with pytest.raises(ModelProviderError) as exc_info:
        BailianEmbeddingAdapter(
            api_key=API_KEY,
            base_url="https://attacker.invalid/compatible-mode/v1",
        )

    assert exc_info.value.kind is ProviderErrorKind.INVALID_REQUEST
    assert exc_info.value.provider_code == "invalid_base_url"


def test_response_model_falls_back_when_provider_echoes_sensitive_value() -> None:
    assert (
        response_model(
            {"model": API_KEY},
            "text-embedding-v4",
            sensitive_values=(API_KEY,),
        )
        == "text-embedding-v4"
    )


@pytest.mark.asyncio
async def test_embedding_uses_compatible_endpoint_and_restores_index_order() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        assert request.url.path == "/compatible-mode/v1/embeddings"
        assert request.headers["Authorization"] == f"Bearer {API_KEY}"
        payload = json.loads(request.content)
        assert payload == {
            "model": "text-embedding-v4",
            "input": ["first", "second"],
            "dimensions": 1_024,
            "encoding_format": "float",
        }
        return httpx.Response(
            200,
            json={
                "id": "req_embedding_1",
                "model": "text-embedding-v4",
                "data": [
                    {"index": 1, "embedding": _vector(2.0)},
                    {"index": 0, "embedding": _vector(1.0)},
                ],
                "usage": {"prompt_tokens": 2, "total_tokens": 2},
            },
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianEmbeddingAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
        )
        result = await adapter.embed_documents(["first", "second"])

    assert len(result.vectors) == 2
    assert all(len(vector) == 1_024 for vector in result.vectors)
    assert result.vectors[0][0] == 1.0
    assert result.vectors[1][0] == 2.0
    assert result.request_id == "req_embedding_1"
    assert result.usage.input_tokens == 2


@pytest.mark.asyncio
async def test_embedding_enforces_10_8192_and_33000_token_boundaries() -> None:
    calls = 0

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        payload = json.loads(request.content)
        return httpx.Response(
            200,
            json={
                "data": [
                    {"index": index, "embedding": _vector()}
                    for index, _ in enumerate(payload["input"])
                ]
            },
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianEmbeddingAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
        )

        await adapter.embed_documents(["x"] * 10)
        await adapter.embed_documents(["x" * 8_192])
        await adapter.embed_documents([*(["x" * 8_192] * 4), "x" * 232])
        assert calls == 3

        invalid_batches = (
            ["x"] * 11,
            ["x" * 8_193],
            [*(["x" * 8_192] * 4), "x" * 233],
        )
        for texts in invalid_batches:
            with pytest.raises(ModelProviderError) as exc_info:
                await adapter.embed_documents(texts)
            assert exc_info.value.kind is ProviderErrorKind.INVALID_REQUEST
        assert calls == 3


@pytest.mark.asyncio
async def test_embedding_rejects_non_1024_dimensional_response() -> None:
    def handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={"data": [{"index": 0, "embedding": [1.0] * 1_023}]},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianEmbeddingAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
        )
        with pytest.raises(ModelProviderError) as exc_info:
            await adapter.embed_query("sensitive geological paragraph")

    assert exc_info.value.kind is ProviderErrorKind.INVALID_RESPONSE
    assert exc_info.value.provider_code == "invalid_embedding_dimensions"


@pytest.mark.asyncio
async def test_embedding_provider_error_does_not_expose_key_or_body() -> None:
    sensitive_text = "CONFIDENTIAL_GEOLOGICAL_BODY"

    def handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(
            400,
            json={
                "id": sensitive_text,
                "error": {
                    "code": API_KEY,
                    "message": f"failed for {sensitive_text} using {API_KEY}",
                },
            },
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianEmbeddingAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
        )
        with pytest.raises(ModelProviderError) as exc_info:
            await adapter.embed_query(sensitive_text)

    rendered = f"{exc_info.value!s} {exc_info.value!r}"
    assert API_KEY not in rendered
    assert sensitive_text not in rendered
    assert exc_info.value.provider_code is None
    assert exc_info.value.request_id is None


@pytest.mark.asyncio
async def test_embedding_retries_timeout_with_a_fixed_bound() -> None:
    calls = 0

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal calls
        calls += 1
        if calls == 1:
            raise httpx.ReadTimeout("simulated timeout", request=request)
        return httpx.Response(
            200,
            json={"data": [{"index": 0, "embedding": _vector()}]},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        adapter = BailianEmbeddingAdapter(
            api_key=API_KEY,
            base_url=BASE_URL,
            token_counter=len,
            http_client=client,
            max_retries=1,
            retry_base_seconds=0,
        )
        result = await adapter.embed_query("query")

    assert len(result.vectors[0]) == 1_024
    assert calls == 2