diff --git a/src/model_capability_readers/__init__.py b/src/model_capability_readers/__init__.py index df5c28581..587ec6507 100644 --- a/src/model_capability_readers/__init__.py +++ b/src/model_capability_readers/__init__.py @@ -25,7 +25,9 @@ from src.model_capability_readers import ( ) from src.model_capability_readers.base import ( CANONICAL_MODEL_SHAPE_VERSION, + CANONICAL_RUNTIME_CONTEXT_SHAPE_VERSION, ModelCapabilityRecord, + RuntimeContextAllocationRecord, VENDOR_ANTHROPIC, VENDOR_CEREBRAS, VENDOR_CHATGPT_SUBSCRIPTION, @@ -250,7 +252,9 @@ def records_from_payload( __all__ = [ "ModelCapabilityRecord", + "RuntimeContextAllocationRecord", "CANONICAL_MODEL_SHAPE_VERSION", + "CANONICAL_RUNTIME_CONTEXT_SHAPE_VERSION", "PLACEHOLDER_VENDOR_IDS", "READER_MODULES", "VENDOR_ANTHROPIC", diff --git a/src/model_capability_readers/base.py b/src/model_capability_readers/base.py index e07b10a55..c9f0aae18 100644 --- a/src/model_capability_readers/base.py +++ b/src/model_capability_readers/base.py @@ -46,6 +46,7 @@ VENDOR_ZAI = "zai" VENDOR_UNKNOWN = "unknown" CANONICAL_MODEL_SHAPE_VERSION = 1 +CANONICAL_RUNTIME_CONTEXT_SHAPE_VERSION = 1 @dataclass(frozen=True) @@ -109,6 +110,53 @@ class ModelCapabilityRecord: return data +@dataclass(frozen=True) +class RuntimeContextAllocationRecord: + """Provider-reported context allocated to one currently loaded model. + + This runtime state is deliberately separate from ``ModelCapability.limits``: + a model's maximum, configured ``num_ctx``, and the allocation reported for a + loaded process are different facts with different lifetimes. + """ + + vendor: str + model_id: str + allocated_context_tokens: int + stable_model_id: str = "" + source: str = mc.SOURCE_PROVIDER_READER + confidence: str = mc.CONFIDENCE_PROVIDER_REPORTED + runtime_shape_id: str = "" + raw: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + context_tokens = int_limit(self.allocated_context_tokens) + if context_tokens is None: + raise ValueError("allocated_context_tokens must be a positive integer") + object.__setattr__(self, "allocated_context_tokens", context_tokens) + if not self.stable_model_id: + object.__setattr__(self, "stable_model_id", stable_model_id_for(self.vendor, self.model_id)) + + def to_dict(self, *, include_raw: bool = False) -> dict[str, Any]: + data = { + "schema_version": CANONICAL_RUNTIME_CONTEXT_SHAPE_VERSION, + "provider": self.vendor, + "model": self.model_id, + "stable_id": self.stable_model_id, + "runtime": { + "allocated_context_tokens": self.allocated_context_tokens, + }, + "evidence": { + "source": self.source, + "confidence": self.confidence, + "shape": self.runtime_shape_id, + "scope": "loaded_model", + }, + } + if include_raw: + data["raw"] = dict(self.raw) + return data + + class CapabilityReader(Protocol): vendor: str diff --git a/src/model_capability_readers/ollama.py b/src/model_capability_readers/ollama.py index 2333b3574..2a8c341e4 100644 --- a/src/model_capability_readers/ollama.py +++ b/src/model_capability_readers/ollama.py @@ -8,6 +8,7 @@ from typing import Any from src import model_capabilities as mc from src.model_capability_readers.base import ( ModelCapabilityRecord, + RuntimeContextAllocationRecord, VENDOR_OLLAMA, as_list, as_mapping, @@ -24,6 +25,9 @@ from src.model_capability_readers.base import ( vendor = VENDOR_OLLAMA +OLLAMA_PS_SHAPE_ID = "ollama.ps.v1" + + _CAPABILITY_MAP = { "completion": None, "completions": None, @@ -110,6 +114,65 @@ def _limits_from_show(raw: Mapping[str, Any]) -> dict[str, Any]: return limits +def _runtime_model_key(value: Any) -> str: + model_id = identity_str(value).casefold() + return model_id.removesuffix(":latest") + + +def runtime_context_from_ps_payload( + model_id: str, + payload: Mapping[str, Any], + *, + endpoint_id: Any = "", + base_url: Any = "", +) -> RuntimeContextAllocationRecord | None: + """Normalize one loaded model's allocation from Ollama ``GET /api/ps``. + + Exact identity is preferred, with only Ollama's implicit ``:latest`` alias + normalized. Conflicting matching rows fail closed instead of choosing an + arbitrary allocation. + """ + + requested_id = identity_str(model_id) + requested_key = _runtime_model_key(requested_id) + if not requested_key: + return None + + matches: list[tuple[int, Mapping[str, Any]]] = [] + for item in as_list(as_mapping(payload).get("models")): + if not isinstance(item, Mapping): + continue + identities = tuple( + identity + for identity in (identity_str(item.get("model")), identity_str(item.get("name"))) + if identity + ) + if not any(_runtime_model_key(identity) == requested_key for identity in identities): + continue + context_tokens = int_limit(item.get("context_length")) + if context_tokens is not None: + matches.append((context_tokens, item)) + + allocations = {context_tokens for context_tokens, _item in matches} + if len(allocations) != 1: + return None + + context_tokens, raw = matches[0] + return RuntimeContextAllocationRecord( + vendor=VENDOR_OLLAMA, + model_id=requested_id, + stable_model_id=stable_model_id_for( + VENDOR_OLLAMA, + requested_id, + endpoint_id=endpoint_id, + base_url=base_url, + ), + allocated_context_tokens=context_tokens, + runtime_shape_id=OLLAMA_PS_SHAPE_ID, + raw=raw, + ) + + def record_from_show_payload( model_id: str, payload: Mapping[str, Any], diff --git a/tests/test_model_capability_readers.py b/tests/test_model_capability_readers.py index ca547b15a..e69cca761 100644 --- a/tests/test_model_capability_readers.py +++ b/tests/test_model_capability_readers.py @@ -404,6 +404,68 @@ def test_ollama_reader_uses_generic_model_info_context_length_when_no_num_ctx(): assert dict(record.capability.limits) == {"context_tokens": 32768} +def test_ollama_ps_reader_keeps_loaded_allocation_in_runtime_shape(): + record = ollama.runtime_context_from_ps_payload( + "hf.co/example/Qwen3:Q6_K", + { + "models": [ + { + "name": "hf.co/example/Qwen3:Q6_K", + "model": "hf.co/example/Qwen3:Q6_K", + "context_length": 65536, + "details": {"family": "qwen3"}, + } + ] + }, + endpoint_id="7", + ) + + assert record is not None + assert record.allocated_context_tokens == 65536 + assert record.stable_model_id == "ollama|endpoint:7|hf.co/example/qwen3:q6_k" + serialized = record.to_dict() + assert serialized["runtime"] == {"allocated_context_tokens": 65536} + assert serialized["evidence"] == { + "source": mc.SOURCE_PROVIDER_READER, + "confidence": mc.CONFIDENCE_PROVIDER_REPORTED, + "shape": "ollama.ps.v1", + "scope": "loaded_model", + } + assert "limits" not in serialized + assert "raw" not in serialized + + +def test_ollama_ps_reader_matches_only_exact_or_latest_identity(): + payload = { + "models": [ + {"model": "qwen3:latest", "context_length": 32768}, + {"model": "qwen3:14b", "context_length": 65536}, + ] + } + + record = ollama.runtime_context_from_ps_payload("qwen3", payload) + + assert record is not None + assert record.allocated_context_tokens == 32768 + assert ollama.runtime_context_from_ps_payload("qwen3:8b", payload) is None + + +def test_ollama_ps_reader_fails_closed_for_invalid_or_conflicting_allocation(): + assert ollama.runtime_context_from_ps_payload( + "qwen3", + {"models": [{"model": "qwen3", "context_length": True}]}, + ) is None + assert ollama.runtime_context_from_ps_payload( + "qwen3", + { + "models": [ + {"model": "qwen3", "context_length": 32768}, + {"name": "qwen3:latest", "context_length": 65536}, + ] + }, + ) is None + + def test_lmstudio_reader_uses_native_v1_capabilities_when_present(): records = lmstudio.records_from_payload( {