File size: 2,383 Bytes
d61821a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | 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()
|