diff --git a/README.md b/README.md index 7ebed3bc..162b1d96 100644 --- a/README.md +++ b/README.md @@ -2,12 +2,12 @@ > **Before an AI writes a new class/id/function, CodeLens must be checked. This is not optional.** -CodeLens is an AI-native code intelligence platform that gives AI agents **full visibility** into a codebase before they write any code. It prevents collision, overwrite of existing logic, security vulnerabilities, and dead code through 70 CLI commands, an MCP server with 68 tools (54 static + 14 dynamic), AST-based taint analysis, live CVE/OSV scanning, a plugin system with OWASP Top 10 + Compliance rule packs, a true graph data model (nodes + edges) for structural code queries, and token-efficient `--format compact` output for high-volume agent workflows (issue #17). +CodeLens is an AI-native code intelligence platform that gives AI agents **full visibility** into a codebase before they write any code. It prevents collision, overwrite of existing logic, security vulnerabilities, and dead code through 71 CLI commands, an MCP server with 69 tools (54 static + 15 dynamic), AST-based taint analysis, live CVE/OSV scanning, a plugin system with OWASP Top 10 + Compliance rule packs, a true graph data model (nodes + edges) for structural code queries, and token-efficient `--format compact` output for high-volume agent workflows (issue #17). ## Features -- **70 CLI Commands** — From basic scan/query to AST taint analysis, CVE scanning, plugin management, auto-fix, dashboards, CI/CD quality gates, and `graph-schema` for cheap graph-shape introspection -- **MCP Server (68 Tools)** — Native AI agent integration via Model Context Protocol (JSON-RPC over stdio), 54 statically-defined tools + 14 dynamically discovered, every tool accepts a `format` parameter (`json`/`markdown`/`ai`/`sarif`/`compact`) +- **71 CLI Commands** — From basic scan/query to AST taint analysis, CVE scanning, plugin management, auto-fix, dashboards, CI/CD quality gates, and `graph-schema` for cheap graph-shape introspection +- **MCP Server (69 Tools)** — Native AI agent integration via Model Context Protocol (JSON-RPC over stdio), 54 statically-defined tools + 15 dynamically discovered, every tool accepts a `format` parameter (`json`/`markdown`/`ai`/`sarif`/`compact`) - **Token-Efficient Compact Output (v8.2, issue #17)** — `--format compact` produces single-char-key JSON with abbreviated types, omitted null fields, and relative paths — ~50% smaller than `json` on real trace output. Combined with `--limit`/`--offset` pagination, 5 structural queries now cost <5k tokens (down from 30-80k) - **AST Taint Engine** — Tree-sitter based taint analysis with return-value propagation, scope hierarchy, and branch condition refinement - **Live CVE/OSV Scanning** — Real-time vulnerability data from OSV.dev API with SQLite cache, 9 ecosystems (PyPI, npm, crates.io, Go, Maven, NuGet, RubyGems, Pub, Hex) @@ -225,8 +225,8 @@ codelens/ │ ├── changelog.md # Older changelog (per-version highlights) │ └── agent-integration.md # AI agent integration guide ├── scripts/ -│ ├── codelens.py # CLI entry point (70 commands registered) -│ ├── mcp_server.py # MCP JSON-RPC server (68 tools) +│ ├── codelens.py # CLI entry point (71 commands registered) +│ ├── mcp_server.py # MCP JSON-RPC server (69 tools) │ ├── registry.py # Registry read/write/build │ ├── persistent_registry.py # SQLite persistent storage (WAL mode) │ ├── base_parser.py # Base tree-sitter parser diff --git a/SKILL-QUICK.md b/SKILL-QUICK.md index f7944813..62c04593 100755 --- a/SKILL-QUICK.md +++ b/SKILL-QUICK.md @@ -114,7 +114,7 @@ $CLI list --limit 5 --offset 10 --format compact # → paginated + co | "Cross-file taint" | `dataflow` | `taint` (taint is single-file, AST-deep) | | "Auto-fix issues" | `fix` | `check` (check just gates, doesn't fix) | -## All 70 Commands +## All 71 Commands ### Setup & Lifecycle (8+) `init` · `scan [--incremental] [--max-files N] [--full]` · `registry-validate` · `detect` · `watch [--debounce SECS] [--git-mode] [--interval SECS]` · `git-status` · `migrate` · `serve` · `lsp-status` (issue #33: `codelens --lsp-status` top-level flag is an alias of `codelens lsp-status` — both delegate to `hybrid_engine.get_lsp_status()` and return the identical payload) @@ -146,9 +146,9 @@ $CLI list --limit 5 --offset 10 --format compact # → paginated + co ### Tooling (1) `plugin ` -**Total: 70 commands** (auto-registered via `commands/__init__.py`; rerun `python3 scripts/sync_command_count.py --apply` after adding/removing a command) +**Total: 71 commands** (auto-registered via `commands/__init__.py`; rerun `python3 scripts/sync_command_count.py --apply` after adding/removing a command) -## MCP Server (68 Tools) +## MCP Server (69 Tools) Start the MCP server for AI agent integration: @@ -156,9 +156,9 @@ Start the MCP server for AI agent integration: python3 scripts/codelens.py serve ``` -Exposes 68 tools as `codelens_` (e.g., `codelens_query`, `codelens_taint`, `codelens_graph_schema`, `codelens_architecture`, `codelens_resolve_types`, `codelens_git_status`): +Exposes 69 tools as `codelens_` (e.g., `codelens_query`, `codelens_taint`, `codelens_graph_schema`, `codelens_architecture`, `codelens_resolve_types`, `codelens_git_status`): - 50 statically-defined tools (full JSON schemas in `mcp_server.py`) -- 14 dynamically-discovered tools (auto-discovered from `COMMAND_REGISTRY`; long-running `watch` and `serve` are excluded) +- 15 dynamically-discovered tools (auto-discovered from `COMMAND_REGISTRY`; long-running `watch` and `serve` are excluded) - Every tool accepts a `format` parameter (`json`/`markdown`/`ai`/`sarif`/`compact`). Use `format: "compact"` for token-efficient responses (~50% smaller than `json`). - `watch` and `serve` itself are excluded (long-running) diff --git a/SKILL.md b/SKILL.md index 77b8be13..64a3ca83 100755 --- a/SKILL.md +++ b/SKILL.md @@ -1,10 +1,10 @@ --- name: codelens description: > - CodeLens — AI-Native Code Intelligence. 70 commands for AI-powered code analysis, + CodeLens — AI-Native Code Intelligence. 71 commands for AI-powered code analysis, security auditing, quality scoring, AST-based taint analysis, live CVE scanning, and pre-write safety checks. Supports 28+ languages with tree-sitter + regex - fallback parsing. MCP server exposes 68 tools for AI agent integration. + fallback parsing. MCP server exposes 69 tools for AI agent integration. For quick command reference with validated output schemas, see SKILL-QUICK.md. For version history, see CHANGELOG.md. --- diff --git a/docs/llm-framework.md b/docs/llm-framework.md new file mode 100644 index 00000000..0e7a2d2e --- /dev/null +++ b/docs/llm-framework.md @@ -0,0 +1,249 @@ +# LLM Integration Framework (issue #63 Phase 1) + +> **Status:** Phase 1 shipped. Phases 2–5 are tracked as follow-up issues. +> **Last updated:** 2026-07-02 + +## Alasan Dibuat + +CodeLens punya banyak fitur yang secara alami bisa di-enhance dengan LLM +reasoning: taint validation (FP check pada path yang dilaporkan sebagai +vulnerable), secret false-positive check, smell justification, dead-code +reason, dan bug explanation. Sebelum Phase 1, setiap fitur yang ingin +memakai LLM harus mengimplementasikan dispatch + retry + timeout-nya +sendiri — duplikasi yang error-prone. + +Phase 1 membangun abstraksi tunggal: + +- 6 provider (OpenAI / Anthropic / Bedrock / Google / DeepSeek / Z.ai GLM) +- Lazy import per provider (SDK hanya di-import saat benar-benar dipakai) +- 60s timeout, 3-retry exponential backoff (1s → 2s → 4s) +- Config via env vars (`CODELENS_LLM_PROVIDER`, `CODELENS_LLM_MODEL`, + `CODELENS_LLM_API_KEY`) atau explicit kwargs +- Error model yang membedakan retryable vs non-retryable + +## Arsitektur + +``` +scripts/llm/ +├── __init__.py # Re-exports public API +├── base_tool.py # LLMTool ABC + LLMToolInput/Output ABCs + errors +└── provider.py # invoke_llm() + resolve_provider() + 6 _call_* wrappers + +scripts/commands/ +└── llm_framework.py # `codelens llm providers|config|ping` command + # (file di-nama `llm_framework.py` — bukan `llm.py` — + # supaya tidak shadow package `scripts/llm/` saat + # `commands/__init__.py` auto-import semua submodule) + +tests/ +└── test_llm.py # 73 tests — semua network-free (SDK calls di-mock) +``` + +## Public API + +```python +from llm import ( + # ABCs untuk bikin tool domain-specific + LLMTool, LLMToolInput, LLMToolOutput, + # High-level entry point + invoke_llm, + # Provider dispatch + resolve_provider, get_provider, + # Errors + LLMError, LLMTimeoutError, + ProviderNotConfiguredError, ProviderNotInstalledError, +) +``` + +### Quick example + +```python +from dataclasses import dataclass +from llm import LLMTool, LLMToolInput, LLMToolOutput, invoke_llm + +@dataclass(frozen=True) +class TaintInput(LLMToolInput): + bug_type: str + tainted_value: str + + def __hash__(self): return hash((self.bug_type, self.tainted_value)) + def __eq__(self, other): + return isinstance(other, TaintInput) and \ + (self.bug_type, self.tainted_value) == (other.bug_type, other.tainted_value) + +@dataclass +class TaintOutput(LLMToolOutput): + is_false_positive: bool + explanation: str + + def __hash__(self): return hash(self.explanation) + def __eq__(self, other): + return isinstance(other, TaintOutput) and self.explanation == other.explanation + + +class TaintValidator(LLMTool): + def _get_prompt(self, inp: TaintInput) -> str: + return f"Is this {inp.bug_type} finding a false positive? Tainted value: {inp.tainted_value}" + + def _parse_response(self, raw: str, inp: TaintInput) -> TaintOutput: + # Parse the model's reply into a typed output. + is_fp = "false positive" in raw.lower() + return TaintOutput(is_false_positive=is_fp, explanation=raw) + + +# Usage — provider + retry + timeout are handled by the framework. +tool = TaintValidator() # reads CODELENS_LLM_MODEL + CODELENS_LLM_API_KEY from env +result = tool.invoke(TaintInput(bug_type="SQL injection", tainted_value="user_input")) +print(result.output.is_false_positive) +print(result.stats.attempts, result.stats.elapsed_seconds) +``` + +Atau langsung tanpa subclass, untuk one-shot call: + +```python +from llm import invoke_llm + +raw, stats = invoke_llm( + prompt="Explain this bug: ...", + model="glm-4.5", # dispatch ke zai_glm provider + api_key="...", # atau set ZAI_API_KEY env var +) +``` + +## Provider dispatch + +Dispatch by model name prefix (case-insensitive, first match wins): + +| Prefix | Provider | SDK | +|---------------------------------|-------------|--------------------| +| `gpt-`, `o1-`, `o3-`, `o4-`, `chatgpt-` | `openai` | `openai` | +| `claude-` | `anthropic` | `anthropic` | +| `bedrock-`, `amazon.` | `bedrock` | `boto3` | +| `gemini-` | `google` | `google-generativeai` | +| `deepseek-` | `deepseek` | `openai` (OpenAI-compatible endpoint) | +| `glm-`, `glm4-`, `zai-` | `zai_glm` | `openai` (OpenAI-compatible endpoint, base URL `https://open.bigmodel.cn/api/paas/v4/`) | + +Force a provider regardless of model name: set `CODELENS_LLM_PROVIDER=openai` +(useful for self-hosted OpenAI-compatible endpoints). + +## Config (env vars) + +| Variable | Purpose | +|---------------------------|-------------------------------------------------------------| +| `CODELENS_LLM_MODEL` | Default model name (e.g. `glm-4.5`). | +| `CODELENS_LLM_API_KEY` | Fallback API key for any provider. | +| `CODELENS_LLM_PROVIDER` | Force a provider (skip prefix dispatch). | +| `OPENAI_API_KEY` | OpenAI-specific key (preferred over the fallback). | +| `ANTHROPIC_API_KEY` | Anthropic-specific key. | +| `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` + `AWS_REGION` | Bedrock credentials. | +| `GOOGLE_API_KEY` / `GEMINI_API_KEY` | Google-specific key. | +| `DEEPSEEK_API_KEY` | DeepSeek-specific key. | +| `ZAI_API_KEY` / `GLM_API_KEY` | Z.ai GLM-specific key. | +| `ZAI_BASE_URL` | Override the Z.ai GLM base URL (for self-hosted endpoints). | + +API key resolution order: explicit kwarg > provider-specific env var > `CODELENS_LLM_API_KEY`. + +## Error model + +``` +LLMError (base) +├── LLMTimeoutError — retryable (transient network slowness) +├── ProviderNotConfiguredError — NOT retryable (missing API key / model) +└── ProviderNotInstalledError — NOT retryable (SDK not importable) +``` + +The `LLMError.retryable` flag drives the retry loop. Non-retryable errors +propagate on the first attempt — no point retrying a missing API key. + +## CLI command + +```bash +# List all 6 providers + their env var hints + SDK pip names +codelens llm providers + +# Show currently resolved config (model, provider, which env vars are set) +# — does NOT print API key values, only which env vars have non-empty values +codelens llm config + +# Send a 1-token smoke prompt to verify provider + API key + SDK end-to-end +codelens llm ping [--model MODEL] [--provider PROVIDER] [--timeout 15] +``` + +## Default behaviour + +| Setting | Default | Override | +|----------------------|--------------------|-----------------------------------------| +| Timeout per call | 60 seconds | `timeout_seconds=` kwarg | +| Max retries | 3 (including first)| `max_retries=` kwarg | +| Backoff | Exponential 1s → 2s → 4s | (not configurable in Phase 1) | + +## Phases 2–5 (deferred) + +| Phase | Scope | Status | +|-------|----------------------------------------------------------|----------------| +| 2 | Disk cache + cost tracking + `llm-cache` command | Not started | +| 3 | `ExplanationGenerator` tool + SARIF embedding | Not started | +| 4 | Reasoning offload for `codelens_explore` MCP tool | Not started | +| 5 | MCP prompts for rule authoring (`write_custom_codelens_rule`) | Not started | + +Phase 2 will add: + +- Cache at `~/.codelens/llm_cache//.json` +- Cache key = SHA-256 of `(tool_name, model_name, input_hash)` — + invalidates automatically when the model changes +- `codelens llm-cache stats` / `clear` subcommands +- `--no-cache` and `--max-cost-usd N` flags +- Auto-evict entries >30 days +- Thread-safe for concurrent agents + +The `LLMToolInput` ABC already requires `__hash__` / `__eq__` so the +Phase 2 cache can key off input objects directly — no API change needed +when the cache lands. + +## Testing + +``` +PYTHONUTF8=1 PYTHONPATH=scripts python3 -m pytest tests/test_llm.py -v +``` + +73 tests, all network-free. Provider SDK calls are mocked so the tests +run in any environment without API keys or SDKs installed. The tests +verify the framework's *logic* (dispatch, retry, config resolution), +not the SDK call shapes — those are validated by the SDK authors. + +## Design decisions + +1. **Why `commands/llm_framework.py` instead of `commands/llm.py`?** + The `commands/__init__.py` auto-imports every `.py` in `commands/` + via `importlib.import_module`. When the file was named `llm.py`, the + import shadowed the `scripts/llm/` package — `import llm` resolved + to `commands/llm.py` (a single module), not the `scripts/llm/` + package. Renaming the file to `llm_framework.py` fixed the + collision. The user-facing command name (`codelens llm ...`) is + unaffected — it's set by `register_command("llm", ...)`. + +2. **Why thread-based timeout instead of `signal.SIGALRM`?** + CodeLens runs on Windows (per CONTEXT.md), and `signal.SIGALRM` is + POSIX-only. A `ThreadPoolExecutor` with `future.result(timeout=)` + works on both platforms. The tradeoff is one idle thread per + timed-out call — acceptable for LLM use where calls are infrequent + and bounded by `max_retries`. + +3. **Why dispatch by prefix, not by an explicit provider field?** + Most users know the model name (`gpt-4o`, `claude-3-7`, `glm-4.5`) + but not which "provider" it belongs to. Prefix dispatch makes the + common case (just set `CODELENS_LLM_MODEL`) work out of the box. + The `CODELENS_LLM_PROVIDER` env var and `provider=` kwarg exist for + the edge case where the user knows better (self-hosted endpoints). + +4. **Why is the cache deferred to Phase 2?** + Phase 1 establishes the abstraction contract — `LLMToolInput` + requires `__hash__` / `__eq__` so cache keys "just work" when the + cache lands. Building the cache before the abstraction would have + meant retrofitting it into N callers later. Phase 1 = foundation, + Phase 2 = persistence. + +5. **Why is the default provider "Z.ai GLM"?** + The issue spec says "Use Z.ai GLM provider as default for CodeLens + (matches existing `z-ai-web-dev-sdk` integration pattern)." Users + who want OpenAI just set `CODELENS_LLM_MODEL=gpt-4o`. diff --git a/pyproject.toml b/pyproject.toml index cfb46efa..975ec607 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "codelens" version = "8.2.0" -description = "Live Codebase Reference Intelligence — 70 commands for AI-powered code analysis, security auditing, and quality scoring" +description = "Live Codebase Reference Intelligence — 71 commands for AI-powered code analysis, security auditing, and quality scoring" readme = "README.md" license = {text = "MIT"} requires-python = ">=3.8" diff --git a/scripts/commands/llm_framework.py b/scripts/commands/llm_framework.py new file mode 100644 index 00000000..679c6cec --- /dev/null +++ b/scripts/commands/llm_framework.py @@ -0,0 +1,262 @@ +"""LLM command — inspect LLM framework config and test provider connectivity. + +Issue #63 Phase 1. Provides three subcommands:: + + codelens llm providers # list known providers + their env vars + codelens llm config # show the currently resolved config (no API key) + codelens llm ping # send a 1-token smoke prompt to verify the chain + +Why a command and not just tests? +--------------------------------- +The LLM framework is opt-in and env-var driven. Users need a way to +answer "did I configure this right?" without reading the source. The +``ping`` subcommand is the canonical end-to-end smoke test — it fails +fast on missing API keys, missing SDKs, and timeouts, with actionable +error messages. + +Phase 2 will add ``codelens llm-cache stats`` / ``clear`` here. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from typing import Any, Dict, List, Optional + +from commands import register_command + + +def add_args(parser: argparse.ArgumentParser) -> None: + sub = parser.add_subparsers(dest="llm_subcommand") + + p_providers = sub.add_parser( + "providers", + help="List known LLM providers and the env vars each one reads.", + ) + p_providers.add_argument( + "--json", + action="store_true", + help="Emit JSON instead of a human-readable table.", + ) + + p_config = sub.add_parser( + "config", + help="Show the currently resolved LLM config (model, provider, key sources).", + ) + p_config.add_argument( + "--json", + action="store_true", + help="Emit JSON instead of a human-readable block.", + ) + + p_ping = sub.add_parser( + "ping", + help="Send a 1-token smoke prompt to verify provider + API key + SDK.", + ) + p_ping.add_argument( + "--model", + default=None, + help="Model name (defaults to CODELENS_LLM_MODEL).", + ) + p_ping.add_argument( + "--provider", + default=None, + help="Force a provider (bypass prefix dispatch).", + ) + p_ping.add_argument( + "--timeout", + type=float, + default=15.0, + help="Per-call timeout in seconds (default: 15).", + ) + p_ping.add_argument( + "--json", + action="store_true", + help="Emit JSON instead of a human-readable line.", + ) + + +def execute(args: argparse.Namespace, workspace: str) -> Dict[str, Any]: + sub = getattr(args, "llm_subcommand", None) + if sub is None: + # No subcommand → behave like `codelens llm config` for discoverability. + return _cmd_config(json_output=False) + if sub == "providers": + return _cmd_providers(json_output=getattr(args, "json", False)) + if sub == "config": + return _cmd_config(json_output=getattr(args, "json", False)) + if sub == "ping": + return _cmd_ping( + model=args.model, + provider=args.provider, + timeout=args.timeout, + json_output=args.json, + ) + return { + "status": "error", + "error": f"unknown llm subcommand: {sub!r}", + "available": ["providers", "config", "ping"], + } + + +# ─── Subcommands ─────────────────────────────────────────────────────────── + + +def _cmd_providers(*, json_output: bool) -> Dict[str, Any]: + """List the 6 providers + their env var hints + SDK pip names.""" + # Lazy import so the command module is importable even if the llm + # package itself failed to load (defensive — shouldn't happen). + try: + from llm.provider import PROVIDER_PREFIX_MAP, _PROVIDER_API_KEY_ENV, _PROVIDER_PIP_NAME + except ImportError as e: + return { + "status": "error", + "error": f"llm framework not importable: {e}", + } + + rows: List[Dict[str, Any]] = [] + for provider in sorted(PROVIDER_PREFIX_MAP): + prefixes = PROVIDER_PREFIX_MAP[provider] + env_vars = list(_PROVIDER_API_KEY_ENV.get(provider, ("CODELENS_LLM_API_KEY",))) + pip_name = _PROVIDER_PIP_NAME.get(provider, "?") + rows.append({ + "provider": provider, + "model_prefixes": list(prefixes), + "api_key_env_vars": env_vars, + "sdk_pip_name": pip_name, + }) + + if not json_output: + print("CodeLens LLM providers (issue #63 Phase 1):") + print() + for r in rows: + print(f" {r['provider']}") + print(f" prefixes: {', '.join(r['model_prefixes'])}") + print(f" api_key_env: {' or '.join(r['api_key_env_vars'])}") + print(f" sdk_install: pip install {r['sdk_pip_name']}") + print() + + return { + "status": "ok", + "providers": rows, + "count": len(rows), + } + + +def _cmd_config(*, json_output: bool) -> Dict[str, Any]: + """Show the currently resolved config (model, provider, key source).""" + try: + from llm.provider import ( + DEFAULT_MAX_RETRIES, + DEFAULT_TIMEOUT_SECONDS, + PROVIDER_PREFIX_MAP, + _PROVIDER_API_KEY_ENV, + resolve_provider, + ) + except ImportError as e: + return {"status": "error", "error": f"llm framework not importable: {e}"} + + model = os.environ.get("CODELENS_LLM_MODEL", "").strip() or None + forced_provider = os.environ.get("CODELENS_LLM_PROVIDER", "").strip().lower() or None + resolved_provider: Optional[str] = None + resolve_error: Optional[str] = None + if model: + try: + resolved_provider = forced_provider or resolve_provider(model) + except ValueError as e: + resolve_error = str(e) + elif forced_provider: + resolved_provider = forced_provider + + # Which env vars have values (do NOT print the values themselves). + key_sources: Dict[str, bool] = {} + if resolved_provider: + for var in _PROVIDER_API_KEY_ENV.get(resolved_provider, ("CODELENS_LLM_API_KEY",)): + key_sources[var] = bool(os.environ.get(var, "").strip()) + + config_block = { + "model": model or "(not set — set CODELENS_LLM_MODEL)", + "model_env_var": "CODELENS_LLM_MODEL", + "provider_forced": forced_provider or "(not set)", + "provider_forced_env_var": "CODELENS_LLM_PROVIDER", + "provider_resolved": resolved_provider or "(could not resolve)", + "provider_resolve_error": resolve_error, + "api_key_sources": key_sources, + "timeout_seconds_default": DEFAULT_TIMEOUT_SECONDS, + "max_retries_default": DEFAULT_MAX_RETRIES, + "known_providers": sorted(PROVIDER_PREFIX_MAP), + } + + if not json_output: + print("CodeLens LLM config (issue #63 Phase 1):") + print() + for k, v in config_block.items(): + print(f" {k}: {v}") + + return {"status": "ok", "config": config_block} + + +def _cmd_ping( + *, + model: Optional[str], + provider: Optional[str], + timeout: float, + json_output: bool, +) -> Dict[str, Any]: + """Send a 1-token smoke prompt to verify the chain end-to-end.""" + try: + from llm.provider import invoke_llm + from llm.base_tool import LLMError, ProviderNotConfiguredError, ProviderNotInstalledError + except ImportError as e: + return {"status": "error", "error": f"llm framework not importable: {e}"} + + smoke_prompt = "Reply with exactly: OK" + try: + raw, stats = invoke_llm( + prompt=smoke_prompt, + model=model, + provider=provider, + timeout_seconds=timeout, + max_retries=1, # ping should fail fast + ) + except ProviderNotConfiguredError as e: + msg = f"NOT CONFIGURED: {e} (provider={e.provider}, env_var={e.env_var})" + if not json_output: + print(msg, file=sys.stderr) + return {"status": "not_configured", "error": str(e), "provider": e.provider, "env_var": e.env_var} + except ProviderNotInstalledError as e: + msg = f"SDK MISSING: {e} (provider={e.provider}, install with: pip install {e.pip_name})" + if not json_output: + print(msg, file=sys.stderr) + return {"status": "sdk_missing", "error": str(e), "provider": e.provider, "pip_name": e.pip_name} + except LLMError as e: + if not json_output: + print(f"LLM ERROR: {e}", file=sys.stderr) + return {"status": "error", "error": str(e)} + except Exception as e: + if not json_output: + print(f"UNEXPECTED ERROR: {e}", file=sys.stderr) + return {"status": "error", "error": f"unexpected: {e}"} + + if not json_output: + print( + f"OK — {stats.provider}/{stats.model} " + f"({stats.attempts} attempt, {stats.elapsed_seconds:.2f}s)" + ) + return { + "status": "ok", + "provider": stats.provider, + "model": stats.model, + "attempts": stats.attempts, + "elapsed_seconds": round(stats.elapsed_seconds, 3), + "raw_preview": stats.raw_response_preview, + } + + +register_command( + "llm", + "LLM framework: list providers, show config, ping model (issue #63)", + add_args, + execute, +) diff --git a/scripts/graph_model.py b/scripts/graph_model.py index a3cc6a08..8b2e8cd5 100644 --- a/scripts/graph_model.py +++ b/scripts/graph_model.py @@ -13,7 +13,7 @@ - New tables `graph_nodes` and `graph_edges` are additive (prefixed `graph_` to avoid colliding with any existing table name). - The flat registry tables and JSON files are untouched. -- All 70 existing CLI commands continue to work unchanged. +- All 71 existing CLI commands continue to work unchanged. Schema: graph_nodes( diff --git a/scripts/llm/__init__.py b/scripts/llm/__init__.py new file mode 100644 index 00000000..c8fa72b0 --- /dev/null +++ b/scripts/llm/__init__.py @@ -0,0 +1,61 @@ +# @WHO: scripts/llm/__init__.py +# @WHAT: LLM integration framework package — multi-provider abstraction (issue #63 Phase 1) +# @PART: llm +# @ENTRY: - +# +# Phase 1 scope (issue #63): +# - ``base_tool.LLMTool`` ABC + ``LLMToolInput`` / ``LLMToolOutput`` ABCs +# - ``provider.invoke_llm`` dispatch by ``model_name`` prefix +# - 6 providers: OpenAI, Anthropic, Bedrock, Google, DeepSeek, Z.ai GLM +# - Lazy import per provider, 60s timeout, 3-retry exponential backoff +# - Config via ``CODELENS_LLM_PROVIDER`` / ``CODELENS_LLM_MODEL`` / +# ``CODELENS_LLM_API_KEY`` env vars +# +# Phases 2-5 (cache, explanation generator, reasoning offload, MCP prompts) +# are deferred to follow-up issues. + +"""LLM integration framework for CodeLens. + +Re-exported entry points:: + + from llm import LLMTool, LLMToolInput, LLMToolOutput, invoke_llm, get_provider + +The high-level :func:`invoke_llm` is the only function most callers need. +Subclass :class:`LLMTool` to build a domain-specific tool (e.g. an +``ExplanationGenerator`` in Phase 3); the framework handles provider +dispatch, retry, and timeout. +""" + +from .base_tool import ( # noqa: F401 + LLMError, + LLMTimeoutError, + LLMTool, + LLMToolInput, + LLMToolOutput, + ProviderNotConfiguredError, + ProviderNotInstalledError, +) +from .provider import ( # noqa: F401 + DEFAULT_TIMEOUT_SECONDS, + DEFAULT_MAX_RETRIES, + PROVIDER_PREFIX_MAP, + invoke_llm, + resolve_provider, + get_provider, +) + +__all__ = [ + "LLMTool", + "LLMToolInput", + "LLMToolOutput", + "LLMError", + "LLMTimeoutError", + "ProviderNotConfiguredError", + "ProviderNotInstalledError", + "invoke_llm", + "resolve_provider", + "get_provider", + "PROVIDER_PREFIX_MAP", + "DEFAULT_TIMEOUT_SECONDS", + "DEFAULT_MAX_RETRIES", +] diff --git a/scripts/llm/base_tool.py b/scripts/llm/base_tool.py new file mode 100644 index 00000000..7694cb1f --- /dev/null +++ b/scripts/llm/base_tool.py @@ -0,0 +1,304 @@ +# @WHO: scripts/llm/base_tool.py +# @WHAT: LLMTool ABC + LLMToolInput / LLMToolOutput ABCs — domain-specific LLM tool base +# @PART: llm +# @ENTRY: LLMTool.invoke() +# +# Issue #63 Phase 1 — LLMTool ABC + provider abstraction. +# +# Design: +# - ``LLMToolInput`` / ``LLMToolOutput`` are ABCs with ``__hash__`` / ``__eq__`` +# so cache keys (Phase 2) work directly off the input object. Subclasses +# must be immutable dataclasses (frozen=True). +# - ``LLMTool`` is the ABC every domain tool (e.g. Phase 3 ``ExplanationGenerator``) +# subclasses. It centralises the provider dispatch + retry + timeout +# logic so subclasses only need to implement ``_get_prompt`` (input → +# prompt string) and ``_parse_response`` (raw model output → typed output). +# - The actual provider call lives in ``provider.invoke_llm`` — keep +# ``LLMTool.invoke`` thin so retry behaviour is testable in isolation. +# +# Error model: +# - ``LLMError`` is the base — never raised directly. Use one of the +# subclasses so callers can do granular ``except`` clauses. +# - ``LLMTimeoutError`` — request exceeded the timeout. Retryable by default. +# - ``ProviderNotConfiguredError`` — no API key / endpoint configured. +# NOT retryable; caller should report a config error. +# - ``ProviderNotInstalledError`` — provider SDK not importable. NOT +# retryable; caller should report an install hint. + +"""LLMTool ABC + LLMToolInput / LLMToolOutput ABCs. + +The framework's contract:: + + class MyTool(LLMTool): + def _get_prompt(self, inp: MyInput) -> str: ... + def _parse_response(self, raw: str, inp: MyInput) -> MyOutput: ... + + tool = MyTool(model="glm-4.5", api_key="...") + result = tool.invoke(MyInput(...)) # → MyOutput + +``invoke`` is the only public method on subclasses — everything else is +protected. The framework handles provider dispatch, retry, and timeout. +""" + +from __future__ import annotations + +import abc +import time +from dataclasses import dataclass +from typing import Any, Dict, Generic, Optional, TypeVar + +# ─── Errors ──────────────────────────────────────────────────────────────── + + +class LLMError(Exception): + """Base error for all LLM framework failures. + + Never raised directly — use one of the subclasses below. All LLM + framework errors inherit from this so callers can catch the whole + family with ``except LLMError``. + """ + + def __init__(self, message: str, *, retryable: bool = False) -> None: + super().__init__(message) + self.retryable = retryable + + +class LLMTimeoutError(LLMError): + """Request exceeded the per-call timeout. + + Retryable by default (transient network slowness). + """ + + def __init__(self, message: str, *, timeout_seconds: float) -> None: + super().__init__(message, retryable=True) + self.timeout_seconds = timeout_seconds + + +class ProviderNotConfiguredError(LLMError): + """Provider has no API key / endpoint configured. + + NOT retryable — caller must surface a config hint. + """ + + def __init__(self, message: str, *, provider: str, env_var: str) -> None: + super().__init__(message, retryable=False) + self.provider = provider + self.env_var = env_var + + +class ProviderNotInstalledError(LLMError): + """Provider SDK is not importable. + + NOT retryable — caller must surface an install hint. + """ + + def __init__(self, message: str, *, provider: str, pip_name: str) -> None: + super().__init__(message, retryable=False) + self.provider = provider + self.pip_name = pip_name + + +# ─── Input / Output ABCs ─────────────────────────────────────────────────── +# +# Subclasses MUST be frozen dataclasses so __hash__ / __eq__ are stable +# for cache keys (Phase 2). The ABCs themselves don't enforce frozen=True +# because Python's dataclass machinery can't enforce it via ABC, but the +# contract is documented and tested. + +InputT = TypeVar("InputT", bound="LLMToolInput") +OutputT = TypeVar("OutputT", bound="LLMToolOutput") + + +class LLMToolInput(abc.ABC): + """Abstract base for typed LLM tool inputs. + + Subclasses MUST be ``@dataclass(frozen=True)`` so instances are + hashable and equality is value-based — this is required for the + disk cache (Phase 2) to key off the input object directly. + + The framework never inspects fields on the input — it just passes + it to ``LLMTool._get_prompt`` and ``LLMTool._parse_response``. + Subclasses are free to model their domain however they like. + """ + + @abc.abstractmethod + def __hash__(self) -> int: # pragma: no cover — abstract + ... + + @abc.abstractmethod + def __eq__(self, other: object) -> bool: # pragma: no cover — abstract + ... + + +class LLMToolOutput(abc.ABC): + """Abstract base for typed LLM tool outputs. + + Subclasses are typically ``@dataclass`` (mutable is fine — outputs + are not used as cache keys). The framework treats the output as + opaque: it is whatever ``_parse_response`` returns. + """ + + @abc.abstractmethod + def __hash__(self) -> int: # pragma: no cover — abstract + ... + + @abc.abstractmethod + def __eq__(self, other: object) -> bool: # pragma: no cover — abstract + ... + + +# ─── Invocation metadata ─────────────────────────────────────────────────── + + +@dataclass +class InvocationStats: + """Per-invocation telemetry, returned alongside the typed output. + + Phase 1 keeps this minimal — Phase 2 will add cache hit/miss, + Phase 3 will add cost (USD). All fields are populated by + ``LLMTool.invoke`` and are read-only from the caller's perspective. + """ + + provider: str + model: str + attempts: int + elapsed_seconds: float + timed_out: bool = False + raw_response_preview: str = "" + + def as_dict(self) -> Dict[str, Any]: + return { + "provider": self.provider, + "model": self.model, + "attempts": self.attempts, + "elapsed_seconds": round(self.elapsed_seconds, 3), + "timed_out": self.timed_out, + "raw_response_preview": self.raw_response_preview, + } + + +@dataclass +class InvocationResult(Generic[OutputT]): + """Wraps the typed output + telemetry for a single invoke call.""" + + output: OutputT + stats: InvocationStats + + +# ─── LLMTool ABC ─────────────────────────────────────────────────────────── + + +class LLMTool(abc.ABC, Generic[InputT, OutputT]): + """Abstract base for domain-specific LLM tools. + + Subclasses implement:: + + _get_prompt(inp) -> str # what to send to the model + _parse_response(raw, inp) -> OutputT # how to interpret the reply + + The base class provides ``invoke`` — the single public entry point. + ``invoke`` resolves the provider from ``model_name``, calls + ``provider.invoke_llm`` (which handles retry + timeout), then hands + the raw response to ``_parse_response``. + + Config resolution order (first non-empty wins): + 1. Explicit kwargs to ``__init__`` / ``invoke`` + 2. ``CODELENS_LLM_PROVIDER`` / ``CODELENS_LLM_MODEL`` / ``CODELENS_LLM_API_KEY`` + 3. Workspace config (``.codelens/codelens.config.json`` — Phase 2) + + Phase 1 deliberately does NOT implement workspace config — env vars + are sufficient for the abstraction. Workspace config lands with the + cache layer in Phase 2. + """ + + def __init__( + self, + *, + model: Optional[str] = None, + api_key: Optional[str] = None, + provider: Optional[str] = None, + timeout_seconds: Optional[float] = None, + max_retries: Optional[int] = None, + ) -> None: + self._model_override = model + self._api_key_override = api_key + self._provider_override = provider + self._timeout_override = timeout_seconds + self._retries_override = max_retries + + # ─── Subclass contract ──────────────────────────────────────── + + @abc.abstractmethod + def _get_prompt(self, inp: InputT) -> str: + """Render the input as a single prompt string for the model.""" + + @abc.abstractmethod + def _parse_response(self, raw: str, inp: InputT) -> OutputT: + """Parse the model's raw text response into a typed output.""" + + # ─── Public API ─────────────────────────────────────────────── + + # @FLOW: LLM_TOOL_INVOKE + # @CALLS: llm.provider.invoke_llm() -> str + # @MUTATES: none (pure — disk cache lands in Phase 2) + # @BEHAVIOR: Retryable on LLMTimeoutError / retryable LLMError. + # After max_retries attempts, the last error is re-raised. + # Non-retryable errors (ProviderNotConfiguredError / + # ProviderNotInstalledError) propagate immediately on + # the first attempt — no point retrying a missing API key. + def invoke(self, inp: InputT) -> InvocationResult[OutputT]: + """Run the tool on the given input. + + Args: + inp: Typed input — must be a frozen dataclass instance. + + Returns: + ``InvocationResult`` wrapping the typed output + telemetry. + + Raises: + LLMError (or subclass) on failure. + """ + # Lazy import — keeps ``base_tool`` importable without the + # provider module loaded (helps test isolation). + from .provider import invoke_llm + + prompt = self._get_prompt(inp) + raw, stats = invoke_llm( + prompt=prompt, + model=self._model_override, + api_key=self._api_key_override, + provider=self._provider_override, + timeout_seconds=self._timeout_override, + max_retries=self._retries_override, + ) + output = self._parse_response(raw, inp) + return InvocationResult(output=output, stats=stats) + + # ─── Introspection (used by `codelens llm config`) ─────────── + + def describe(self) -> Dict[str, Any]: + """Return a JSON-serialisable description of the tool's config. + + Used by ``codelens llm config`` to show what model / provider / + timeout a tool will use. Does NOT include the API key. + """ + return { + "tool": self.__class__.__name__, + "model": self._model_override or "(env default)", + "provider": self._provider_override or "(auto from model)", + "timeout_seconds": self._timeout_override or "(default)", + "max_retries": self._retries_override or "(default)", + } + + +__all__ = [ + "LLMError", + "LLMTimeoutError", + "ProviderNotConfiguredError", + "ProviderNotInstalledError", + "LLMToolInput", + "LLMToolOutput", + "LLMTool", + "InvocationStats", + "InvocationResult", +] diff --git a/scripts/llm/provider.py b/scripts/llm/provider.py new file mode 100644 index 00000000..d887ab71 --- /dev/null +++ b/scripts/llm/provider.py @@ -0,0 +1,593 @@ +# @WHO: scripts/llm/provider.py +# @WHAT: Multi-provider LLM dispatch — OpenAI / Anthropic / Bedrock / Google / DeepSeek / Z.ai GLM +# @PART: llm +# @ENTRY: invoke_llm(), resolve_provider(), get_provider() +# +# Issue #63 Phase 1 — provider abstraction. +# +# Dispatch model: +# provider = resolve_provider(model_name="glm-4.5") # → "zai_glm" +# raw_text, stats = invoke_llm(prompt="...", model="glm-4.5", api_key="...") +# +# Provider → model prefix mapping (first match wins, case-insensitive): +# gpt-*, o1-*, o3-* → openai +# claude-* → anthropic +# bedrock-* / amazon.* → bedrock (AWS) +# gemini-* → google +# deepseek-* → deepseek +# glm-*, glm4-*, zai-* → zai_glm +# +# Lazy import per provider: +# - The provider SDK is only imported when a call actually targets that +# provider. A user with only the OpenAI SDK installed can still import +# ``llm`` and dispatch to OpenAI; calls to Anthropic fail with +# ``ProviderNotInstalledError`` (not ``ImportError`` at module load). +# +# Config: +# - ``CODELENS_LLM_PROVIDER`` — force a provider (skip prefix dispatch) +# - ``CODELENS_LLM_MODEL`` — default model name +# - ``CODELENS_LLM_API_KEY`` — default API key (provider-agnostic fallback) +# - Per-provider API key env vars also honoured (see ``_PROVIDER_API_KEY_ENV``). +# +# Retry + timeout: +# - Default 60s per call, 3 retries with exponential backoff (1s, 2s, 4s). +# - Only retryable errors trigger retry (``LLMError.retryable=True``). +# - Non-retryable errors (missing config, missing SDK) propagate immediately. + +"""Multi-provider LLM dispatch. + +The public entry point is :func:`invoke_llm`. Most callers don't need +to call :func:`resolve_provider` or :func:`get_provider` directly — +``invoke_llm`` does it internally and returns the parsed text + telemetry. +""" + +from __future__ import annotations + +import os +import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError +from typing import Any, Callable, Dict, Optional, Tuple + +from utils import logger + +from .base_tool import ( + InvocationStats, + LLMError, + LLMTimeoutError, + ProviderNotConfiguredError, + ProviderNotInstalledError, +) + +# ─── Constants ───────────────────────────────────────────────────────────── + +DEFAULT_TIMEOUT_SECONDS = 60.0 +DEFAULT_MAX_RETRIES = 3 + +# Base backoff in seconds for the first retry. Subsequent retries double it +# (1s → 2s → 4s ...). Kept small because LLM timeouts are usually transient. +_BASE_BACKOFF_SECONDS = 1.0 + +# Provider name → tuple of model-name prefixes (lowercase, no separator). +# First match wins, so put the most specific prefixes first within each tuple. +# (No two providers share a prefix today, but if they ever do, order matters.) +PROVIDER_PREFIX_MAP: Dict[str, Tuple[str, ...]] = { + "openai": ("gpt-", "o1-", "o3-", "o4-", "chatgpt-"), + "anthropic": ("claude-",), + "bedrock": ("bedrock-", "amazon."), + "google": ("gemini-",), + "deepseek": ("deepseek-",), + "zai_glm": ("glm-", "glm4-", "zai-"), +} + +# Reverse index: prefix → provider. Built once at import time. +_PREFIX_TO_PROVIDER: Dict[str, str] = { + prefix: provider + for provider, prefixes in PROVIDER_PREFIX_MAP.items() + for prefix in prefixes +} + +# Per-provider API key env var names. The first non-empty one wins. +# ``CODELENS_LLM_API_KEY`` is the catch-all fallback for any provider. +_PROVIDER_API_KEY_ENV: Dict[str, Tuple[str, ...]] = { + "openai": ("OPENAI_API_KEY", "CODELENS_LLM_API_KEY"), + "anthropic": ("ANTHROPIC_API_KEY", "CODELENS_LLM_API_KEY"), + "bedrock": ("AWS_ACCESS_KEY_ID", "CODELENS_LLM_API_KEY"), # secret via AWS_SECRET_ACCESS_KEY + "google": ("GOOGLE_API_KEY", "GEMINI_API_KEY", "CODELENS_LLM_API_KEY"), + "deepseek": ("DEEPSEEK_API_KEY", "CODELENS_LLM_API_KEY"), + "zai_glm": ("ZAI_API_KEY", "GLM_API_KEY", "CODELENS_LLM_API_KEY"), +} + +# Per-provider pip-install name (for the install hint when SDK is missing). +_PROVIDER_PIP_NAME: Dict[str, str] = { + "openai": "openai", + "anthropic": "anthropic", + "bedrock": "boto3", + "google": "google-generativeai", + "deepseek": "openai", # DeepSeek API is OpenAI-compatible — same SDK + "zai_glm": "openai", # Z.ai GLM API is OpenAI-compatible — same SDK +} + + +# ─── Provider resolution ─────────────────────────────────────────────────── + + +def resolve_provider(model_name: str) -> str: + """Resolve which provider a model name belongs to. + + Dispatch is by prefix (case-insensitive). If no prefix matches, + raises ``ValueError`` — callers should validate the model name + before reaching this point. + + The ``CODELENS_LLM_PROVIDER`` env var, when set, overrides this + function entirely (returns that value verbatim). This lets users + force a provider even if the model name doesn't match a known prefix + (e.g. self-hosted OpenAI-compatible endpoints). + """ + forced = os.environ.get("CODELENS_LLM_PROVIDER", "").strip().lower() + if forced: + if forced not in PROVIDER_PREFIX_MAP: + raise ValueError( + f"CODELENS_LLM_PROVIDER={forced!r} is not a known provider " + f"(known: {sorted(PROVIDER_PREFIX_MAP)})" + ) + return forced + + if not model_name: + raise ValueError("model_name is required when CODELENS_LLM_PROVIDER is not set") + + lowered = model_name.lower() + for prefix, provider in _PREFIX_TO_PROVIDER.items(): + if lowered.startswith(prefix): + return provider + + raise ValueError( + f"Could not resolve provider for model {model_name!r}. " + f"Known prefixes: {sorted(_PREFIX_TO_PROVIDER)}. " + f"Set CODELENS_LLM_PROVIDER to force one." + ) + + +def get_provider(model_name: str) -> str: + """Alias for :func:`resolve_provider` — kept for callers that prefer the + shorter name. Identical behaviour.""" + return resolve_provider(model_name) + + +# ─── Config helpers ──────────────────────────────────────────────────────── + + +def _resolve_api_key(provider: str, explicit: Optional[str]) -> Optional[str]: + """Resolve the API key for a provider. + + Order: explicit kwarg > per-provider env var(s) > CODELENS_LLM_API_KEY. + Returns ``None`` if no key is found. + """ + if explicit: + return explicit + env_names = _PROVIDER_API_KEY_ENV.get(provider, ("CODELENS_LLM_API_KEY",)) + for name in env_names: + val = os.environ.get(name, "").strip() + if val: + return val + return None + + +def _resolve_model(explicit: Optional[str]) -> str: + """Resolve the model name. ``CODELENS_LLM_MODEL`` env var is the fallback.""" + if explicit: + return explicit + env_val = os.environ.get("CODELENS_LLM_MODEL", "").strip() + if env_val: + return env_val + raise ProviderNotConfiguredError( + "No model name configured. Set CODELENS_LLM_MODEL or pass model= explicitly.", + provider="(unknown)", + env_var="CODELENS_LLM_MODEL", + ) + + +# ─── Per-provider call wrappers ──────────────────────────────────────────── +# +# Each ``_call_`` function: +# - Imports the provider SDK lazily (so a missing SDK doesn't break import +# of ``provider.py`` — only the call fails). +# - Builds the request in the SDK's expected shape. +# - Returns the raw text response. +# - Raises ``LLMError`` (or subclass) on any failure. +# +# All wrappers take the same kwargs so the dispatch site is uniform: +# (prompt: str, model: str, api_key: str) -> str + + +def _call_openai(prompt: str, model: str, api_key: str) -> str: + try: + from openai import OpenAI # type: ignore + except ImportError as e: + raise ProviderNotInstalledError( + f"OpenAI SDK not installed: {e}", + provider="openai", + pip_name=_PROVIDER_PIP_NAME["openai"], + ) from e + + try: + client = OpenAI(api_key=api_key) + # ``o1-`` / ``o3-`` reasoning models don't accept ``temperature``; + # the SDK auto-handles this since v1.40+, but we pass the minimal + # shape to be safe. + resp = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + ) + return resp.choices[0].message.content or "" + except Exception as e: + # Re-raise non-LLMError exceptions as LLMError so the retry loop + # can decide whether to retry. + if isinstance(e, LLMError): + raise + raise LLMError(f"[llm.openai] request failed: {e}", retryable=True) from e + + +def _call_anthropic(prompt: str, model: str, api_key: str) -> str: + try: + import anthropic # type: ignore + except ImportError as e: + raise ProviderNotInstalledError( + f"Anthropic SDK not installed: {e}", + provider="anthropic", + pip_name=_PROVIDER_PIP_NAME["anthropic"], + ) from e + + try: + client = anthropic.Anthropic(api_key=api_key) + # Anthropic separates ``max_tokens`` (required) from the prompt. + # Use 4096 as a sane default — callers can override later via + # tool subclasses if they need more. + msg = client.messages.create( + model=model, + max_tokens=4096, + messages=[{"role": "user", "content": prompt}], + ) + # ``msg.content`` is a list of content blocks; concatenate text blocks. + parts = [b.text for b in msg.content if getattr(b, "type", None) == "text"] + return "".join(parts) + except Exception as e: + if isinstance(e, LLMError): + raise + raise LLMError(f"[llm.anthropic] request failed: {e}", retryable=True) from e + + +def _call_bedrock(prompt: str, model: str, api_key: str) -> str: + """Call AWS Bedrock via boto3. + + Bedrock uses AWS credentials (``AWS_ACCESS_KEY_ID`` + + ``AWS_SECRET_ACCESS_KEY`` + optional ``AWS_REGION``) rather than a + single API key. ``api_key`` here is the access key ID; the secret + must be in ``AWS_SECRET_ACCESS_KEY``. + """ + try: + import boto3 # type: ignore + except ImportError as e: + raise ProviderNotInstalledError( + f"boto3 not installed: {e}", + provider="bedrock", + pip_name=_PROVIDER_PIP_NAME["bedrock"], + ) from e + + try: + region = os.environ.get("AWS_REGION", "us-east-1") + secret = os.environ.get("AWS_SECRET_ACCESS_KEY", "").strip() + if not secret: + raise ProviderNotConfiguredError( + "AWS_SECRET_ACCESS_KEY not set (required for Bedrock)", + provider="bedrock", + env_var="AWS_SECRET_ACCESS_KEY", + ) + client = boto3.client( + "bedrock-runtime", + region_name=region, + aws_access_key_id=api_key, + aws_secret_access_key=secret, + ) + # Bedrock's InvokeModel API takes a JSON body whose shape depends + # on the underlying model (Anthropic / Meta / etc.). Strip the + # ``bedrock-`` prefix to get the real model ID, then send a minimal + # Anthropic-style payload if the model ID looks like Claude. + model_id = model[len("bedrock-"):] if model.startswith("bedrock-") else model + if model_id.startswith("anthropic."): + body = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 4096, + "messages": [{"role": "user", "content": prompt}], + } + import json as _json + resp = client.invoke_model( + modelId=model_id, + body=_json.dumps(body), + contentType="application/json", + accept="application/json", + ) + payload = _json.loads(resp["body"].read()) + parts = [b.get("text", "") for b in payload.get("content", []) if b.get("type") == "text"] + return "".join(parts) + # Non-Claude Bedrock models: surface a clear error rather than + # guess the body shape. Phase 2 can add per-model body builders. + raise LLMError( + f"[llm.bedrock] model {model_id!r} body shape not implemented in Phase 1 " + f"(only anthropic.* models supported).", + retryable=False, + ) + except LLMError: + raise + except Exception as e: + raise LLMError(f"[llm.bedrock] request failed: {e}", retryable=True) from e + + +def _call_google(prompt: str, model: str, api_key: str) -> str: + try: + import google.generativeai as genai # type: ignore + except ImportError as e: + raise ProviderNotInstalledError( + f"google-generativeai SDK not installed: {e}", + provider="google", + pip_name=_PROVIDER_PIP_NAME["google"], + ) from e + + try: + genai.configure(api_key=api_key) + gm = genai.GenerativeModel(model) + resp = gm.generate_content(prompt) + # ``resp.text`` raises if the response was blocked — fall back to + # an empty string and let the caller's ``_parse_response`` decide. + return getattr(resp, "text", "") or "" + except Exception as e: + if isinstance(e, LLMError): + raise + raise LLMError(f"[llm.google] request failed: {e}", retryable=True) from e + + +def _call_deepseek(prompt: str, model: str, api_key: str) -> str: + """DeepSeek's API is OpenAI-compatible — reuse the OpenAI SDK pointed at + DeepSeek's base URL.""" + try: + from openai import OpenAI # type: ignore + except ImportError as e: + raise ProviderNotInstalledError( + f"DeepSeek requires the OpenAI SDK (not installed): {e}", + provider="deepseek", + pip_name=_PROVIDER_PIP_NAME["deepseek"], + ) from e + + try: + client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com") + resp = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + ) + return resp.choices[0].message.content or "" + except Exception as e: + if isinstance(e, LLMError): + raise + raise LLMError(f"[llm.deepseek] request failed: {e}", retryable=True) from e + + +def _call_zai_glm(prompt: str, model: str, api_key: str) -> str: + """Z.ai GLM API is OpenAI-compatible — reuse the OpenAI SDK pointed at + Z.ai's base URL. + + Default base URL: ``https://open.bigmodel.cn/api/paas/v4/``. Override + via ``ZAI_BASE_URL`` env var for self-hosted endpoints. + """ + try: + from openai import OpenAI # type: ignore + except ImportError as e: + raise ProviderNotInstalledError( + f"Z.ai GLM requires the OpenAI SDK (not installed): {e}", + provider="zai_glm", + pip_name=_PROVIDER_PIP_NAME["zai_glm"], + ) from e + + try: + base_url = os.environ.get( + "ZAI_BASE_URL", + "https://open.bigmodel.cn/api/paas/v4/", + ) + client = OpenAI(api_key=api_key, base_url=base_url) + resp = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + ) + return resp.choices[0].message.content or "" + except Exception as e: + if isinstance(e, LLMError): + raise + raise LLMError(f"[llm.zai_glm] request failed: {e}", retryable=True) from e + + +_PROVIDER_CALL_TABLE: Dict[str, Callable[[str, str, str], str]] = { + "openai": _call_openai, + "anthropic": _call_anthropic, + "bedrock": _call_bedrock, + "google": _call_google, + "deepseek": _call_deepseek, + "zai_glm": _call_zai_glm, +} + + +# ─── Timeout helper ──────────────────────────────────────────────────────── + + +def _run_with_timeout(fn: Callable[[], str], timeout_seconds: float) -> str: + """Run ``fn`` in a worker thread with a hard timeout. + + Uses ``ThreadPoolExecutor`` so the timeout works on Windows as well + as POSIX (``signal.SIGALRM`` is POSIX-only). On timeout, raises + ``LLMTimeoutError`` — the worker thread continues to completion in + the background, but the caller is unblocked immediately. + + Why a thread and not ``signal.SIGALRM``? CodeLens must run on + Windows (see CONTEXT.md / pre-flight SKILL.md), and SIGALRM is + POSIX-only. The thread-based approach is portable at the cost of + one idle thread per timed-out call — acceptable for LLM use where + calls are infrequent and bounded by max_retries. + """ + with ThreadPoolExecutor(max_workers=1) as ex: + future = ex.submit(fn) + try: + return future.result(timeout=timeout_seconds) + except FuturesTimeoutError as e: + raise LLMTimeoutError( + f"LLM call exceeded {timeout_seconds}s timeout", + timeout_seconds=timeout_seconds, + ) from e + + +# ─── Public entry point ──────────────────────────────────────────────────── + + +# @FLOW: LLM_INVOKE +# @CALLS: _PROVIDER_CALL_TABLE[provider](prompt, model, api_key) -> str +# _run_with_timeout(call_fn, timeout) -> str +# @MUTATES: none +# @BEHAVIOR: Retries on LLMTimeoutError and any retryable LLMError. +# Non-retryable errors (ProviderNotConfiguredError, +# ProviderNotInstalledError) propagate on the first attempt. +# Returns ``(raw_text, InvocationStats)`` even on the last +# attempt's success — stats.attempts reflects how many calls +# were made (1 = first try succeeded, 3 = first two failed). +def invoke_llm( + *, + prompt: str, + model: Optional[str] = None, + api_key: Optional[str] = None, + provider: Optional[str] = None, + timeout_seconds: Optional[float] = None, + max_retries: Optional[int] = None, +) -> Tuple[str, InvocationStats]: + """Send a prompt to an LLM and return the raw response text + stats. + + Args: + prompt: The user prompt to send. Required. + model: Model name (e.g. ``"gpt-4o"``, ``"claude-3-7-sonnet"``, + ``"glm-4.5"``). Falls back to ``CODELENS_LLM_MODEL`` if omitted. + api_key: API key for the resolved provider. Falls back to + provider-specific env vars (e.g. ``OPENAI_API_KEY``) and + finally ``CODELENS_LLM_API_KEY``. + provider: Force a provider, bypassing prefix dispatch. Must be a + key in :data:`PROVIDER_PREFIX_MAP`. + timeout_seconds: Per-call timeout. Defaults to 60s. + max_retries: Max attempts (including the first). Defaults to 3. + + Returns: + ``(raw_text, InvocationStats)``. ``raw_text`` is the model's text + response; ``stats`` records provider/model/attempts/elapsed. + + Raises: + ProviderNotConfiguredError: No model / API key configured. + ProviderNotInstalledError: Provider SDK not importable. + LLMTimeoutError: All retries exhausted due to timeouts. + LLMError: All retries exhausted due to other retryable errors. + """ + if not prompt: + raise ValueError("prompt is required") + + resolved_model = _resolve_model(model) + resolved_provider = provider or resolve_provider(resolved_model) + resolved_timeout = timeout_seconds if timeout_seconds is not None else DEFAULT_TIMEOUT_SECONDS + resolved_retries = max_retries if max_retries is not None else DEFAULT_MAX_RETRIES + + resolved_key = _resolve_api_key(resolved_provider, api_key) + if not resolved_key: + env_hint = " or ".join(_PROVIDER_API_KEY_ENV.get(resolved_provider, ("CODELENS_LLM_API_KEY",))) + raise ProviderNotConfiguredError( + f"No API key for provider {resolved_provider!r}. " + f"Set {env_hint} or pass api_key= explicitly.", + provider=resolved_provider, + env_var=env_hint, + ) + + call_fn = _PROVIDER_CALL_TABLE.get(resolved_provider) + if call_fn is None: # pragma: no cover — defensive, resolve_provider validates + raise LLMError( + f"Provider {resolved_provider!r} has no call wrapper (internal bug).", + retryable=False, + ) + + start = time.monotonic() + last_err: Optional[Exception] = None + timed_out = False + attempts = 0 + + for attempt in range(1, resolved_retries + 1): + attempts = attempt + try: + raw = _run_with_timeout( + lambda: call_fn(prompt, resolved_model, resolved_key), + resolved_timeout, + ) + elapsed = time.monotonic() - start + stats = InvocationStats( + provider=resolved_provider, + model=resolved_model, + attempts=attempts, + elapsed_seconds=elapsed, + timed_out=False, + raw_response_preview=raw[:200], + ) + return raw, stats + except LLMTimeoutError as e: + last_err = e + timed_out = True + logger.warning( + f"[llm.invoke] attempt {attempt}/{resolved_retries} timed out " + f"({resolved_timeout}s) for {resolved_provider}/{resolved_model}" + ) + except LLMError as e: + last_err = e + if not e.retryable: + # Non-retryable: propagate immediately. + raise + logger.warning( + f"[llm.invoke] attempt {attempt}/{resolved_retries} failed " + f"for {resolved_provider}/{resolved_model}: {e}" + ) + except Exception as e: + # Unexpected exception — wrap and propagate as non-retryable. + last_err = e + raise LLMError( + f"[llm.invoke] unexpected error for {resolved_provider}/{resolved_model}: {e}", + retryable=False, + ) from e + + # Backoff before the next retry (skip on the last attempt). + if attempt < resolved_retries: + backoff = _BASE_BACKOFF_SECONDS * (2 ** (attempt - 1)) + time.sleep(backoff) + + # All retries exhausted. + elapsed = time.monotonic() - start + if timed_out and isinstance(last_err, LLMTimeoutError): + raise LLMTimeoutError( + f"LLM call to {resolved_provider}/{resolved_model} timed out " + f"{resolved_retries}× (last timeout: {resolved_timeout}s)", + timeout_seconds=resolved_timeout, + ) + if last_err is None: # pragma: no cover — defensive + raise LLMError( + f"LLM call to {resolved_provider}/{resolved_model} failed with no error captured.", + retryable=False, + ) + raise LLMError( + f"LLM call to {resolved_provider}/{resolved_model} failed after " + f"{resolved_retries} attempts: {last_err}", + retryable=False, + ) from last_err + + +__all__ = [ + "DEFAULT_TIMEOUT_SECONDS", + "DEFAULT_MAX_RETRIES", + "PROVIDER_PREFIX_MAP", + "invoke_llm", + "resolve_provider", + "get_provider", +] diff --git a/skill.json b/skill.json index 76595063..d236d4ae 100755 --- a/skill.json +++ b/skill.json @@ -1,7 +1,7 @@ { "name": "codelens", "version": "8.2.0", - "description": "Live Codebase Reference Intelligence. 70 commands for AI-powered code analysis, security auditing, quality scoring, and pre-write safety checks. Supports 28+ languages with regex+AST hybrid parsing. Must activate before writing/editing/deleting any class, id, or function.", + "description": "Live Codebase Reference Intelligence. 71 commands for AI-powered code analysis, security auditing, quality scoring, and pre-write safety checks. Supports 28+ languages with regex+AST hybrid parsing. Must activate before writing/editing/deleting any class, id, or function.", "author": "codelens", "command_categories": { "setup": [ diff --git a/tests/test_llm.py b/tests/test_llm.py new file mode 100644 index 00000000..1d419123 --- /dev/null +++ b/tests/test_llm.py @@ -0,0 +1,824 @@ +"""Tests for the LLM integration framework (issue #63 Phase 1). + +Scope: + +* ``llm.provider.resolve_provider`` — model-name → provider dispatch. +* ``llm.provider.invoke_llm`` — retry semantics, timeout handling, + config resolution (env vars + explicit kwargs), error propagation. +* ``llm.base_tool.LLMTool`` — ABC contract, ``invoke`` wires prompt + rendering → provider call → response parsing. +* ``commands.llm_framework`` — CLI registration + subcommand dispatch. + +All tests are **network-free**: provider SDK calls are mocked so the +tests run in any environment without API keys or SDKs installed. The +goal is to verify the framework's *logic* (dispatch, retry, config +resolution), not the SDK call shapes — those are validated by the SDK +authors. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import time +from dataclasses import dataclass +from typing import Any, Optional +from unittest import mock + +import pytest + +# ─── Path setup (mirror other tests) ─────────────────────────────────────── + +_THIS_DIR = os.path.dirname(os.path.abspath(__file__)) +_SCRIPTS_DIR = os.path.join(os.path.dirname(_THIS_DIR), "scripts") +if _SCRIPTS_DIR not in sys.path: + sys.path.insert(0, _SCRIPTS_DIR) + +from commands import COMMAND_REGISTRY # noqa: E402 +from llm import ( # noqa: E402 + DEFAULT_MAX_RETRIES, + DEFAULT_TIMEOUT_SECONDS, + PROVIDER_PREFIX_MAP, + LLMError, + LLMTimeoutError, + LLMTool, + LLMToolInput, + LLMToolOutput, + ProviderNotConfiguredError, + ProviderNotInstalledError, + invoke_llm, + resolve_provider, +) +from llm import provider as provider_mod # noqa: E402 +from llm import base_tool as base_tool_mod # noqa: E402 + + +# ─── Constants ────────────────────────────────────────────────────────────── + + +ALL_KNOWN_PROVIDERS = {"openai", "anthropic", "bedrock", "google", "deepseek", "zai_glm"} + + +# ─── Test fixtures ────────────────────────────────────────────────────────── + + +@pytest.fixture +def clean_env(monkeypatch): + """Strip all LLM env vars so each test starts from a known state. + + Tests that need specific env vars set them via ``monkeypatch.setenv`` + after this fixture runs — later ``setenv`` calls win. + """ + for var in ( + "CODELENS_LLM_PROVIDER", + "CODELENS_LLM_MODEL", + "CODELENS_LLM_API_KEY", + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "GOOGLE_API_KEY", + "GEMINI_API_KEY", + "DEEPSEEK_API_KEY", + "ZAI_API_KEY", + "GLM_API_KEY", + "ZAI_BASE_URL", + ): + monkeypatch.delenv(var, raising=False) + yield + + +# ─── Provider prefix dispatch ─────────────────────────────────────────────── + + +class TestResolveProvider: + """``resolve_provider`` maps model names to providers by prefix.""" + + @pytest.mark.parametrize( + "model,expected", + [ + ("gpt-4o", "openai"), + ("gpt-3.5-turbo", "openai"), + ("o1-preview", "openai"), + ("o3-mini", "openai"), + ("claude-3-7-sonnet", "anthropic"), + ("claude-opus-4", "anthropic"), + ("bedrock-anthropic.claude-3", "bedrock"), + ("amazon.nova-pro", "bedrock"), + ("gemini-1.5-pro", "google"), + ("gemini-2.0-flash", "google"), + ("deepseek-chat", "deepseek"), + ("deepseek-reasoner", "deepseek"), + ("glm-4.5", "zai_glm"), + ("glm4-plus", "zai_glm"), + ("zai-glm-4", "zai_glm"), + ], + ) + def test_known_prefixes(self, clean_env, model, expected): + assert resolve_provider(model) == expected + + def test_case_insensitive(self, clean_env): + assert resolve_provider("GPT-4o") == "openai" + assert resolve_provider("Claude-3-7") == "anthropic" + assert resolve_provider("GLM-4.5") == "zai_glm" + + def test_unknown_prefix_raises(self, clean_env): + with pytest.raises(ValueError, match="Could not resolve provider"): + resolve_provider("some-unknown-model-12345") + + def test_empty_model_raises(self, clean_env): + with pytest.raises(ValueError, match="model_name is required"): + resolve_provider("") + + def test_forced_provider_env_var(self, clean_env, monkeypatch): + monkeypatch.setenv("CODELENS_LLM_PROVIDER", "openai") + # Even with an unknown model name, the forced provider wins. + assert resolve_provider("totally-unknown-model") == "openai" + + def test_forced_provider_unknown_raises(self, clean_env, monkeypatch): + monkeypatch.setenv("CODELENS_LLM_PROVIDER", "not_a_real_provider") + with pytest.raises(ValueError, match="not a known provider"): + resolve_provider("gpt-4o") + + def test_all_six_providers_have_prefixes(self, clean_env): + """All 6 providers from issue #63 must be in the dispatch table.""" + assert set(PROVIDER_PREFIX_MAP.keys()) == ALL_KNOWN_PROVIDERS + + def test_no_two_providers_share_a_prefix(self, clean_env): + """Prefixes must be unique across providers (no ambiguity).""" + all_prefixes = [] + for prefixes in PROVIDER_PREFIX_MAP.values(): + all_prefixes.extend(prefixes) + assert len(all_prefixes) == len(set(all_prefixes)), ( + f"Duplicate prefixes: {all_prefixes}" + ) + + +# ─── Config resolution ───────────────────────────────────────────────────── + + +class TestConfigResolution: + """API key + model resolution order.""" + + def test_explicit_api_key_wins(self, clean_env, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "env-key") + assert provider_mod._resolve_api_key("openai", "explicit-key") == "explicit-key" + + def test_provider_specific_env_var(self, clean_env, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "openai-env-key") + assert provider_mod._resolve_api_key("openai", None) == "openai-env-key" + + def test_fallback_to_codelens_api_key(self, clean_env, monkeypatch): + monkeypatch.setenv("CODELENS_LLM_API_KEY", "fallback-key") + assert provider_mod._resolve_api_key("openai", None) == "fallback-key" + + def test_provider_env_takes_precedence_over_fallback(self, clean_env, monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "provider-specific") + monkeypatch.setenv("CODELENS_LLM_API_KEY", "fallback") + assert provider_mod._resolve_api_key("openai", None) == "provider-specific" + + def test_no_key_returns_none(self, clean_env): + assert provider_mod._resolve_api_key("openai", None) is None + + def test_model_explicit_wins(self, clean_env, monkeypatch): + monkeypatch.setenv("CODELENS_LLM_MODEL", "env-model") + assert provider_mod._resolve_model("explicit-model") == "explicit-model" + + def test_model_env_fallback(self, clean_env, monkeypatch): + monkeypatch.setenv("CODELENS_LLM_MODEL", "env-model") + assert provider_mod._resolve_model(None) == "env-model" + + def test_model_missing_raises_not_configured(self, clean_env): + with pytest.raises(ProviderNotConfiguredError) as exc_info: + provider_mod._resolve_model(None) + assert exc_info.value.env_var == "CODELENS_LLM_MODEL" + + +# ─── invoke_llm: config errors ───────────────────────────────────────────── + + +class TestInvokeLlmConfigErrors: + """``invoke_llm`` must fail fast on config errors.""" + + def test_missing_model_raises_not_configured(self, clean_env): + with pytest.raises(ProviderNotConfiguredError): + invoke_llm(prompt="hello", api_key="k") + + def test_missing_api_key_raises_not_configured(self, clean_env, monkeypatch): + monkeypatch.setenv("CODELENS_LLM_MODEL", "glm-4.5") + with pytest.raises(ProviderNotConfiguredError) as exc_info: + invoke_llm(prompt="hello") + # Error must mention the provider + at least one env var name. + assert exc_info.value.provider == "zai_glm" + assert "ZAI_API_KEY" in exc_info.value.env_var + + def test_empty_prompt_raises_value_error(self, clean_env, monkeypatch): + monkeypatch.setenv("CODELENS_LLM_MODEL", "glm-4.5") + monkeypatch.setenv("ZAI_API_KEY", "k") + with pytest.raises(ValueError, match="prompt is required"): + invoke_llm(prompt="", api_key="k", model="glm-4.5") + + +# ─── invoke_llm: provider dispatch ───────────────────────────────────────── + + +class TestInvokeLlmDispatch: + """``invoke_llm`` calls the right provider wrapper.""" + + def test_dispatches_to_zai_glm(self, clean_env, monkeypatch): + captured = {} + + def fake_call(prompt, model, api_key): + captured["prompt"] = prompt + captured["model"] = model + captured["api_key"] = api_key + return "GLM reply" + + monkeypatch.setattr(provider_mod, "_PROVIDER_CALL_TABLE", {"zai_glm": fake_call}) + monkeypatch.setenv("CODELENS_LLM_MODEL", "glm-4.5") + raw, stats = invoke_llm(prompt="hi", api_key="k") + assert raw == "GLM reply" + assert captured["model"] == "glm-4.5" + assert captured["api_key"] == "k" + assert stats.provider == "zai_glm" + assert stats.attempts == 1 + assert stats.timed_out is False + + def test_dispatches_to_openai_via_prefix(self, clean_env, monkeypatch): + monkeypatch.setattr( + provider_mod, + "_PROVIDER_CALL_TABLE", + {"openai": lambda p, m, k: "openai reply"}, + ) + raw, stats = invoke_llm(prompt="hi", model="gpt-4o", api_key="k") + assert raw == "openai reply" + assert stats.provider == "openai" + + def test_dispatches_to_anthropic_via_prefix(self, clean_env, monkeypatch): + monkeypatch.setattr( + provider_mod, + "_PROVIDER_CALL_TABLE", + {"anthropic": lambda p, m, k: "anthropic reply"}, + ) + raw, stats = invoke_llm(prompt="hi", model="claude-3-7", api_key="k") + assert raw == "anthropic reply" + assert stats.provider == "anthropic" + + def test_explicit_provider_overrides_prefix(self, clean_env, monkeypatch): + """``provider=`` kwarg bypasses prefix dispatch entirely.""" + monkeypatch.setattr( + provider_mod, + "_PROVIDER_CALL_TABLE", + {"openai": lambda p, m, k: "openai reply"}, + ) + # Use a model name that would normally resolve to anthropic, + # but force provider=openai. + raw, stats = invoke_llm( + prompt="hi", + model="claude-3-7", + api_key="k", + provider="openai", + ) + assert stats.provider == "openai" + + +# ─── invoke_llm: retry semantics ─────────────────────────────────────────── + + +class TestInvokeLlmRetry: + """Retry behaviour: retryable errors trigger retry, non-retryable don't.""" + + def test_retries_on_timeout(self, clean_env, monkeypatch): + call_count = {"n": 0} + + def fake_call(prompt, model, api_key): + call_count["n"] += 1 + if call_count["n"] < 3: + raise LLMTimeoutError("timeout", timeout_seconds=1.0) + return "success on attempt 3" + + monkeypatch.setattr(provider_mod, "_PROVIDER_CALL_TABLE", {"zai_glm": fake_call}) + monkeypatch.setattr(provider_mod.time, "sleep", lambda s: None) # no real backoff + raw, stats = invoke_llm( + prompt="hi", + model="glm-4.5", + api_key="k", + max_retries=3, + timeout_seconds=1.0, + ) + assert raw == "success on attempt 3" + assert stats.attempts == 3 + assert call_count["n"] == 3 + + def test_gives_up_after_max_retries(self, clean_env, monkeypatch): + def fake_call(prompt, model, api_key): + raise LLMTimeoutError("always timeout", timeout_seconds=1.0) + + monkeypatch.setattr(provider_mod, "_PROVIDER_CALL_TABLE", {"zai_glm": fake_call}) + monkeypatch.setattr(provider_mod.time, "sleep", lambda s: None) + with pytest.raises(LLMTimeoutError, match="timed out"): + invoke_llm( + prompt="hi", + model="glm-4.5", + api_key="k", + max_retries=2, + timeout_seconds=1.0, + ) + + def test_non_retryable_error_propagates_immediately(self, clean_env, monkeypatch): + call_count = {"n": 0} + + def fake_call(prompt, model, api_key): + call_count["n"] += 1 + raise ProviderNotInstalledError( + "missing SDK", provider="zai_glm", pip_name="openai" + ) + + monkeypatch.setattr(provider_mod, "_PROVIDER_CALL_TABLE", {"zai_glm": fake_call}) + with pytest.raises(ProviderNotInstalledError): + invoke_llm( + prompt="hi", + model="glm-4.5", + api_key="k", + max_retries=3, + ) + # Must NOT have retried — non-retryable. + assert call_count["n"] == 1 + + def test_retries_on_retryable_llm_error(self, clean_env, monkeypatch): + call_count = {"n": 0} + + def fake_call(prompt, model, api_key): + call_count["n"] += 1 + if call_count["n"] < 2: + raise LLMError("transient", retryable=True) + return "ok" + + monkeypatch.setattr(provider_mod, "_PROVIDER_CALL_TABLE", {"zai_glm": fake_call}) + monkeypatch.setattr(provider_mod.time, "sleep", lambda s: None) + raw, stats = invoke_llm( + prompt="hi", + model="glm-4.5", + api_key="k", + max_retries=3, + ) + assert raw == "ok" + assert stats.attempts == 2 + + def test_max_retries_one_means_no_retry(self, clean_env, monkeypatch): + call_count = {"n": 0} + + def fake_call(prompt, model, api_key): + call_count["n"] += 1 + raise LLMTimeoutError("timeout", timeout_seconds=1.0) + + monkeypatch.setattr(provider_mod, "_PROVIDER_CALL_TABLE", {"zai_glm": fake_call}) + monkeypatch.setattr(provider_mod.time, "sleep", lambda s: None) + with pytest.raises(LLMTimeoutError): + invoke_llm( + prompt="hi", + model="glm-4.5", + api_key="k", + max_retries=1, + timeout_seconds=1.0, + ) + assert call_count["n"] == 1 + + def test_backoff_doubles_between_retries(self, clean_env, monkeypatch): + """Verify exponential backoff: 1s, 2s, 4s, ...""" + sleeps = [] + + def fake_call(prompt, model, api_key): + raise LLMTimeoutError("timeout", timeout_seconds=1.0) + + monkeypatch.setattr(provider_mod, "_PROVIDER_CALL_TABLE", {"zai_glm": fake_call}) + monkeypatch.setattr(provider_mod.time, "sleep", lambda s: sleeps.append(s)) + with pytest.raises(LLMTimeoutError): + invoke_llm( + prompt="hi", + model="glm-4.5", + api_key="k", + max_retries=4, + timeout_seconds=1.0, + ) + # 3 sleeps between 4 attempts (no sleep after the last attempt). + assert sleeps == [1.0, 2.0, 4.0] + + +# ─── invoke_llm: timeout behaviour ───────────────────────────────────────── + + +class TestInvokeLlmTimeout: + """The ``_run_with_timeout`` helper must enforce the timeout.""" + + def test_slow_call_raises_timeout(self, clean_env, monkeypatch): + def slow_call(): + time.sleep(0.5) + return "should not get here" + + with pytest.raises(LLMTimeoutError) as exc_info: + provider_mod._run_with_timeout(slow_call, timeout_seconds=0.1) + assert exc_info.value.timeout_seconds == 0.1 + + def test_fast_call_returns_normally(self, clean_env): + def fast_call(): + return "ok" + + assert provider_mod._run_with_timeout(fast_call, timeout_seconds=5.0) == "ok" + + +# ─── invoke_llm: stats recording ─────────────────────────────────────────── + + +class TestInvokeLlmStats: + """``InvocationStats`` records the right fields.""" + + def test_stats_contain_provider_and_model(self, clean_env, monkeypatch): + monkeypatch.setattr( + provider_mod, + "_PROVIDER_CALL_TABLE", + {"zai_glm": lambda p, m, k: "reply"}, + ) + _, stats = invoke_llm(prompt="hi", model="glm-4.5", api_key="k") + assert stats.provider == "zai_glm" + assert stats.model == "glm-4.5" + assert stats.attempts == 1 + assert stats.elapsed_seconds >= 0.0 + assert stats.timed_out is False + assert stats.raw_response_preview == "reply" + + def test_stats_preview_truncated_to_200_chars(self, clean_env, monkeypatch): + long_reply = "x" * 500 + monkeypatch.setattr( + provider_mod, + "_PROVIDER_CALL_TABLE", + {"zai_glm": lambda p, m, k: long_reply}, + ) + _, stats = invoke_llm(prompt="hi", model="glm-4.5", api_key="k") + assert len(stats.raw_response_preview) == 200 + + def test_stats_as_dict_round_trips(self, clean_env, monkeypatch): + monkeypatch.setattr( + provider_mod, + "_PROVIDER_CALL_TABLE", + {"zai_glm": lambda p, m, k: "reply"}, + ) + _, stats = invoke_llm(prompt="hi", model="glm-4.5", api_key="k") + d = stats.as_dict() + assert set(d.keys()) == { + "provider", + "model", + "attempts", + "elapsed_seconds", + "timed_out", + "raw_response_preview", + } + + +# ─── Provider call wrappers (mocked) ─────────────────────────────────────── + + +class TestProviderCallWrappers: + """Each ``_call_`` wrapper handles import + call correctly.""" + + def test_openai_missing_sdk_raises_not_installed(self, clean_env, monkeypatch): + # Force the import to fail. + monkeypatch.setitem(sys.modules, "openai", None) + with pytest.raises(ProviderNotInstalledError) as exc_info: + provider_mod._call_openai("hi", "gpt-4o", "k") + assert exc_info.value.provider == "openai" + assert exc_info.value.pip_name == "openai" + + def test_anthropic_missing_sdk_raises_not_installed(self, clean_env, monkeypatch): + monkeypatch.setitem(sys.modules, "anthropic", None) + with pytest.raises(ProviderNotInstalledError) as exc_info: + provider_mod._call_anthropic("hi", "claude-3", "k") + assert exc_info.value.provider == "anthropic" + assert exc_info.value.pip_name == "anthropic" + + def test_bedrock_missing_sdk_raises_not_installed(self, clean_env, monkeypatch): + monkeypatch.setitem(sys.modules, "boto3", None) + with pytest.raises(ProviderNotInstalledError) as exc_info: + provider_mod._call_bedrock("hi", "bedrock-anthropic.claude-3", "k") + assert exc_info.value.provider == "bedrock" + assert exc_info.value.pip_name == "boto3" + + def test_google_missing_sdk_raises_not_installed(self, clean_env, monkeypatch): + monkeypatch.setitem(sys.modules, "google", None) + monkeypatch.setitem(sys.modules, "google.generativeai", None) + with pytest.raises(ProviderNotInstalledError) as exc_info: + provider_mod._call_google("hi", "gemini-1.5-pro", "k") + assert exc_info.value.provider == "google" + assert exc_info.value.pip_name == "google-generativeai" + + def test_bedrock_missing_secret_raises_not_configured( + self, clean_env, monkeypatch + ): + """Bedrock needs AWS_SECRET_ACCESS_KEY in addition to the access key.""" + # Mock boto3 so we get past the import, then verify the secret check. + fake_boto3 = mock.MagicMock() + monkeypatch.setitem(sys.modules, "boto3", fake_boto3) + # No AWS_SECRET_ACCESS_KEY set (clean_env stripped it). + with pytest.raises(ProviderNotConfiguredError) as exc_info: + provider_mod._call_bedrock( + "hi", + "bedrock-anthropic.claude-3", + "AKIATEST", + ) + assert exc_info.value.env_var == "AWS_SECRET_ACCESS_KEY" + + def test_bedrock_non_anthropic_model_raises(self, clean_env, monkeypatch): + """Phase 1 only supports anthropic.* models on Bedrock.""" + fake_boto3 = mock.MagicMock() + monkeypatch.setitem(sys.modules, "boto3", fake_boto3) + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret") + with pytest.raises(LLMError, match="body shape not implemented"): + provider_mod._call_bedrock( + "hi", + "bedrock-amazon.nova-pro", + "AKIATEST", + ) + + def test_openai_sdk_exception_wrapped_as_llm_error( + self, clean_env, monkeypatch + ): + """A non-LLMError exception from the SDK becomes a retryable LLMError.""" + fake_openai_mod = mock.MagicMock() + fake_client = mock.MagicMock() + fake_client.chat.completions.create.side_effect = RuntimeError("API error") + fake_openai_mod.OpenAI.return_value = fake_client + monkeypatch.setitem(sys.modules, "openai", fake_openai_mod) + with pytest.raises(LLMError) as exc_info: + provider_mod._call_openai("hi", "gpt-4o", "k") + assert exc_info.value.retryable is True + assert "[llm.openai]" in str(exc_info.value) + + def test_zai_glm_uses_zai_base_url_env(self, clean_env, monkeypatch): + """Z.ai GLM respects the ``ZAI_BASE_URL`` override.""" + fake_openai_mod = mock.MagicMock() + fake_client = mock.MagicMock() + fake_response = mock.MagicMock() + fake_response.choices = [mock.MagicMock()] + fake_response.choices[0].message.content = "glm reply" + fake_client.chat.completions.create.return_value = fake_response + fake_openai_mod.OpenAI.return_value = fake_client + monkeypatch.setitem(sys.modules, "openai", fake_openai_mod) + monkeypatch.setenv("ZAI_BASE_URL", "https://custom.example.com/v1/") + provider_mod._call_zai_glm("hi", "glm-4.5", "k") + fake_openai_mod.OpenAI.assert_called_once() + _, kwargs = fake_openai_mod.OpenAI.call_args + assert kwargs["base_url"] == "https://custom.example.com/v1/" + + def test_zai_glm_default_base_url(self, clean_env, monkeypatch): + fake_openai_mod = mock.MagicMock() + fake_client = mock.MagicMock() + fake_response = mock.MagicMock() + fake_response.choices = [mock.MagicMock()] + fake_response.choices[0].message.content = "glm reply" + fake_client.chat.completions.create.return_value = fake_response + fake_openai_mod.OpenAI.return_value = fake_client + monkeypatch.setitem(sys.modules, "openai", fake_openai_mod) + provider_mod._call_zai_glm("hi", "glm-4.5", "k") + _, kwargs = fake_openai_mod.OpenAI.call_args + assert "open.bigmodel.cn" in kwargs["base_url"] + + +# ─── LLMTool ABC ─────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class _FakeInput(LLMToolInput): + text: str + + def __hash__(self) -> int: + return hash(self.text) + + def __eq__(self, other: object) -> bool: + return isinstance(other, _FakeInput) and self.text == other.text + + +@dataclass +class _FakeOutput(LLMToolOutput): + reply: str + + def __hash__(self) -> int: + return hash(self.reply) + + def __eq__(self, other: object) -> bool: + return isinstance(other, _FakeOutput) and self.reply == other.reply + + +class _EchoTool(LLMTool): + """Test double — echoes the prompt back as the response.""" + + def _get_prompt(self, inp: _FakeInput) -> str: + return f"echo: {inp.text}" + + def _parse_response(self, raw: str, inp: _FakeInput) -> _FakeOutput: + return _FakeOutput(reply=raw) + + +class TestLLMTool: + """The ``LLMTool`` ABC wires prompt → provider → response parsing.""" + + def test_invoke_calls_get_prompt_then_parse_response( + self, clean_env, monkeypatch + ): + captured = {} + + def fake_invoke_llm(**kwargs): + captured["prompt"] = kwargs["prompt"] + stats = base_tool_mod.InvocationStats( + provider="zai_glm", + model="glm-4.5", + attempts=1, + elapsed_seconds=0.01, + ) + return "RAW MODEL OUTPUT", stats + + # Patch the lazy import inside LLMTool.invoke. + monkeypatch.setattr( + "llm.provider.invoke_llm", fake_invoke_llm, raising=True + ) + + tool = _EchoTool(model="glm-4.5", api_key="k") + result = tool.invoke(_FakeInput(text="hello")) + assert captured["prompt"] == "echo: hello" + assert result.output.reply == "RAW MODEL OUTPUT" + assert result.stats.provider == "zai_glm" + + def test_describe_does_not_leak_api_key(self, clean_env): + tool = _EchoTool(model="glm-4.5", api_key="secret-key-12345") + desc = tool.describe() + # API key must never appear in describe() output. + assert "secret-key-12345" not in str(desc) + assert desc["model"] == "glm-4.5" + assert desc["tool"] == "_EchoTool" + + def test_cannot_instantiate_abc_directly(self, clean_env): + with pytest.raises(TypeError): + LLMTool() # type: ignore[abstract] + + def test_subclass_must_implement_get_prompt(self, clean_env): + class _Incomplete(LLMTool): + def _parse_response(self, raw, inp): + return None + with pytest.raises(TypeError): + _Incomplete() # type: ignore[abstract] + + def test_subclass_must_implement_parse_response(self, clean_env): + class _Incomplete(LLMTool): + def _get_prompt(self, inp): + return "" + with pytest.raises(TypeError): + _Incomplete() # type: ignore[abstract] + + def test_frozen_input_is_hashable(self, clean_env): + """Inputs MUST be hashable so the Phase 2 cache can key off them.""" + inp = _FakeInput(text="hello") + assert hash(inp) == hash(_FakeInput(text="hello")) + assert inp == _FakeInput(text="hello") + assert inp != _FakeInput(text="world") + # Set membership works (this is what cache keys need). + assert inp in {_FakeInput(text="hello")} + + +# ─── CLI command ─────────────────────────────────────────────────────────── + + +class TestLlmCommand: + """The ``codelens llm`` CLI command is registered and dispatches correctly.""" + + def test_command_is_registered(self, clean_env): + assert "llm" in COMMAND_REGISTRY + info = COMMAND_REGISTRY["llm"] + assert callable(info["execute"]) + assert callable(info["add_args"]) + + def test_providers_subcommand_returns_six_providers(self, clean_env): + from commands import llm_framework as cmd + args = mock.MagicMock() + args.llm_subcommand = "providers" + args.json = True + result = cmd.execute(args, "") + assert result["status"] == "ok" + assert result["count"] == 6 + names = {p["provider"] for p in result["providers"]} + assert names == ALL_KNOWN_PROVIDERS + + def test_config_subcommand_shows_env_state(self, clean_env, monkeypatch): + from commands import llm_framework as cmd + monkeypatch.setenv("CODELENS_LLM_MODEL", "glm-4.5") + monkeypatch.setenv("ZAI_API_KEY", "k") + args = mock.MagicMock() + args.llm_subcommand = "config" + args.json = True + result = cmd.execute(args, "") + assert result["status"] == "ok" + cfg = result["config"] + assert cfg["model"] == "glm-4.5" + assert cfg["provider_resolved"] == "zai_glm" + assert cfg["api_key_sources"]["ZAI_API_KEY"] is True + + def test_config_does_not_leak_api_key_value(self, clean_env, monkeypatch): + from commands import llm_framework as cmd + monkeypatch.setenv("CODELENS_LLM_MODEL", "glm-4.5") + monkeypatch.setenv("ZAI_API_KEY", "secret-value-xyz") + args = mock.MagicMock() + args.llm_subcommand = "config" + args.json = True + result = cmd.execute(args, "") + # The actual key value must never appear in the output — only + # which env vars are set/unset. + assert "secret-value-xyz" not in str(result) + + def test_ping_subcommand_reports_not_configured(self, clean_env, monkeypatch): + from commands import llm_framework as cmd + monkeypatch.setenv("CODELENS_LLM_MODEL", "glm-4.5") + # No API key set. + args = mock.MagicMock() + args.llm_subcommand = "ping" + args.model = None + args.provider = None + args.timeout = 5.0 + args.json = True + result = cmd.execute(args, "") + assert result["status"] == "not_configured" + assert result["provider"] == "zai_glm" + + def test_ping_subcommand_reports_sdk_missing(self, clean_env, monkeypatch): + from commands import llm_framework as cmd + monkeypatch.setenv("CODELENS_LLM_MODEL", "glm-4.5") + monkeypatch.setenv("ZAI_API_KEY", "k") + # openai SDK is not installed in test env, so we should get sdk_missing. + args = mock.MagicMock() + args.llm_subcommand = "ping" + args.model = None + args.provider = None + args.timeout = 5.0 + args.json = True + result = cmd.execute(args, "") + assert result["status"] == "sdk_missing" + assert result["pip_name"] == "openai" + + def test_no_subcommand_defaults_to_config(self, clean_env, monkeypatch): + from commands import llm_framework as cmd + args = mock.MagicMock() + args.llm_subcommand = None + result = cmd.execute(args, "") + assert result["status"] == "ok" + assert "config" in result + + def test_unknown_subcommand_returns_error(self, clean_env): + from commands import llm_framework as cmd + args = mock.MagicMock() + args.llm_subcommand = "bogus" + result = cmd.execute(args, "") + assert result["status"] == "error" + assert "bogus" in result["error"] + + +# ─── CLI subprocess smoke test ───────────────────────────────────────────── + + +class TestCLISmoke: + """End-to-end: invoke ``codelens llm `` as a real subprocess.""" + + def _run_cli(self, *extra_args): + env = os.environ.copy() + env["PYTHONPATH"] = _SCRIPTS_DIR + env["PYTHONUTF8"] = "1" + return subprocess.run( + [sys.executable, os.path.join(_SCRIPTS_DIR, "codelens.py"), "llm", *extra_args], + capture_output=True, + text=True, + env=env, + timeout=30, + ) + + def test_llm_providers_runs_cleanly(self): + result = self._run_cli("providers", "--json") + assert result.returncode == 0, ( + f"exit={result.returncode}\nstdout={result.stdout}\nstderr={result.stderr}" + ) + + def test_llm_config_runs_cleanly(self): + # Strip all LLM env vars for this subprocess so config is in default state. + env = os.environ.copy() + env["PYTHONPATH"] = _SCRIPTS_DIR + env["PYTHONUTF8"] = "1" + for var in ( + "CODELENS_LLM_PROVIDER", + "CODELENS_LLM_MODEL", + "CODELENS_LLM_API_KEY", + ): + env.pop(var, None) + result = subprocess.run( + [sys.executable, os.path.join(_SCRIPTS_DIR, "codelens.py"), "llm", "config", "--json"], + capture_output=True, + text=True, + env=env, + timeout=30, + ) + assert result.returncode == 0