1"""Unit tests: formats, estimator, compatibility engine, OpenAI helpers."""23from __future__ import annotations45from pathlib import Path67from llm_api.models import compat, formats8from llm_api.models.estimator import estimate, kv_bytes_per_token, recommended_context9from llm_api.worker.openai_types import StopMatcher, ThinkSplitter, parse_tool_calls1011from .conftest import make_gguf, make_mlx_model121314def test_parse_param_count():15 assert formats.parse_param_count_from_name("Qwen3-30B-A3B-4bit") == (30_000_000_000, 3_000_000_000)16 assert formats.parse_param_count_from_name("gemma-3-4b-it-qat-4bit")[0] == 4_000_000_00017 assert formats.parse_param_count_from_name("Qwen3-Embedding-0.6B-8bit")[0] == 600_000_00018 assert formats.parse_param_count_from_name("all-MiniLM-L6-v2-4bit") == (None, None)19 assert formats.parse_param_count_from_name("embeddinggemma-300m-4bit")[0] == 300_000_000202122def test_parse_quant():23 assert formats.parse_quant_from_name("Qwen3-4B-Instruct-2507-4bit") == ("4bit", 4.0)24 assert formats.parse_quant_from_name("gpt-oss-20b-MXFP4-Q8")[0] == "MXFP4"25 q, b = formats.parse_quant_from_name("gemma-3-1b-it-Q4_K_M")26 assert q == "Q4_K_M" and b == 4.527 assert formats.parse_quant_from_name("Qwen3-Embedding-0.6B-4bit-DWQ") == ("DWQ-4bit", 4.0)28 assert formats.parse_quant_from_name("Qwen3.8-27B-bf16") == ("bf16", 16)293031def test_family():32 assert formats.guess_family("Qwen3.6-35B-A3B-4bit") == "qwen"33 assert formats.guess_family("gemma-4-26b-a4b-it-4bit") == "gemma"34 assert formats.guess_family("Devstral-Small-2-24B") == "mistral"35 assert formats.guess_family("gpt-oss-20b-MXFP4-Q8") == "gpt-oss"36 assert formats.guess_family("Something-Unknown", "phi3") == "phi"373839def test_size_class():40 assert formats.size_class(3, 45) == "TINY"41 assert formats.size_class(7, 45) == "SMALL"42 assert formats.size_class(15, 45) == "MEDIUM"43 assert formats.size_class(30, 45) == "LARGE"44 assert formats.size_class(40, 45) == "XL"45 assert formats.size_class(50, 45) == "TOO_LARGE"464748def test_kv_and_estimate():49 kv = kv_bytes_per_token(36, 8, 128, 16)50 assert kv == 2 * 36 * 8 * 128 * 251 est = estimate(2_400_000_000, "mlx", kv, 32768)52 assert 9 < est.total_gb < 1153 ctx, e2 = recommended_context(2_400_000_000, "mlx", kv, 262144, 45)54 assert ctx == 32768 # capped at the preferred context55 ctx3, _ = recommended_context(40 * 1024**3, "mlx", kv * 4, 131072, 45)56 assert ctx3 is not None and ctx3 <= 8192575859def test_compat_statuses(monkeypatch):60 monkeypatch.setattr(compat, "mlx_available", lambda: True)61 monkeypatch.setattr(compat, "mlx_lm_model_types", lambda: {"qwen3"})62 kv = kv_bytes_per_token(36, 8, 128)63 ok = compat.evaluate(runtime="mlx", weights_bytes=2_400_000_000, kv_per_token=kv, max_context=131072, model_type="qwen3",64 architecture="Qwen3ForCausalLM", vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,65 llamacpp_available=True)66 assert ok.status == compat.COMPATIBLE and ok.recommended_context == 3276867 too_big = compat.evaluate(runtime="mlx", weights_bytes=60 * 1024**3, kv_per_token=kv, max_context=131072, model_type="qwen3",68 architecture=None, vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,69 llamacpp_available=True)70 assert too_big.status == compat.INCOMPATIBLE and not too_big.compatible71 swap = compat.evaluate(runtime="mlx", weights_bytes=44 * 1024**3, kv_per_token=kv, max_context=131072, model_type="qwen3",72 architecture=None, vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,73 llamacpp_available=True)74 assert swap.status == compat.NOT_RECOMMENDED75 unknown = compat.evaluate(runtime="mlx", weights_bytes=1e9, kv_per_token=kv, max_context=8192, model_type="mystery",76 architecture=None, vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,77 llamacpp_available=True)78 assert unknown.status == compat.INCOMPATIBLE79 no_llama = compat.evaluate(runtime="llamacpp", weights_bytes=1e9, kv_per_token=kv, max_context=8192, model_type="llama",80 architecture="llama", vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,81 llamacpp_available=False, weights_file="x.gguf")82 assert no_llama.status == compat.INCOMPATIBLE83 exp = compat.evaluate(runtime="llamacpp", weights_bytes=1e9, kv_per_token=kv, max_context=8192, model_type="weird",84 architecture="weird", vision=False, embedding=False, reranker=False, budget_gb=45, absolute_gb=50,85 llamacpp_available=True, weights_file="x.gguf")86 assert exp.status == compat.EXPERIMENTAL878889def test_gguf_reader(tmp_path: Path):90 p = tmp_path / "m-Q4_K_M.gguf"91 make_gguf(p, arch="llama", layers=8, kv_heads=4, head_dim=64, ctx=4096)92 info = formats.read_gguf(p)93 assert info.architecture == "llama" and info.n_layers == 8 and info.n_kv_heads == 4 and info.head_dim == 6494 assert info.context_length == 4096 and info.file_type_label == "Q4_K_M"95 assert info.param_count == 1_000_000969798def test_safetensors_params(tmp_path: Path):99 d = tmp_path / "m"100 make_mlx_model(d, layers=2, hidden=256, bits=4)101 params, nbytes = formats.count_safetensors_params([d / "model.safetensors"], 4)102 # 2 layers * (1024 x 256) + embeddings 1000*256103 assert params == 2 * 1024 * 256 + 1000 * 256104 assert nbytes > 0105106107def test_hf_config_parse():108 cfg = {"model_type": "qwen3", "architectures": ["Qwen3ForCausalLM"], "num_hidden_layers": 36, "num_attention_heads": 32,109 "num_key_value_heads": 8, "hidden_size": 2560, "head_dim": 128, "max_position_embeddings": 262144,110 "quantization": {"bits": 4, "group_size": 64}}111 p = formats.parse_hf_config(cfg)112 assert p["n_layers"] == 36 and p["n_kv_heads"] == 8 and p["head_dim"] == 128 and p["quant_bits"] == 4 and not p["vision"]113 v = formats.parse_hf_config({"model_type": "qwen3_vl", "text_config": {"num_hidden_layers": 2}, "vision_config": {}})114 assert v["vision"] and v["n_layers"] == 2115116117def test_stop_matcher():118 sm = StopMatcher(["</s>", "User:"])119 out = sm.feed("Hello <")120 assert out == "Hello " and not sm.done121 out += sm.feed("/s> tail")122 assert out == "Hello " and sm.done123 sm2 = StopMatcher(["STOP"])124 assert sm2.feed("abc ST") == "abc "125 assert sm2.feed("art") == "STart"126 assert sm2.flush() == ""127128129def test_think_splitter():130 ts = ThinkSplitter("<think>", "</think>")131 r, c = ts.feed("<think>reasoning here</thi")132 assert r == "reasoning here" and c == ""133 r2, c2 = ts.feed("nk>\n\nanswer")134 assert r2 == "" and c2 == "answer"135136137def test_tool_calls():138 rest, calls = parse_tool_calls('Sure.\n<tool_call>\n{"name": "get_weather", "arguments": {"city": "Montreal"}}\n</tool_call>')139 assert rest == "Sure." and len(calls) == 1 and calls[0]["function"]["name"] == "get_weather"140 assert '"city": "Montreal"' in calls[0]["function"]["arguments"]141