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