diff options
| -rw-r--r-- | llamachat/models.py | 188 | ||||
| -rwxr-xr-x | test_llamachat.py | 98 |
2 files changed, 286 insertions, 0 deletions
diff --git a/llamachat/models.py b/llamachat/models.py new file mode 100644 index 0000000..76be1d7 --- /dev/null +++ b/llamachat/models.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: GPL-2.0-only +# +# llamachat - a small native chat client for a local llama.cpp router +# Copyright (C) 2026 Danilo M. <danix@danix.xyz> +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License version 2 as +# published by the Free Software Foundation. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +"""Per-model metadata: what presets.ini cannot answer for a cloud model.""" + +import configparser +import os +from dataclasses import dataclass +from pathlib import Path + +# Recorded when the dialog is cancelled, so a model tried once never nags. +# The value is meaningless, only the key's presence is read: it exists to +# keep an otherwise empty section from being written with no keys at all, +# which some editors and merge tools drop. +SKIPPED = "offered" + + +@dataclass +class ModelInfo: + """What a model can tell us. Every field may be unknown.""" + + ctx_size: int | None = None + vision: bool | None = None + price_in: float | None = None + price_out: float | None = None + + def is_empty(self) -> bool: + return all( + v is None + for v in (self.ctx_size, self.vision, self.price_in, self.price_out) + ) + + +def _get(section, key, cast): + raw = section.get(key, "").strip() + if not raw: + return None + try: + return cast(raw) + except ValueError: + # A hand-edited "32k" or "3.7" for an int degrades to unknown rather + # than raising: the file is user-editable, and one bad line must not + # take down every other model's metadata with it. + return None + + +def _get_bool(section, key): + raw = section.get(key, "").strip().lower() + if raw in ("true", "yes", "1", "on"): + return True + if raw in ("false", "no", "0", "off"): + return False + return None + + +class ModelStore: + """models.ini, keyed by full model id. + + configparser rather than TOML because this file is written by the app, + and tomllib is read-only in the standard library. + """ + + def __init__(self, path: Path): + self.path = Path(path) + # Model ids contain colons and slashes, so no key/value delimiter + # may be inferred from a section name. Sections are safe as-is: + # brackets, colons, equals and percent signs all round-trip, and + # interpolation is off so a "%" in an id or value is literal. + self.parser = self._empty() + if self.path.exists(): + try: + self.parser.read(self.path, encoding="utf-8") + except (configparser.Error, OSError, UnicodeDecodeError): + # ponytail: a corrupt file reads as empty; the dialog can + # rewrite it. Failing to start over metadata is worse. The + # cost is that the next save silently discards whatever was + # readable, which is acceptable for a cache of prices. + self.parser = self._empty() + + @staticmethod + def _empty(): + # default_section is renamed away from "DEFAULT" because a real model + # id can never contain a newline but could plausibly be the literal + # string "DEFAULT", which add_section rejects. More importantly, a + # [DEFAULT] section in a hand-edited file leaks its keys into every + # other section, so get() would invent a ctx_size for models that + # have none. Reserving an id-shaped name no model can have keeps + # every section independent. + return configparser.ConfigParser( + interpolation=None, default_section="\x00llamachat-default" + ) + + def get(self, model_id: str) -> ModelInfo | None: + """Stored metadata, or None when there is none worth having. + + None and was_offered() True together mean "asked, learned nothing": + either the dialog was cancelled or every field was left blank. The + caller wants exactly that distinction, so it is not a trap: get() + answers "what do we know", was_offered() answers "should we ask". + """ + if not self.parser.has_section(model_id): + return None + section = self.parser[model_id] + info = ModelInfo( + ctx_size=_get(section, "ctx_size", int), + vision=_get_bool(section, "vision"), + price_in=_get(section, "price_in", float), + price_out=_get(section, "price_out", float), + ) + return None if info.is_empty() else info + + def was_offered(self, model_id: str) -> bool: + """Whether the dialog has already been shown for this model.""" + return self.parser.has_section(model_id) + + def save(self, model_id: str, info: ModelInfo) -> None: + section = self._section(model_id) + for key, value in ( + ("ctx_size", info.ctx_size), + ("vision", info.vision), + ("price_in", info.price_in), + ("price_out", info.price_out), + ): + if value is None: + section.pop(key, None) + elif isinstance(value, bool): + # Before the numeric branch: bool is a subclass of int, so + # str(False) would otherwise write "False" for vision. + section[key] = "true" if value else "false" + else: + section[key] = str(value) + # An all-blank save is still an answer, so the marker keeps the + # section from being written empty and re-read as never offered. + if any(k != SKIPPED for k in section): + section.pop(SKIPPED, None) + else: + section[SKIPPED] = "true" + self._write() + + def mark_skipped(self, model_id: str) -> None: + """Record a cancelled dialog: offered, declined, do not ask again. + + Never clears fields already stored. Cancelling a dialog opened over + a model we know about must not delete what we know. + """ + section = self._section(model_id) + if not any(k != SKIPPED for k in section): + section[SKIPPED] = "true" + self._write() + + def _section(self, model_id: str): + if not self.parser.has_section(model_id): + self.parser.add_section(model_id) + return self.parser[model_id] + + def _write(self) -> None: + """Rewrite the whole file from the in-memory parser. + + ponytail: last writer wins. Each store holds a full parser, so two + live instances would drop each other's sections rather than merge. + Not defended against because ipc.py enforces a single instance and + the window holds exactly one store, making a second writer a case + that cannot arise. The upgrade path, if the app ever grows a second + writer, is to reread before each save and merge. + """ + self.path.parent.mkdir(parents=True, exist_ok=True) + # Written to a sibling and renamed: os.replace is atomic within a + # directory, so a crash mid-write leaves the old file intact instead + # of a truncated one. Cheap here, and the alternative is losing every + # recorded price to one bad moment. + tmp = self.path.with_name(self.path.name + f".{os.getpid()}.tmp") + try: + with open(tmp, "w", encoding="utf-8") as fh: + self.parser.write(fh) + os.replace(tmp, self.path) + except OSError: + tmp.unlink(missing_ok=True) + raise diff --git a/test_llamachat.py b/test_llamachat.py index 29f305f..8b53d55 100755 --- a/test_llamachat.py +++ b/test_llamachat.py @@ -1380,6 +1380,103 @@ def test_config_providers(): print("ok config provider table") +def test_models_store(): + """models.ini round-trips per-model metadata and the cancel record.""" + from llamachat import models + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "models.ini" + store = models.ModelStore(path) + + # Nothing recorded yet. + assert store.get("together:Qwen/Qwen2.5") is None + assert store.was_offered("together:Qwen/Qwen2.5") is False + + store.save( + "together:Qwen/Qwen2.5", + models.ModelInfo( + ctx_size=32768, vision=False, price_in=1.2, price_out=1.2 + ), + ) + # A cancelled dialog records that it was offered, nothing more. + store.mark_skipped("together:Llama-Vision-Free") + + # Reread from disk, not from memory: this is the round trip. + fresh = models.ModelStore(path) + info = fresh.get("together:Qwen/Qwen2.5") + assert info.ctx_size == 32768 + assert info.vision is False + assert info.price_in == 1.2 + assert info.price_out == 1.2 + assert fresh.was_offered("together:Qwen/Qwen2.5") is True + + assert fresh.get("together:Llama-Vision-Free") is None + assert fresh.was_offered("together:Llama-Vision-Free") is True + + # Partial entries are legal: prices may be left blank. + fresh.save("together:cheap", models.ModelInfo(ctx_size=8192)) + again = models.ModelStore(path) + partial = again.get("together:cheap") + assert partial.ctx_size == 8192 + assert partial.price_in is None + assert partial.vision is None + + # A model id with a colon must survive being an ini section name. + again.save("together:org/name:v2", models.ModelInfo(ctx_size=4096)) + assert models.ModelStore(path).get("together:org/name:v2").ctx_size == 4096 + + # Real ids carry slashes, dots and mixed case. Section names are + # case sensitive, unlike keys, so the id must come back verbatim. + for real in ( + "together:Qwen/Qwen2.5-72B-Instruct-Turbo", + "together:deepseek-ai/DeepSeek-V3", + "together:meta-llama/Llama-3.3-70B", + ): + again.save(real, models.ModelInfo(ctx_size=128000, vision=True)) + back = models.ModelStore(path).get(real) + assert back.ctx_size == 128000, real + assert back.vision is True, real + + # vision=False must not be written as "False" and read back as None: + # bool is an int subclass, so the write order matters. + again.save("together:novision", models.ModelInfo(vision=False)) + assert models.ModelStore(path).get("together:novision").vision is False + + # Cancelling a dialog over a model we already know must not erase it. + again.mark_skipped("together:Qwen/Qwen2.5") + assert models.ModelStore(path).get("together:Qwen/Qwen2.5").ctx_size == 32768 + + # An all-blank save is still an answer: asked, learned nothing. It + # must not read back as never offered, or the dialog reopens forever. + # configparser does round-trip a keyless section, so the marker is + # belt and braces against an empty section being dropped by hand. + again.save("together:blank", models.ModelInfo()) + reread = models.ModelStore(path) + assert reread.get("together:blank") is None + assert reread.was_offered("together:blank") is True + + # A hand-edited file must degrade, not raise: "32k" is not an int. + # The [DEFAULT] value is deliberately a *valid* int, so this catches + # the leak itself rather than an unparseable value hiding it. + path.write_text( + "[DEFAULT]\nctx_size = 999\n\n" + "[together:junk]\nctx_size = 32k\nprice_in = free\n\n" + "[together:empty]\n", + encoding="utf-8", + ) + edited = models.ModelStore(path) + assert edited.get("together:junk") is None + assert edited.was_offered("together:junk") is True + # Would be ctx_size 999, inherited from [DEFAULT], if the store used + # configparser's real default section. + assert edited.get("together:empty") is None + + # A file that is not ini at all reads as empty rather than raising. + path.write_text("this is not an ini file\n", encoding="utf-8") + assert models.ModelStore(path).was_offered("together:junk") is False + print("ok models.ini storage") + + class _FakeResponse: """Enough of an http.client response for urlopen's context manager.""" @@ -2123,6 +2220,7 @@ if __name__ == "__main__": test_model_ids_and_filtering() test_key_resolution() test_config_providers() + test_models_store() test_search_tool_schema() test_search_results_sanitising() test_tool_call_accumulation() |