ai-resume/corpus/tests/test_embeddings.py

"""Providers: determinism offline, and an OpenAI-compatible wire shape through OmniRoute.

There is no direct vendor SDK call anywhere (AGENTS.md §4.5), so the provider
test asserts the HTTP request shape, not a client library's behaviour.
"""

from __future__ import annotations

import math
import unittest

from corpus import embeddings
from corpus.tests import support


class LocalHashingTest(unittest.TestCase):
    def test_deterministic_and_unit_length(self):
        provider = embeddings.LocalHashingProvider(dim=64)
        first = provider.embed(["zigbee mesh latency"])[0]
        second = embeddings.LocalHashingProvider(dim=64).embed(["zigbee mesh latency"])[0]
        self.assertEqual(first, second)
        self.assertEqual(len(first), 64)
        self.assertAlmostEqual(math.sqrt(sum(value * value for value in first)), 1.0, places=9)

    def test_different_text_differs(self):
        provider = embeddings.LocalHashingProvider(dim=64)
        a, b = provider.embed(["zigbee", "lora"])
        self.assertNotEqual(a, b)

    def test_empty_text_still_returns_a_unit_vector(self):
        provider = embeddings.LocalHashingProvider(dim=16)
        vector = provider.embed([""])[0]
        self.assertEqual(len(vector), 16)
        self.assertAlmostEqual(math.sqrt(sum(value * value for value in vector)), 1.0, places=9)

    def test_stats_count_calls_and_texts(self):
        provider = embeddings.LocalHashingProvider(dim=16)
        provider.embed(["a", "b", "c"])
        provider.embed(["d"])
        self.assertEqual(provider.stats(), {"calls": 2, "texts": 4})
        provider.reset_stats()
        self.assertEqual(provider.stats(), {"calls": 0, "texts": 0})

    def test_rejects_a_zero_dimension(self):
        with self.assertRaises(embeddings.EmbeddingError):
            embeddings.LocalHashingProvider(dim=0)


class OmniRouteTest(unittest.TestCase):
    def responder_factory(self, dim=8, status_plan=None):
        plan = list(status_plan or [])
        state = {"n": 0}

        def responder(method, path, body):
            state["n"] += 1
            if plan:
                status = plan.pop(0)
                if status != 200:
                    return status, {"error": "planned failure"}
            vectors = [
                {"index": index, "embedding": [0.1] * dim}
                for index, _text in enumerate(body.get("input") or [])
            ]
            return 200, {"model": body.get("model"), "data": vectors}

        return responder

    def test_posts_to_v1_embeddings_with_the_configured_model(self):
        with support.StubServer(self.responder_factory()) as server:
            provider = embeddings.OmniRouteProvider(
                api_url=server.url, api_key="test-key", model="text-embedding-3-small", dim=8
            )
            vectors = provider.embed(["one", "two"])
        self.assertEqual(len(vectors), 2)
        request = server.requests[0]
        self.assertEqual(request["method"], "POST")
        self.assertEqual(request["path"], "/v1/embeddings")
        self.assertEqual(request["body"]["model"], "text-embedding-3-small")
        self.assertEqual(request["body"]["input"], ["one", "two"])
        self.assertEqual(request["headers"]["Authorization"], "Bearer test-key")
        self.assertEqual(provider.stats(), {"calls": 1, "texts": 2})

    def test_batches_requests(self):
        with support.StubServer(self.responder_factory()) as server:
            provider = embeddings.OmniRouteProvider(
                api_url=server.url, api_key="k", model="m", dim=8, batch_size=2
            )
            provider.embed(["a", "b", "c", "d", "e"])
        self.assertEqual(len(server.requests), 3)
        self.assertEqual(provider.stats(), {"calls": 3, "texts": 5})

    def test_dimension_mismatch_is_refused(self):
        with support.StubServer(self.responder_factory(dim=4)) as server:
            provider = embeddings.OmniRouteProvider(
                api_url=server.url, api_key="k", model="m", dim=8
            )
            with self.assertRaises(embeddings.EmbeddingError) as caught:
                provider.embed(["one"])
        self.assertIn("dimension", str(caught.exception))

    def test_retries_then_succeeds(self):
        with support.StubServer(self.responder_factory(status_plan=[429, 200])) as server:
            provider = embeddings.OmniRouteProvider(
                api_url=server.url, api_key="k", model="m", dim=8, backoff=0.01
            )
            vectors = provider.embed(["one"])
        self.assertEqual(len(vectors), 1)
        self.assertEqual(len(server.requests), 2)

    def test_client_error_is_not_retried(self):
        with support.StubServer(self.responder_factory(status_plan=[400])) as server:
            provider = embeddings.OmniRouteProvider(
                api_url=server.url, api_key="k", model="m", dim=8, backoff=0.01
            )
            with self.assertRaises(embeddings.EmbeddingError):
                provider.embed(["one"])
        self.assertEqual(len(server.requests), 1)

    def test_missing_key_is_refused_before_any_call(self):
        with self.assertRaises(embeddings.EmbeddingError):
            embeddings.OmniRouteProvider(api_url="http://127.0.0.1:1", api_key="", model="m", dim=8)


if __name__ == "__main__":
    unittest.main()