diff options
| -rw-r--r-- | llamachat/models.py | 87 | ||||
| -rwxr-xr-x | test_llamachat.py | 183 |
2 files changed, 270 insertions, 0 deletions
diff --git a/llamachat/models.py b/llamachat/models.py index 988dc43..a99b6a6 100644 --- a/llamachat/models.py +++ b/llamachat/models.py @@ -19,6 +19,8 @@ import os from dataclasses import dataclass from pathlib import Path +from . import providers as providers_mod + # configparser reads a literal [DEFAULT] as inherited defaults for every # other section. A model genuinely named DEFAULT is stored under this name # instead, so the file stays safe for the stock-configparser readers that @@ -210,3 +212,88 @@ class ModelStore: except OSError: tmp.unlink(missing_ok=True) raise + + +def resolve(model_id: str, table: dict, store: "ModelStore") -> ModelInfo: + """Merge the three metadata layers, most specific first. + + Per field, not per source: a models.ini entry that sets only ctx_size + still inherits the provider's prices. + """ + stored = store.get(model_id) or ModelInfo() + name, _ = providers_mod.split(model_id, table) + provider = table.get(name) + + def pick(from_store, from_provider): + return from_store if from_store is not None else from_provider + + if provider is None: + return stored + return ModelInfo( + ctx_size=pick(stored.ctx_size, provider.ctx_size), + vision=pick(stored.vision, provider.vision), + price_in=pick(stored.price_in, provider.price_in), + price_out=pick(stored.price_out, provider.price_out), + ) + + +def is_billable(model_id: str, table: dict) -> bool: + """Whether this model costs money, regardless of prices being known. + + Keyed on the provider having an api_key: that is what distinguishes a + free local model from a cloud one whose price was never entered. + """ + name, _ = providers_mod.split(model_id, table) + provider = table.get(name) + return bool(provider and not provider.is_local and provider.api_key) + + +def is_priced(model_id: str, table: dict, store: "ModelStore") -> bool: + """Whether a cost can be computed for this model.""" + info = resolve(model_id, table, store) + return info.price_in is not None or info.price_out is not None + + +def message_cost(row, table: dict, store: "ModelStore") -> float: + """Cost of one stored assistant row, priced at the model that made it.""" + model_id = row["model"] if row["model"] else "" + if not model_id: + return 0.0 + info = resolve(model_id, table, store) + prompt = row["prompt_tokens"] or 0 + completion = row["completion_tokens"] or 0 + total = 0.0 + if info.price_in is not None: + total += prompt * info.price_in + if info.price_out is not None: + total += completion * info.price_out + return total / 1_000_000 + + +def conversation_cost(rows, table: dict, store: "ModelStore") -> float: + """Everything spent in one conversation so far. + + Rows predating the token columns carry NULLs and contribute zero, so an + old conversation reads as free rather than as a fabricated number. + """ + return sum(message_cost(row, table, store) for row in rows) + + +def projected_cost( + tokens: int, model_id: str, table: dict, store: "ModelStore" +) -> float: + """What sending `tokens` of input to this model would cost. + + Only the input rate: the length of the reply is unknowable in advance. + """ + info = resolve(model_id, table, store) + if info.price_in is None: + return 0.0 + return tokens * info.price_in / 1_000_000 + + +def format_cost(amount: float) -> str: + """Money, with enough precision to see a cheap turn move the number.""" + if amount < 1: + return f"${amount:.3f}" + return f"${amount:.2f}" diff --git a/test_llamachat.py b/test_llamachat.py index ae59d1d..695a303 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -1495,6 +1495,188 @@ def test_models_store(): print("ok models.ini storage") +def test_metadata_and_cost(): + """models.ini beats provider defaults beats unknown; cost sums per model.""" + from llamachat import models, providers + + table = providers.parse( + { + "providers": { + "local": {"base_url": "http://localhost:8181"}, + "together": { + "base_url": "https://api.example.org", + "api_key": "env:X", + "ctx_size": 32768, + "price_in": 0.6, + "price_out": 0.9, + }, + "free": {"base_url": "https://api3.example.org"}, + } + } + ) + + with tempfile.TemporaryDirectory() as tmp: + store = models.ModelStore(Path(tmp) / "models.ini") + store.save( + "together:specific", + models.ModelInfo(ctx_size=8192, vision=True, price_in=5.0), + ) + + # Layer 1: models.ini wins where it has a value. + info = models.resolve("together:specific", table, store) + assert info.ctx_size == 8192 + assert info.vision is True + assert info.price_in == 5.0 + # Layer 2 fills the gap models.ini left: price_out was never set. + assert info.price_out == 0.9 + + # Layer 2 alone for a model with no models.ini entry. + other = models.resolve("together:other", table, store) + assert other.ctx_size == 32768 + assert other.price_in == 0.6 + assert other.vision is None # layer 3: still unknown + + # Layer 3 throughout for a provider that configured nothing. + bare = models.resolve("free:anything", table, store) + assert bare.ctx_size is None + assert bare.price_in is None + + # A local model resolves against the local provider rather than + # falling through to an unrelated one: an unprefixed id and an + # unconfigured prefix both split to local, and local configured + # nothing here, so nothing may be inherited. + assert models.resolve("gemma4", table, store).ctx_size is None + # A models.ini entry is keyed by the full id, so a local model's + # own metadata must still come back when no provider supplies any. + store.save("gemma4", models.ModelInfo(ctx_size=4096)) + assert models.resolve("gemma4", table, store).ctx_size == 4096 + + # A config with cloud providers but no local one is legal, and + # split() still answers "local" for a bare id, so resolve() looks up + # a provider that is not in the table. The stored metadata must + # survive that, and nothing may be inherited from an unrelated + # provider: charging a local model together's prices would invent + # money that was never spent. + cloud_only = providers.parse( + { + "providers": { + "together": { + "base_url": "https://api.example.org", + "api_key": "env:X", + "ctx_size": 32768, + "price_in": 0.6, + } + } + } + ) + orphan = models.resolve("gemma4", cloud_only, store) + assert orphan.ctx_size == 4096 + assert orphan.price_in is None + assert models.is_billable("gemma4", cloud_only) is False + assert models.conversation_cost( + [{"model": "gemma4", "prompt_tokens": 1_000_000, + "completion_tokens": 1_000_000}], cloud_only, store + ) == 0.0 + + # Cost: prompt at the input rate, completion at the output rate. + rows = [ + {"model": "together:other", "prompt_tokens": 1_000_000, + "completion_tokens": 1_000_000}, + # A pre-migration row: no counts, no model. Contributes zero. + {"model": None, "prompt_tokens": None, "completion_tokens": None}, + # Counts but no model, which is what a reply interrupted before + # it recorded its model leaves behind. There is no rate to apply, + # so it must contribute zero rather than borrow another row's. + {"model": None, "prompt_tokens": 1_000_000, + "completion_tokens": 1_000_000}, + ] + assert models.conversation_cost(rows, table, store) == 1.5 + + # A mixed conversation prices each reply at what produced it. + mixed = [ + {"model": "together:other", "prompt_tokens": 1_000_000, + "completion_tokens": 0}, + {"model": "together:specific", "prompt_tokens": 1_000_000, + "completion_tokens": 0}, + ] + assert models.conversation_cost(mixed, table, store) == 5.6 + + # An unpriced model contributes nothing rather than guessing. + assert models.conversation_cost( + [{"model": "free:anything", "prompt_tokens": 1_000_000, + "completion_tokens": 0}], table, store + ) == 0.0 + + # A model priced on only one side still charges that side, rather + # than being all-or-nothing: together:specific has no price_out of + # its own but inherits one, so this uses a store-only provider. + assert models.conversation_cost( + [{"model": "free:half", "prompt_tokens": 1_000_000, + "completion_tokens": 1_000_000}], table, store + ) == 0.0 + store.save("free:half", models.ModelInfo(price_in=2.0)) + assert models.conversation_cost( + [{"model": "free:half", "prompt_tokens": 1_000_000, + "completion_tokens": 1_000_000}], table, store + ) == 2.0 + # The other side alone, so neither rate is quietly gated on the + # other being known. + store.save("free:outonly", models.ModelInfo(price_out=3.0)) + assert models.conversation_cost( + [{"model": "free:outonly", "prompt_tokens": 1_000_000, + "completion_tokens": 1_000_000}], table, store + ) == 3.0 + assert models.is_priced("free:outonly", table, store) is True + + # The projection prices input only: the reply's length is unknown + # until it arrives, so guessing it would overstate every turn. + assert models.projected_cost( + 1_000_000, "together:other", table, store + ) == 0.6 + assert models.projected_cost( + 1_000_000, "free:anything", table, store + ) == 0.0 + + # Whether a model can be priced at all decides ? versus blank. + assert models.is_priced("together:other", table, store) is True + assert models.is_priced("free:anything", table, store) is False + # Local is free, never unpriced. + assert models.is_billable("gemma4", table) is False + assert models.is_billable("free:anything", table) is False # no api_key + assert models.is_billable("together:other", table) is True + + # A local router behind an authenticating proxy is a legal config, + # and it is still free. Without this the assertion above passes + # only because the local provider happens to carry no api_key, + # which is what the is_local test actually exists to cover. + keyed_local = providers.parse( + { + "providers": { + "local": { + "base_url": "http://localhost:8181", + "api_key": "env:X", + "price_in": 9.0, + } + } + } + ) + assert models.is_billable("gemma4", keyed_local) is False + # A modelless row must stay free even when the local provider it + # would otherwise resolve to carries prices, which is what makes + # the empty-id guard in message_cost() load-bearing rather than + # decorative. + assert models.conversation_cost( + [{"model": None, "prompt_tokens": 1_000_000, + "completion_tokens": 1_000_000}], keyed_local, store + ) == 0.0 + + # Formatting: cents matter, so three decimals below a dollar. + assert models.format_cost(0.0) == "$0.000" + assert models.format_cost(1.5) == "$1.50" + assert models.format_cost(12.345) == "$12.35" + print("ok metadata resolution and cost") + + class _FakeResponse: """Enough of an http.client response for urlopen's context manager.""" @@ -2239,6 +2421,7 @@ if __name__ == "__main__": test_key_resolution() test_config_providers() test_models_store() + test_metadata_and_cost() test_search_tool_schema() test_search_results_sanitising() test_tool_call_accumulation() |