tests/test_lm_studio_embeddings.py

from __future__ import annotations

from pathlib import Path
import unittest

from agent_harness.lm_studio_embeddings import (
    EmbeddingStudioError,
    LMStudioEmbeddingClient,
    validate_embedding_record,
)
from agent_harness.specs import load_embeddings


ROOT = Path(__file__).resolve().parents[1]


def matching_record() -> dict[str, object]:
    return {
        "type": "embedding",
        "publisher": "Qwen",
        "key": "text-embedding-qwen3-embedding-0.6b",
        "display_name": "Qwen3 Embedding 0.6B",
        "format": "gguf",
        "quantization": {"name": "Q8_0", "bits_per_weight": 8},
        "size_bytes": 639150592,
        "max_context_length": 32768,
        "loaded_instances": [{"config": {"context_length": 8192}}],
    }


class LMStudioEmbeddingTests(unittest.TestCase):
    def setUp(self) -> None:
        self.spec = load_embeddings(ROOT)["EMB001"]

    def test_pinned_record_matches(self) -> None:
        validate_embedding_record(self.spec, matching_record())

    def test_quantization_mismatch_is_fatal(self) -> None:
        record = matching_record()
        record["quantization"] = {"name": "Q4_K_M"}
        with self.assertRaises(EmbeddingStudioError):
            validate_embedding_record(self.spec, record)

    def test_loaded_context_mismatch_is_fatal(self) -> None:
        record = matching_record()
        record["loaded_instances"] = [{"config": {"context_length": 4096}}]
        with self.assertRaises(EmbeddingStudioError):
            validate_embedding_record(self.spec, record)

    def test_probe_validates_dimensions_normalization_and_distinctness(self) -> None:
        client = LMStudioEmbeddingClient(self.spec)
        positive = [1.0 / 32.0] * 1024
        negative = [-1.0 / 32.0] * 1024
        response = {
            "model": self.spec.model_key,
            "data": [
                {"index": 0, "embedding": positive},
                {"index": 1, "embedding": negative},
            ],
            "usage": {"prompt_tokens": 0, "total_tokens": 0},
        }
        client._request = lambda *args, **kwargs: response  # type: ignore[method-assign]
        result = client.probe()
        self.assertEqual(result.vector_count, 2)
        self.assertEqual(result.vector_dimension, 1024)
        self.assertAlmostEqual(result.pairwise_cosine, -1.0)


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