diff --git a/README.md b/README.md index 56eae9c..fc62bbf 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,7 @@ uv run pytest uv run gitmate --help uv run gitmate config set-key # stored in the OS credential store, never in a file uv run gitmate config show -uv run gitmate config set model gemini-flash +uv run gitmate config set model gemini-3.5-flash-lite uv run gitmate config set ignore_globs '*.gen.py,*.tmp' # comma-separated, empty clears uv run gitmate config set max_context_tokens 8000 # override model context limit uv run gitmate config set reserved_output_tokens 1500 # reserve headroom for model reply diff --git a/pyproject.toml b/pyproject.toml index f675997..0acf543 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "tomli-w>=1.0", "diskcache>=5.6", "tenacity>=8.2", + "httpx>=0.27", "google-genai>=0.3", # default provider (Gemini 3 Flash / Flash-Lite) "anthropic>=0.34", # secondary provider, same LLMProvider interface ] diff --git a/src/gitmate/cli.py b/src/gitmate/cli.py index 40b4741..4c08871 100644 --- a/src/gitmate/cli.py +++ b/src/gitmate/cli.py @@ -2,6 +2,8 @@ from __future__ import annotations +import dataclasses + import typer from rich.console import Console from rich.table import Table @@ -123,38 +125,68 @@ def config_show() -> None: table.add_column("field") table.add_column("value") table.add_row("config file", str(config_mod.config_path())) - table.add_row("model", cfg.model) - table.add_row("commit_style", cfg.commit_style) - table.add_row( - "budget_cap_usd", str(cfg.budget_cap_usd) if cfg.budget_cap_usd is not None else "none" - ) - table.add_row("ignore_globs", ", ".join(cfg.ignore_globs) if cfg.ignore_globs else "none") + for f in dataclasses.fields(cfg): + val = getattr(cfg, f.name) + if val is None: + display_val = "none" + elif isinstance(val, list): + display_val = ", ".join(val) if val else "none" + else: + display_val = str(val) + table.add_row(f.name, display_val) table.add_row("api_key", "set" if key_set else "not set") console.print(table) @config_app.command("set") def config_set(field: str, value: str) -> None: - """Set one field: model, commit_style, budget_cap_usd, ignore_globs.""" + """Set any config field: model, commit_style, budget_cap_usd, ignore_globs, cache_dir, etc.""" try: cfg = config_mod.load_config() except config_mod.ConfigError as exc: console.print(f"[red]error:[/red] {exc}") raise typer.Exit(1) from None - match field: - case "model": - cfg.model = value - case "commit_style": - cfg.commit_style = value - case "budget_cap_usd": - cfg.budget_cap_usd = None if value == "" else _parse_budget(value) - case "ignore_globs": - cfg.ignore_globs = [g.strip() for g in value.split(",") if g.strip()] - case _: - raise typer.BadParameter( - f"unknown field {field!r}; expected model, commit_style, " - "budget_cap_usd, or ignore_globs." - ) + + field_map = {f.name: f for f in dataclasses.fields(config_mod.GitmateConfig)} + if field not in field_map: + valid_fields = ", ".join(field_map.keys()) + raise typer.BadParameter(f"unknown field {field!r}; expected one of: {valid_fields}.") + + if field == "ignore_globs": + cfg.ignore_globs = [g.strip() for g in value.split(",") if g.strip()] + elif field == "budget_cap_usd": + cfg.budget_cap_usd = None if value in ("", "none") else _parse_budget(value) + elif field in ("max_context_tokens", "reserved_output_tokens", "template_overhead"): + if field == "max_context_tokens" and value in ("", "none"): + cfg.max_context_tokens = None + else: + try: + int_val = int(value) + if int_val <= 0: + raise ValueError + except ValueError: + raise typer.BadParameter(f"'{field}' must be a positive integer.") from None + setattr(cfg, field, int_val) + elif field == "cache_dir": + cfg.cache_dir = None if value in ("", "none") else value.strip() or None + elif field in ("model", "commit_style"): + stripped = value.strip() + if not stripped: + raise typer.BadParameter(f"'{field}' must not be empty.") + setattr(cfg, field, stripped) + else: + setattr(cfg, field, value) + + # Validate overall config integrity + max_ctx = cfg.max_context_tokens + reserved = cfg.reserved_output_tokens + overhead = cfg.template_overhead + if max_ctx is not None and max_ctx <= (reserved + overhead): + raise typer.BadParameter( + f"'max_context_tokens' ({max_ctx}) must be greater than " + f"'reserved_output_tokens' + 'template_overhead' ({reserved + overhead})." + ) + config_mod.save_config(cfg) console.print(f"[green]Set {field}.[/green]") diff --git a/src/gitmate/config.py b/src/gitmate/config.py index 24f20da..57f0571 100644 --- a/src/gitmate/config.py +++ b/src/gitmate/config.py @@ -16,15 +16,18 @@ API_KEY_ACCOUNT = "api-key" CONFIG_ENV_VAR = "GITMATE_CONFIG_DIR" -# Opaque default; exact provider model IDs are verified in Phase 3. -DEFAULT_MODEL = "gemini-flash" +# Default model ID. This is the literal google-genai SDK identifier. Gemini +# model IDs churn fast (2.0 Flash is retired; Flash-Lite previews deprecate +# within months), so keep this configurable and re-check before assuming it +# still resolves. +DEFAULT_MODEL = "gemini-3.5-flash-lite" DEFAULT_COMMIT_STYLE = "conventional" #: Context window limits for known models; users override via `max_context_tokens`. DEFAULT_MODEL_CONTEXT_WINDOWS: dict[str, int] = { - "gemini-flash": 1_048_576, - "gemini-3-flash": 1_048_576, "gemini-3.5-flash-lite": 1_048_576, + "gemini-3-flash": 1_048_576, + "gemini-flash": 1_048_576, "claude-3-5-sonnet": 200_000, "claude-3-5-haiku": 200_000, } @@ -51,6 +54,7 @@ class GitmateConfig: max_context_tokens: int | None = None reserved_output_tokens: int = DEFAULT_RESERVED_OUTPUT_TOKENS template_overhead: int = DEFAULT_TEMPLATE_OVERHEAD + cache_dir: str | None = None class SecretStore(Protocol): @@ -175,6 +179,14 @@ def load_config(path: Path | None = None) -> GitmateConfig: f"than 'reserved_output_tokens' + 'template_overhead' ({reserved + overhead})" ) + cache_dir_raw = raw.get("cache_dir") + if cache_dir_raw is None: + cache_dir: str | None = None + elif not isinstance(cache_dir_raw, str) or not cache_dir_raw: + raise ConfigError(f"cannot parse {resolved}: 'cache_dir' must be a non-empty string") + else: + cache_dir = cache_dir_raw + return GitmateConfig( model=model, commit_style=style, @@ -183,6 +195,7 @@ def load_config(path: Path | None = None) -> GitmateConfig: max_context_tokens=max_ctx, reserved_output_tokens=reserved, template_overhead=overhead, + cache_dir=cache_dir, ) @@ -201,4 +214,6 @@ def save_config(cfg: GitmateConfig, path: Path | None = None) -> None: data["reserved_output_tokens"] = cfg.reserved_output_tokens if cfg.template_overhead != DEFAULT_TEMPLATE_OVERHEAD: data["template_overhead"] = cfg.template_overhead + if cfg.cache_dir is not None: + data["cache_dir"] = cfg.cache_dir resolved.write_text(tomli_w.dumps(data), encoding="utf-8") diff --git a/src/gitmate/providers/base.py b/src/gitmate/providers/base.py new file mode 100644 index 0000000..a7c33fa --- /dev/null +++ b/src/gitmate/providers/base.py @@ -0,0 +1,33 @@ +"""Provider abstraction: LLMProvider protocol, response types, errors.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + + +class ProviderUnavailable(Exception): + """Raised when the LLM backend cannot produce a result after retries.""" + + +class RetryableProviderError(Exception): + """Internal marker for transient failures worth retrying before giving up.""" + + +@dataclass +class LLMResponse: + """Text plus optional token usage from one provider call.""" + + text: str + model: str + input_tokens: int | None = None + output_tokens: int | None = None + + +@runtime_checkable +class LLMProvider(Protocol): + """Generates text; every backend can be swapped behind this boundary.""" + + def generate(self, prompt: str, model: str) -> LLMResponse: + """Produce a response for the prompt using the given model.""" + ... diff --git a/src/gitmate/providers/cache.py b/src/gitmate/providers/cache.py new file mode 100644 index 0000000..5c20800 --- /dev/null +++ b/src/gitmate/providers/cache.py @@ -0,0 +1,63 @@ +"""Caching wrapper that short-circuits provider calls on a hash key.""" + +from __future__ import annotations + +import hashlib +import os +from pathlib import Path + +from diskcache import Cache # type: ignore[import-untyped] + +from gitmate.providers.base import LLMProvider, LLMResponse + +#: Bumped in Phase 4 when real prompt templates land; kept separate so those +#: templates can thread the same key without migration. Different templates = +#: different (and intentionally colder) cache entries. +DEFAULT_TEMPLATE_VERSION = "1" + + +def default_cache_dir() -> Path: + """XDG cache dir for gitmate, overridable via config.cache_dir.""" + override = os.environ.get("GITMATE_CACHE_DIR") + if override: + return Path(override) + base = os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache") + return Path(base) / "gitmate" + + +def cache_key(prompt: str, model: str, template_version: str) -> str: + """Stable hash of the inputs that change a response's correctness.""" + material = f"{template_version}\x00{model}\x00{prompt}" + return hashlib.sha256(material.encode("utf-8")).hexdigest() + + +class CachedProvider: + """LLMProvider that consults diskcache before calling the wrapped provider.""" + + def __init__( + self, + provider: LLMProvider, + cache_dir: Path | str | None = None, + ttl: float | None = None, + template_version: str = DEFAULT_TEMPLATE_VERSION, + ) -> None: + self._provider = provider + resolved = str(cache_dir) if cache_dir is not None else str(default_cache_dir()) + self._cache = Cache(resolved) + self._ttl = ttl + self._template_version = template_version + + def generate(self, prompt: str, model: str) -> LLMResponse: + """Return a cached response when present, else generate and store.""" + key = cache_key(prompt, model, self._template_version) + cached: LLMResponse | None = self._cache.get(key, default=None) + if cached is not None: + return cached + response = self._provider.generate(prompt, model) + if response.text.strip(): + self._cache.set(key, response, expire=self._ttl) + return response + + def close(self) -> None: + """Close the backing cache.""" + self._cache.close() diff --git a/src/gitmate/providers/gemini.py b/src/gitmate/providers/gemini.py new file mode 100644 index 0000000..7ded756 --- /dev/null +++ b/src/gitmate/providers/gemini.py @@ -0,0 +1,84 @@ +"""Gemini provider over the google-genai SDK, with retry on transient errors.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential + +from gitmate.providers.base import LLMResponse, ProviderUnavailable, RetryableProviderError + +if TYPE_CHECKING: + from google import genai + + +def _is_transient(exc: BaseException) -> bool: + """True for failures worth retrying: 5xx, 429 rate limits, transport.""" + from google.genai.errors import ClientError, ServerError + + if isinstance(exc, ServerError): + return True + if isinstance(exc, ClientError): + return getattr(exc, "code", None) == 429 + import httpx + + return isinstance(exc, httpx.TransportError) + + +class GeminiProvider: + """LLMProvider backed by Google's Gemini models via google-genai.""" + + def __init__(self, api_key: str | None = None, client: genai.Client | None = None) -> None: + self._client = client + self._api_key = api_key + + def _get_client(self) -> genai.Client: + if self._client is not None: + return self._client + if not self._api_key: + raise ProviderUnavailable( + "Gemini API key is required. Run 'gitmate config set-key' first." + ) + from google import genai + + self._client = genai.Client(api_key=self._api_key) + return self._client + + def generate(self, prompt: str, model: str) -> LLMResponse: + """Generate text for prompt, retrying transient failures then raising.""" + try: + return self._generate_with_retry(self._get_client(), prompt, model) + except RetryableProviderError as exc: + raise ProviderUnavailable(f"Gemini unavailable after retries: {exc}") from exc + + @retry( + retry=retry_if_exception_type(RetryableProviderError), + stop=stop_after_attempt(3), + wait=wait_exponential(multiplier=1, min=1, max=8), + reraise=True, + ) + def _generate_with_retry(self, client: genai.Client, prompt: str, model: str) -> LLMResponse: + try: + resp = client.models.generate_content(model=model, contents=prompt) + except Exception as exc: + if _is_transient(exc): + raise RetryableProviderError(str(exc)) from exc + raise ProviderUnavailable(f"Gemini request failed: {exc}") from exc + text = (resp.text or "").strip() + if not text: + finish_reason = None + candidates = getattr(resp, "candidates", None) + if candidates and len(candidates) > 0: + finish_reason = getattr(candidates[0], "finish_reason", None) + if finish_reason is not None: + reason_name = getattr(finish_reason, "name", str(finish_reason)) + raise ProviderUnavailable( + f"Gemini declined to respond (finish_reason: {reason_name})." + ) + raise ProviderUnavailable("Gemini returned an empty response.") + return LLMResponse( + text=text, + model=model, + input_tokens=getattr(resp.usage_metadata, "prompt_token_count", None), + output_tokens=getattr(resp.usage_metadata, "candidates_token_count", None), + ) diff --git a/src/gitmate/token_budget.py b/src/gitmate/token_budget.py index af1a448..1c33e84 100644 --- a/src/gitmate/token_budget.py +++ b/src/gitmate/token_budget.py @@ -10,6 +10,8 @@ DEFAULT_MODEL, DEFAULT_RESERVED_OUTPUT_TOKENS, DEFAULT_TEMPLATE_OVERHEAD, + GitmateConfig, + get_context_window, ) from gitmate.diff_extractor import FileDiff @@ -166,6 +168,21 @@ def __init__( self.reserved_output_tokens = reserved_output_tokens self.default_template_overhead = default_template_overhead + @classmethod + def from_config(cls, cfg: GitmateConfig, counter: TokenCounter) -> TokenBudgetManager: + """Build a manager whose context window comes from the configured model. + + Replaces the constructor's 1M default so non-Gemini models (e.g. a + 200k Claude window) are never silently over-budgeted. + """ + return cls( + counter=counter, + model=cfg.model, + context_window=get_context_window(cfg), + reserved_output_tokens=cfg.reserved_output_tokens, + default_template_overhead=cfg.template_overhead, + ) + def count_diff_tokens(self, diffs: list[FileDiff]) -> int: """Calculate token count for a list of FileDiffs.""" if not diffs: diff --git a/tests/test_cli.py b/tests/test_cli.py index 386e992..d4cee13 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -174,3 +174,44 @@ def test_debug_diff_git_command_error_clean_error(tmp_config_dir: Path, git_repo assert result.exit_code == 1 assert "error:" in result.output assert "Traceback" not in result.output + + +def test_config_set_and_show_all_dataclass_fields( + tmp_config_dir: Path, mem_store: InMemorySecretStore +) -> None: + runner.invoke(cli.app, ["config", "set", "max_context_tokens", "16000"]) + runner.invoke(cli.app, ["config", "set", "reserved_output_tokens", "1000"]) + runner.invoke(cli.app, ["config", "set", "template_overhead", "400"]) + runner.invoke(cli.app, ["config", "set", "cache_dir", "/custom/cache"]) + + cfg = load_config() + assert cfg.max_context_tokens == 16000 + assert cfg.reserved_output_tokens == 1000 + assert cfg.template_overhead == 400 + assert cfg.cache_dir == "/custom/cache" + + res = runner.invoke(cli.app, ["config", "show"]) + assert res.exit_code == 0 + assert "max_context_tokens" in res.output + assert "16000" in res.output + assert "cache_dir" in res.output + assert "/custom/cache" in res.output + + +def test_config_set_budget_fields_validation(tmp_config_dir: Path) -> None: + res = runner.invoke(cli.app, ["config", "set", "max_context_tokens", "0"]) + assert res.exit_code != 0 + assert "positive integer" in res.output + + res = runner.invoke(cli.app, ["config", "set", "max_context_tokens", "--", "-5"]) + assert res.exit_code != 0 + assert "positive integer" in res.output + + res = runner.invoke(cli.app, ["config", "set", "reserved_output_tokens", "not-a-number"]) + assert res.exit_code != 0 + assert "positive integer" in res.output + + # max_context_tokens <= reserved + overhead (default 2048 + 500 = 2548) + res = runner.invoke(cli.app, ["config", "set", "max_context_tokens", "2000"]) + assert res.exit_code != 0 + assert "must be greater than" in res.output diff --git a/tests/test_config.py b/tests/test_config.py index 5978e14..8d4e27a 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -74,7 +74,20 @@ def test_ignore_globs_wrong_type_raises(tmp_config_dir: Path) -> None: load_config() +def test_cache_dir_round_trip(tmp_config_dir: Path) -> None: + save_config(GitmateConfig(cache_dir="/tmp/gitmate-cache")) + assert load_config().cache_dir == "/tmp/gitmate-cache" + + +def test_cache_dir_wrong_type_raises(tmp_config_dir: Path) -> None: + config_path().parent.mkdir(parents=True, exist_ok=True) + config_path().write_text('cache_dir = ""\n', encoding="utf-8") + with pytest.raises(ConfigError, match="'cache_dir' must be a non-empty string"): + load_config() + + def test_get_context_window_known_models() -> None: + assert get_context_window(GitmateConfig(model="gemini-3.5-flash-lite")) == 1_048_576 assert get_context_window(GitmateConfig(model="gemini-flash")) == 1_048_576 assert get_context_window(GitmateConfig(model="claude-3-5-sonnet")) == 200_000 diff --git a/tests/test_providers.py b/tests/test_providers.py new file mode 100644 index 0000000..01db236 --- /dev/null +++ b/tests/test_providers.py @@ -0,0 +1,203 @@ +"""Tests for the provider abstraction, cache wrapper, and Gemini provider.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import pytest +from tenacity import wait_none + +from gitmate.providers.base import LLMProvider, LLMResponse, ProviderUnavailable +from gitmate.providers.cache import CachedProvider, cache_key +from gitmate.providers.gemini import GeminiProvider + + +@pytest.fixture(autouse=True) +def fast_retries(monkeypatch: pytest.MonkeyPatch) -> None: + """Eliminate tenacity exponential backoff delays during tests.""" + retry_obj = GeminiProvider._generate_with_retry.retry # type: ignore[attr-defined] + monkeypatch.setattr(retry_obj, "wait", wait_none()) + + +class FakeProvider: + """In-memory LLMProvider recording every call it actually receives.""" + + def __init__(self, text: str = "generated") -> None: + self.text = text + self.calls: list[tuple[str, str]] = [] + + def generate(self, prompt: str, model: str) -> LLMResponse: + self.calls.append((prompt, model)) + return LLMResponse(text=self.text, model=model) + + +# --- cache --- + + +def test_cache_consulted_before_generate(tmp_path: Path) -> None: + provider = FakeProvider() + cached = CachedProvider(provider, cache_dir=tmp_path, ttl=None) + + first = cached.generate("prompt", "gemini-3.5-flash-lite") + second = cached.generate("prompt", "gemini-3.5-flash-lite") + + assert first == second + # The whole point of caching: one underlying call, not two. + assert len(provider.calls) == 1 + cached.close() + + +def test_cache_miss_then_hit_across_prompts(tmp_path: Path) -> None: + provider = FakeProvider() + cached = CachedProvider(provider, cache_dir=tmp_path) + cached.generate("a", "m") + cached.generate("b", "m") + assert len(provider.calls) == 2 + cached.close() + + +def test_cache_key_changes_with_model_and_template_version() -> None: + base = cache_key("p", "m", "1") + assert cache_key("p", "m", "2") != base + assert cache_key("p", "other", "1") != base + assert cache_key("q", "m", "1") != base + assert cache_key("p", "m", "1") == base + + +def test_template_version_bump_misses_old_entry(tmp_path: Path) -> None: + provider = FakeProvider() + v1 = CachedProvider(provider, cache_dir=tmp_path, template_version="1") + v1.generate("p", "m") + assert len(provider.calls) == 1 + v1.close() + + v2 = CachedProvider(FakeProvider("v2"), cache_dir=tmp_path, template_version="2") + resp = v2.generate("p", "m") + assert resp.text == "v2" + v2.close() + + +def test_cache_dir_override(tmp_path: Path) -> None: + custom = tmp_path / "custom" + provider = FakeProvider() + cached = CachedProvider(provider, cache_dir=custom) + cached.generate("p", "m") + assert (custom / "cache.db").exists() + cached.close() + + +# --- Gemini provider --- + + +def _resp(text: str | None, finish_reason: Any = None) -> MagicMock: + r = MagicMock() + r.text = text + meta = MagicMock() + meta.prompt_token_count = 10 + meta.candidates_token_count = 5 + r.usage_metadata = meta + if finish_reason is not None: + cand = MagicMock() + cand.finish_reason = finish_reason + r.candidates = [cand] + else: + r.candidates = [] + return r + + +def test_cache_does_not_store_empty_response(tmp_path: Path) -> None: + provider = FakeProvider(text="") + cached = CachedProvider(provider, cache_dir=tmp_path) + resp1 = cached.generate("p", "m") + assert resp1.text == "" + assert len(provider.calls) == 1 + + # Empty response should not be cached; second call hits the provider again + provider.text = "now valid" + resp2 = cached.generate("p", "m") + assert resp2.text == "now valid" + assert len(provider.calls) == 2 + cached.close() + + +def test_gemini_provider_success() -> None: + client = MagicMock() + client.models.generate_content.return_value = _resp(" hello ") + provider = GeminiProvider(client=client) + resp = provider.generate("write a message", "gemini-3.5-flash-lite") + assert resp.text == "hello" + assert resp.model == "gemini-3.5-flash-lite" + assert resp.input_tokens == 10 + assert resp.output_tokens == 5 + client.models.generate_content.assert_called_once_with( + model="gemini-3.5-flash-lite", contents="write a message" + ) + + +def test_gemini_provider_empty_response_raises() -> None: + client = MagicMock() + client.models.generate_content.return_value = _resp(None) + provider = GeminiProvider(client=client) + with pytest.raises(ProviderUnavailable, match="Gemini returned an empty response"): + provider.generate("prompt", "gemini-3.5-flash-lite") + + +def test_gemini_provider_safety_declined_raises_with_finish_reason() -> None: + client = MagicMock() + reason = MagicMock() + reason.name = "SAFETY" + client.models.generate_content.return_value = _resp("", finish_reason=reason) + provider = GeminiProvider(client=client) + with pytest.raises( + ProviderUnavailable, match=r"Gemini declined to respond \(finish_reason: SAFETY\)" + ): + provider.generate("prompt", "gemini-3.5-flash-lite") + + +def test_gemini_provider_no_api_key_or_client() -> None: + with pytest.raises(ProviderUnavailable, match="API key is required"): + GeminiProvider().generate("prompt", "gemini-3.5-flash-lite") + + +def test_gemini_provider_retries_transient_then_succeeds() -> None: + from google.genai.errors import ServerError + + client = MagicMock() + client.models.generate_content.side_effect = [ + ServerError(500, None, None), + _resp("ok"), + ] + provider = GeminiProvider(client=client) + assert provider.generate("p", "m").text == "ok" + assert client.models.generate_content.call_count == 2 + + +def test_gemini_provider_exhausts_retries() -> None: + from google.genai.errors import ServerError + + client = MagicMock() + client.models.generate_content.side_effect = ServerError(500, None, None) + provider = GeminiProvider(client=client) + with pytest.raises(ProviderUnavailable, match="unavailable after retries"): + provider.generate("p", "m") + assert client.models.generate_content.call_count == 3 + + +def test_gemini_provider_non_transient_raises_without_retry() -> None: + from google.genai.errors import ClientError + + client = MagicMock() + client.models.generate_content.side_effect = ClientError(400, None, None) + provider = GeminiProvider(client=client) + with pytest.raises(ProviderUnavailable, match="request failed"): + provider.generate("p", "m") + assert client.models.generate_content.call_count == 1 + + +def test_is_like_provider_protocol(tmp_path: Path) -> None: + assert isinstance(GeminiProvider(), LLMProvider) + cached = CachedProvider(FakeProvider(), cache_dir=tmp_path) + assert isinstance(cached, LLMProvider) + cached.close() diff --git a/tests/test_token_budget.py b/tests/test_token_budget.py index ddbe34a..abdb635 100644 --- a/tests/test_token_budget.py +++ b/tests/test_token_budget.py @@ -6,6 +6,7 @@ import pytest +from gitmate.config import GitmateConfig, UnknownModelError from gitmate.diff_extractor import FileDiff from gitmate.token_budget import ( BudgetDecision, @@ -518,3 +519,24 @@ def test_mixed_truncated_and_guarded_diffs_trigger_needs_chunking() -> None: assert [d.path for d in decision.omitted_diffs] == ["big.py"] # Ratio format: 1 of 3 files truncated, NOT "even with all 1 files truncated" assert "even with 1/3 files truncated; requires chunking" in decision.summary_note + + +def test_from_config_uses_model_context_window() -> None: + counter = FakeTokenCounter() + gemini = TokenBudgetManager.from_config(GitmateConfig(model="gemini-3.5-flash-lite"), counter) + assert gemini.context_window == 1_048_576 + assert gemini.model == "gemini-3.5-flash-lite" + + claude = TokenBudgetManager.from_config(GitmateConfig(model="claude-3-5-sonnet"), counter) + assert claude.context_window == 200_000 + + +def test_from_config_override_and_unknown_model() -> None: + counter = FakeTokenCounter() + overridden = TokenBudgetManager.from_config( + GitmateConfig(model="anything", max_context_tokens=32_000), counter + ) + assert overridden.context_window == 32_000 + + with pytest.raises(UnknownModelError): + TokenBudgetManager.from_config(GitmateConfig(model="no-such-model"), counter) diff --git a/uv.lock b/uv.lock index d0e4d75..330ffe2 100644 --- a/uv.lock +++ b/uv.lock @@ -604,6 +604,7 @@ dependencies = [ { name = "anthropic" }, { name = "diskcache" }, { name = "google-genai" }, + { name = "httpx" }, { name = "keyring" }, { name = "rich" }, { name = "tenacity" }, @@ -624,6 +625,7 @@ requires-dist = [ { name = "anthropic", specifier = ">=0.34" }, { name = "diskcache", specifier = ">=5.6" }, { name = "google-genai", specifier = ">=0.3" }, + { name = "httpx", specifier = ">=0.27" }, { name = "keyring", specifier = ">=25.0" }, { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.11" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, @@ -765,7 +767,7 @@ name = "importlib-metadata" version = "9.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "zipp" }, + { name = "zipp", marker = "(python_full_version < '3.14' and sys_platform != 'emscripten') or (python_full_version < '3.12' and sys_platform == 'emscripten')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/7e/1e7e8dc30634b93ebb3d58a3dea569ad146e656218d3960ab04f62047b29/importlib_metadata-9.0.1.tar.gz", hash = "sha256:ab830580bc0ef3db61ce8fae716389e5462b67e033018bab6d8f80ef17172f99", size = 59124, upload-time = "2026-08-28T15:30:34.646Z" } wheels = [ @@ -1425,8 +1427,8 @@ name = "secretstorage" version = "3.5.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cryptography" }, - { name = "jeepney" }, + { name = "cryptography", marker = "python_full_version < '3.12' or python_full_version >= '3.15' or sys_platform != 'emscripten'" }, + { name = "jeepney", marker = "python_full_version < '3.12' or python_full_version >= '3.15' or sys_platform != 'emscripten'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } wheels = [