diff --git a/.gitignore b/.gitignore index 5c8e9fa..984b1ad 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,6 @@ desktop.ini # contextzip output files (don't commit test zips) *_context_*.zip + +# contextzip workspace +.contextzip/ diff --git a/README.md b/README.md index 7c7a24d..af61b1d 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ contextzip eliminates that entirely. Run it from your project root, it detects y - **Smart framework detection** : automatically identifies Node.js, Next.js, Python, Django, FastAPI, Rust, Go, Ruby and applies the right exclusion rules for each - **Respects your `.gitignore`** : patterns from your existing gitignore are honoured automatically - **Git-aware packaging** : package only modified, staged, unstaged, and untracked files with `--git-changes` — perfect for AI review sessions, incremental debugging, and PR workflows +- **AI-powered file selection** : describe your task in plain English with `--prompt` and Gemini selects the minimum relevant files automatically — no manual hunting required - **Persistent workspace** : all generated ZIPs land in a `.contextzip/` directory at your project root, not a random temp folder — discoverable, reusable, and git-ignored automatically - **Warns before it's a problem** : flags large files (≥ 1 MB) and binary files that AI tools can't read, before you waste an upload - **Handles the messy stuff** : dangling symlinks, unreadable files, and files outside the project tree are all caught and reported, never silently dropped @@ -83,10 +84,12 @@ contextzip [OPTIONS] | Option | Description | |---|---| +| `-p`, `--prompt TEXT` | Describe your task in plain English. Gemini AI selects only the relevant files. Requires a free API key (guided setup on first use). | | `-i`, `--include PATH` | Only include files under this path. Repeatable. | | `-e`, `--exclude PATTERN` | Extra exclusion patterns (gitignore syntax). Repeatable. | | `exclude` | Subcommand: exclude specific files/patterns. `contextzip exclude CHANGELOG.md LICENSE .github/` | | `include` | Subcommand: include only specific paths. `contextzip include src/ app/` | +| `config` | Subcommand: manage configuration (API key). `contextzip config --reset-key` | | `--git-changes` | Only include files reported by git as modified, staged, or untracked. | | `-n`, `--dry-run` | Preview what would be included, no ZIP created. | | `-o`, `--output FILE` | Write ZIP to a custom path. Bypasses the `.contextzip/` workspace entirely. | @@ -150,6 +153,78 @@ contextzip --dry-run --verbose contextzip --include src --exclude "**/*.test.ts" ``` +**Let AI select only the files relevant to your task:** +```bash +contextzip --prompt "Change toast color on failed login" +``` + +**Preview AI file selection without creating a ZIP:** +```bash +contextzip --prompt "Refactor auth middleware" --dry-run +``` + +--- + +## AI-powered file selection + +The `--prompt` flag lets you describe a task in plain English. contextzip scans your project, builds a lightweight file map, and asks Gemini to select the minimum set of files needed for that task. The result is a tightly scoped ZIP — only what you'd actually open to make the change. + +```bash +contextzip --prompt "Change toast color on failed login" +``` + +Example selected files: +``` +components/ui/toast.tsx +app/login/page.tsx +lib/auth.ts +``` + +The ZIP also includes a `prompt.txt` that describes the task and lists the selected files. When you drop the ZIP into Claude, ChatGPT, or any other AI tool, it immediately knows what you're trying to do. + +### First-time setup + +`--prompt` requires a free Gemini API key from [Google AI Studio](https://aistudio.google.com/apikey) — no credit card needed. On first use, contextzip will guide you through getting and saving one: + +``` +╭─ Gemini API Key Required ──────────────────────────────────╮ +│ │ +│ --prompt requires a free Gemini API key. │ +│ contextzip uses Google AI Studio — no credit card needed. │ +│ │ +│ Get your free key at: │ +│ https://aistudio.google.com/apikey │ +│ │ +╰─────────────────────────────────────────────────────────────╯ + + Open Google AI Studio in your browser now? [Y/n]: +``` + +Your key is saved to `~/.config/contextzip/config.json` (Linux/macOS) or `%APPDATA%\contextzip\config.json` (Windows) and reused automatically on every subsequent run. + +You can also set the key as an environment variable to skip the config file entirely: + +```bash +export GEMINI_API_KEY=AIza... +``` + +### Managing your API key + +```bash +contextzip config # show current key status +contextzip config --reset-key # clear stored key and re-run setup +contextzip config --show-key-path # print the config file location +``` + +### How file selection works + +1. contextzip scans the project and applies all standard exclusions (no `node_modules`, no build artifacts, no `.env` files) +2. A lightweight file map (relative paths + file sizes) is sent to Gemini along with your task description and detected framework +3. Gemini returns a ranked list of the most relevant files — typically 2–5, never more than 10 +4. Only those files are packaged, plus a `prompt.txt` describing the task + +The goal is minimum viable context. Sending fewer, more relevant files gets better AI responses and wastes fewer tokens. + --- ## Workspace directory @@ -168,11 +243,12 @@ The output filename is determined by how you invoke contextzip: |---|---| | `contextzip` | `.contextzip/codebase.zip` | | `contextzip --git-changes` | `.contextzip/changes.zip` | +| `contextzip --prompt "..."` | `.contextzip/codebase.zip` (AI-scoped) | | `contextzip --output PATH` | `PATH` (workspace bypassed entirely) | All other flags (`--include`, `--exclude`, subcommands) do not affect the filename — they refine what goes into `codebase.zip`. -Each run overwrites the previous file of the same name, so `.contextzip/codebase.zip` is always your latest full snapshot and `.contextzip/changes.zip` is always your latest git-diff context. +Each run overwrites the previous file of the same name, so `.contextzip/codebase.zip` is always your latest snapshot. ### `.gitignore` management @@ -269,12 +345,17 @@ contextzip surfaces issues before they waste your time: contextzip/ ├── contextzip/ │ ├── __init__.py # version string -│ ├── cli.py # Click entry point, all flags, subcommands (exclude, include), rich output +│ ├── cli.py # Click entry point, all flags, subcommands (exclude, include, config), rich output +│ ├── config.py # API key storage (~/.config/contextzip/config.json) │ ├── detector.py # framework/language detection engine │ ├── filters.py # pathspec-based file filtering + ResolveResult │ ├── git.py # git status parsing + changed-file detection │ ├── packager.py # ZIP creation, workspace management, PackageResult │ ├── clipboard.py # tiered clipboard strategy (all platforms) +│ ├── ai/ +│ │ ├── __init__.py +│ │ ├── gemini.py # Gemini API client (file relevance selection) +│ │ └── selector.py # project map builder + AI orchestration + prompt.txt │ └── rules/ │ ├── base.py # universal exclusions (includes .contextzip/) │ ├── node.py # Node.js / Next.js / Vite @@ -329,8 +410,9 @@ That's it. Detection, filtering, and CLI output all pick it up automatically. | [`click`](https://click.palletsprojects.com/) | CLI framework | | [`rich`](https://rich.readthedocs.io/) | Terminal output, panels, progress bars | | [`pathspec`](https://github.com/cpburnz/python-pathspec) | Gitignore-style pattern matching | +| [`httpx`](https://www.python-httpx.org/) | HTTP client for Gemini API calls (required for `--prompt`) | -No other runtime dependencies. Clipboard and folder-open use only stdlib (`subprocess`, `shutil`, `platform`). +`httpx` is only used when `--prompt` is invoked. All other contextzip features work without it, but it must be installed for AI-powered file selection. --- diff --git a/contextzip/ai/__init__.py b/contextzip/ai/__init__.py new file mode 100644 index 0000000..db6f47e --- /dev/null +++ b/contextzip/ai/__init__.py @@ -0,0 +1 @@ +"""contextzip.ai — AI-powered file selection for --prompt mode.""" \ No newline at end of file diff --git a/contextzip/ai/gemini.py b/contextzip/ai/gemini.py new file mode 100644 index 0000000..627886f --- /dev/null +++ b/contextzip/ai/gemini.py @@ -0,0 +1,283 @@ +""" +ai/gemini.py — Thin Gemini API client for contextzip's prompt-aware mode. + +Makes a single POST to the Gemini generateContent endpoint and returns +a ranked list of file paths relevant to the user's task description. + +Design principles: + - Raw httpx calls only — no Google SDK dependency + - Strict JSON output from the model — no markdown, no prose + - Validates every returned path against the real file tree + - Single responsibility: call API, parse response, validate paths +""" + +from __future__ import annotations + +import json +from pathlib import Path + +try: + import httpx + _HTTPX_AVAILABLE = True +except ImportError: + _HTTPX_AVAILABLE = False + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_API_URL = ( + "https://generativelanguage.googleapis.com/v1beta/models" + "/{model}:generateContent?key={key}" +) + +DEFAULT_MODEL = "gemini-2.5-flash-lite" + +# Hard cap: never return more than this many files regardless of model output. +# Minimum context is the goal — the prompt enforces this, but we double-guard. +_MAX_FILES = 12 + +_TIMEOUT_SECONDS = 30 + + +# --------------------------------------------------------------------------- +# Exceptions +# --------------------------------------------------------------------------- + +class GeminiError(Exception): + """Raised when the Gemini API call fails for any reason.""" + + +class GeminiUnavailable(GeminiError): + """Raised when httpx is not installed.""" + + +class GeminiRateLimitError(GeminiError): + """Raised specifically on HTTP 429 — allows callers to trigger fallback.""" + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def select_files( + *, + api_key: str, + prompt: str, + file_tree: list[tuple[str, int]], # [(rel_path, size_bytes), ...] + ecosystem: str, + model: str = DEFAULT_MODEL, +) -> list[str]: + """ + Ask Gemini which files are relevant to *prompt* and return their paths. + + Parameters + ---------- + api_key: + Gemini API key from Google AI Studio. + prompt: + The user's natural-language task description. + file_tree: + All candidate files as (relative_posix_path, size_bytes) tuples. + These should already have contextzip's standard exclusions applied — + no node_modules, no build artifacts, no .env files. + ecosystem: + Human-readable detected framework string, e.g. "Next.js + TypeScript". + Gives the model important context for relevance scoring. + model: + Gemini model identifier. Defaults to gemini-2.0-flash-lite. + + Returns + ------- + list[str] + Relative POSIX paths of the selected files, ordered by relevance + (most relevant first). Always a strict subset of the input file_tree + paths — hallucinated or non-existent paths are silently dropped. + + Raises + ------ + GeminiUnavailable + If httpx is not installed. + GeminiError + On any API or parsing failure. + """ + if not _HTTPX_AVAILABLE: + raise GeminiUnavailable( + "httpx is required for AI-powered selection. " + "Install it with: pip install httpx" + ) + + system_prompt = _build_system_prompt() + user_message = _build_user_message(prompt, file_tree, ecosystem) + + url = _API_URL.format(model=model, key=api_key) + + payload = { + "contents": [ + { + "role": "user", + "parts": [{"text": f"{system_prompt}\n\n{user_message}"}], + } + ], + "generationConfig": { + "temperature": 0.0, # deterministic — this is a ranking task + "maxOutputTokens": 512, # a list of paths needs very few tokens + "responseMimeType": "application/json", + }, + } + + try: + response = httpx.post(url, json=payload, timeout=_TIMEOUT_SECONDS) + except httpx.TimeoutException: + raise GeminiError("Request timed out after 30 seconds.") + except httpx.RequestError as exc: + raise GeminiError(f"Network error: {exc}") + + if response.status_code == 400: + raise GeminiError("Invalid request — check your API key format.") + if response.status_code == 401 or response.status_code == 403: + raise GeminiError( + "API key rejected. Run [cyan]contextzip config --reset-key[/] to update it." + ) + if response.status_code == 429: + raise GeminiRateLimitError( + "Rate limit reached on the free tier (15 req/min). " + "Wait a moment and try again, or your new key may still be activating — " + "Google can take up to 60 seconds after key creation." + ) + if response.status_code != 200: + raise GeminiError( + f"Gemini API returned HTTP {response.status_code}: " + f"{response.text[:200]}" + ) + + return _parse_response(response.json(), file_tree) + + +# --------------------------------------------------------------------------- +# Prompt construction +# --------------------------------------------------------------------------- + +def _build_system_prompt() -> str: + return """\ +You are a precise file relevance assistant for software projects. + +Your only job: given a developer's task description and a project file tree, +return the MINIMUM set of files a developer would need to open to complete +that task. Think like a senior engineer doing a surgical code change — +open only what you must touch or read to understand the change. + +Rules you must follow: +- Return ONLY a JSON array of file path strings. No explanation, no markdown, + no extra keys. Example: ["src/auth.ts", "components/Toast.tsx"] +- Be ruthless about exclusion. If a file is not directly relevant to the + stated task, leave it out. Err heavily on the side of fewer files. +- Prefer files that will be MODIFIED over files that are merely referenced. +- Config files, test files, and documentation should only appear if the + task explicitly concerns them. +- Never return more than 10 files. For most tasks 2–5 files is correct. +- Order by relevance: most directly relevant file first.\ +""" + + +def _build_user_message( + prompt: str, + file_tree: list[tuple[str, int]], + ecosystem: str, +) -> str: + tree_lines = "\n".join( + f"{path} ({_human_size(size)})" + for path, size in file_tree + ) + return f"""\ +Framework: {ecosystem} + +Task: {prompt} + +Project files (exclusions already applied): +{tree_lines} + +Return only a JSON array of the most relevant file paths for this task.\ +""" + + +# --------------------------------------------------------------------------- +# Response parsing and validation +# --------------------------------------------------------------------------- + +def _parse_response( + data: dict, + file_tree: list[tuple[str, int]], +) -> list[str]: + """ + Extract and validate the file list from the Gemini API response. + + - Parses the JSON array from the model's text output + - Drops any path the model hallucinated (not in the real file tree) + - Enforces the _MAX_FILES hard cap + - Warns (via exception) if the model returned mostly invalid paths + """ + # Navigate the Gemini response structure + try: + text = ( + data["candidates"][0]["content"]["parts"][0]["text"] + ) + except (KeyError, IndexError) as exc: + raise GeminiError(f"Unexpected API response structure: {exc}\n{data}") + + # Strip any accidental markdown fences the model might add + text = text.strip().strip("`").strip() + if text.startswith("json"): + text = text[4:].strip() + + try: + raw_paths: list = json.loads(text) + except json.JSONDecodeError as exc: + raise GeminiError( + f"Model returned non-JSON output: {exc}\nRaw output: {text[:300]}" + ) + + if not isinstance(raw_paths, list): + raise GeminiError( + f"Expected a JSON array, got {type(raw_paths).__name__}: {text[:200]}" + ) + + # Build a set of valid paths for O(1) lookup + valid_paths: set[str] = {path for path, _ in file_tree} + + validated: list[str] = [] + hallucinated = 0 + + for item in raw_paths: + if not isinstance(item, str): + continue + # Normalise separators (model may return backslashes on Windows prompts) + normalised = item.replace("\\", "/").strip() + if normalised in valid_paths: + validated.append(normalised) + else: + hallucinated += 1 + + # Warn if the model was mostly making things up + total_returned = len(raw_paths) + if total_returned > 0 and hallucinated / total_returned > 0.5: + raise GeminiError( + f"Model returned {hallucinated}/{total_returned} non-existent paths. " + "This may indicate a model or prompt issue. Try a more specific prompt." + ) + + # Enforce hard cap + return validated[:_MAX_FILES] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _human_size(n: int) -> str: + for unit in ("B", "KB", "MB"): + if n < 1024: + return f"{n:.0f} {unit}" + n /= 1024 + return f"{n:.1f} GB" diff --git a/contextzip/ai/heuristic.py b/contextzip/ai/heuristic.py new file mode 100644 index 0000000..ee0ebd3 --- /dev/null +++ b/contextzip/ai/heuristic.py @@ -0,0 +1,214 @@ +""" +ai/heuristic.py — Keyword-based file relevance scorer. + +Used as a fallback when the Gemini API is unavailable (rate limit, network +error, no key). Scores each file purely from the prompt text and file path — +no API calls, no dependencies beyond stdlib. + +This is intentionally simple and transparent. It will never be as accurate +as Gemini, but it is always available and produces reasonable results for +common cases like "update the login page" or "fix the toast component". + +Scoring model +───────────── +Each file receives a score based on: + + 1. Token overlap — how many meaningful words from the prompt appear in + the file's path components (stem + parent dirs). + Longer matches score higher. + + 2. Directory bias — source-like directories (src/, app/, lib/, components/, + utils/, routes/, pages/, server/, api/) score higher. + Test and documentation directories score lower. + + 3. Extension bias — code files score higher than config, lock, or data files. + +Files with a score of zero are excluded entirely. The top-N results are +returned, where N is capped at MAX_FILES. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +# Never return more files than this regardless of score +MAX_FILES = 8 + +# Words that carry no signal — filtered out before scoring +_STOPWORDS = frozenset({ + "i", "a", "an", "the", "to", "in", "on", "at", "of", "for", + "and", "or", "but", "is", "it", "be", "do", "my", "this", + "that", "with", "from", "want", "need", "make", "update", + "change", "fix", "add", "get", "use", "set", "new", "old", + "file", "code", "function", "method", "class", "variable", +}) + +# Directory names that suggest source code worth including +_SOURCE_DIRS = frozenset({ + "src", "app", "lib", "libs", "components", "component", + "utils", "util", "helpers", "helper", "routes", "route", + "pages", "page", "server", "api", "services", "service", + "hooks", "hook", "store", "stores", "context", "contexts", + "middleware", "handlers", "handler", "controllers", "controller", + "views", "view", "models", "model", "core", "common", "shared", +}) + +# Directory names that suggest low relevance for most coding tasks +_NOISE_DIRS = frozenset({ + "test", "tests", "__tests__", "spec", "specs", + "docs", "doc", "documentation", "examples", "example", + "fixtures", "mocks", "mock", "stubs", "stub", + "scripts", "bin", "dist", "build", "out", "coverage", +}) + +# Extensions that suggest runnable/editable source code +_CODE_EXTENSIONS = frozenset({ + ".py", ".ts", ".tsx", ".js", ".jsx", ".go", ".rs", ".rb", + ".java", ".kt", ".swift", ".cs", ".cpp", ".c", ".h", + ".vue", ".svelte", ".astro", +}) + +# Extensions that are config-like (lower relevance unless prompt mentions them) +_CONFIG_EXTENSIONS = frozenset({ + ".json", ".toml", ".yaml", ".yml", ".ini", ".cfg", ".env", + ".md", ".txt", ".lock", ".sum", +}) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def select_files( + *, + prompt: str, + file_tree: list[tuple[str, int]], +) -> list[str]: + """ + Score and rank *file_tree* entries by relevance to *prompt*. + + Parameters + ---------- + prompt: + The user's natural-language task description. + file_tree: + Candidate files as (relative_posix_path, size_bytes) tuples. + Standard exclusions must already have been applied by the caller. + + Returns + ------- + list[str] + Relative POSIX paths of the top-scoring files, best first. + Files scoring zero are excluded. Result is capped at MAX_FILES. + """ + tokens = _tokenize(prompt) + if not tokens: + # Degenerate prompt — return nothing rather than random files + return [] + + scored: list[tuple[float, str]] = [] + + for rel_path, size_bytes in file_tree: + score = _score(rel_path, tokens, size_bytes) + if score > 0: + scored.append((score, rel_path)) + + # Sort descending by score, then alphabetically for determinism on ties + scored.sort(key=lambda x: (-x[0], x[1])) + + return [path for _, path in scored[:MAX_FILES]] + + +# --------------------------------------------------------------------------- +# Scoring +# --------------------------------------------------------------------------- + +def _score(rel_path: str, tokens: set[str], size_bytes: int) -> float: + """Compute a relevance score for a single file path.""" + path = Path(rel_path) + stem = path.stem.lower() + suffix = path.suffix.lower() + parts = [p.lower() for p in path.parts] + dirs = parts[:-1] # everything except the filename itself + + score = 0.0 + + # ── 1. Token overlap ───────────────────────────────────────────────────── + # Split the stem by common separators (camelCase, kebab-case, snake_case) + stem_words = set(_split_identifier(stem)) + dir_words: set[str] = set() + for d in dirs: + dir_words.update(_split_identifier(d)) + + # Direct hits in filename stem score highest + stem_hits = tokens & stem_words + score += len(stem_hits) * 3.0 + + # Partial substring matches in stem (e.g. "toast" in "useToast") + for token in tokens: + if len(token) >= 4 and token in stem and token not in stem_hits: + score += 1.0 + + # Hits in parent directory names score lower + dir_hits = tokens & dir_words + score += len(dir_hits) * 1.5 + + # No token overlap at all → zero score, file is irrelevant + if score == 0.0: + return 0.0 + + # ── 2. Directory bias ──────────────────────────────────────────────────── + dir_set = set(dirs) + if dir_set & _SOURCE_DIRS: + score *= 1.4 + if dir_set & _NOISE_DIRS: + score *= 0.4 + + # ── 3. Extension bias ──────────────────────────────────────────────────── + if suffix in _CODE_EXTENSIONS: + score *= 1.2 + elif suffix in _CONFIG_EXTENSIONS: + score *= 0.7 + + # ── 4. Penalise very large files ───────────────────────────────────────── + # Files over 100 KB are less likely to be what you want to change + if size_bytes > 100_000: + score *= 0.6 + + return score + + +# --------------------------------------------------------------------------- +# Text helpers +# --------------------------------------------------------------------------- + +def _tokenize(text: str) -> set[str]: + """ + Extract meaningful lowercase tokens from *text*. + + - Splits on whitespace and punctuation + - Removes stopwords + - Keeps only tokens of 3+ characters + """ + raw = re.findall(r"[a-zA-Z]+", text.lower()) + return {w for w in raw if len(w) >= 3 and w not in _STOPWORDS} + + +def _split_identifier(name: str) -> list[str]: + """ + Split a file/directory name into component words. + + Handles: kebab-case, snake_case, camelCase, PascalCase. + + Examples: + "LoginPage" → ["login", "page"] + "use-toast" → ["use", "toast"] + "auth_utils" → ["auth", "utils"] + "apiClient" → ["api", "client"] + """ + # Insert space before uppercase letters following lowercase (camelCase) + spaced = re.sub(r"([a-z])([A-Z])", r"\1 \2", name) + # Split on non-alphanumeric separators + parts = re.split(r"[^a-zA-Z0-9]+", spaced) + return [p.lower() for p in parts if p] diff --git a/contextzip/ai/selector.py b/contextzip/ai/selector.py new file mode 100644 index 0000000..10f191f --- /dev/null +++ b/contextzip/ai/selector.py @@ -0,0 +1,162 @@ +""" +ai/selector.py — Orchestrates AI-powered file selection for --prompt mode. + +Responsibilities: + 1. Build a lean project map from the already-filtered file list + 2. Call the Gemini client (with heuristic fallback on rate limit) + 3. Return the selected subset as Path objects ready for packaging + 4. Generate the prompt.txt content to include in the ZIP +""" + +from __future__ import annotations + +from pathlib import Path + +from contextzip.ai.gemini import select_files, GeminiError, GeminiRateLimitError +from contextzip.filters import ResolveResult + +# Sentinels — tell the CLI which selection path was taken +USED_GEMINI = "gemini" +USED_HEURISTIC = "heuristic" + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def ai_select( + *, + resolved: ResolveResult, + project_dir: Path, + prompt: str, + ecosystem: str, + api_key: str, +) -> tuple[list[Path], str, str]: + """ + Select the minimum relevant files for *prompt* from *resolved.included*. + + Returns + ------- + (selected_paths, prompt_txt, method) + selected_paths : list[Path] of absolute paths chosen + prompt_txt : string content ready to write as prompt.txt in the ZIP + method : USED_GEMINI or USED_HEURISTIC + """ + from contextzip.ai import heuristic as _heuristic + + file_tree = _build_file_tree(resolved.included, project_dir) + method = USED_GEMINI + + try: + selected_rel = select_files( + api_key=api_key, + prompt=prompt, + file_tree=file_tree, + ecosystem=ecosystem, + ) + except GeminiRateLimitError: + # Confirmed HTTP 429 only — fall back to keyword heuristic + selected_rel = _heuristic.select_files( + prompt=prompt, + file_tree=file_tree, + ) + method = USED_HEURISTIC + # All other GeminiErrors (bad key, network, unexpected status) propagate + # up to _run_ai_selection in cli.py which displays the error and exits. + + # Map relative path strings back to absolute Path objects + rel_to_abs: dict[str, Path] = { + p.relative_to(project_dir).as_posix(): p + for p in resolved.included + } + selected_paths = [ + rel_to_abs[rel] + for rel in selected_rel + if rel in rel_to_abs + ] + + prompt_txt = _build_prompt_txt(prompt, selected_rel, ecosystem, method) + + return selected_paths, prompt_txt, method + + +def build_prompt_only_txt(prompt: str, ecosystem: str) -> str: + """Build a minimal prompt.txt when no AI selection was performed.""" + return _build_prompt_txt(prompt, [], ecosystem, USED_GEMINI) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +def _build_file_tree( + included: list[Path], + project_dir: Path, +) -> list[tuple[str, int]]: + """ + Convert the included file list into (relative_posix_path, size_bytes) tuples. + + Binary files and files that can't be stat'd are silently skipped — + the model can't reason about them anyway. + """ + tree: list[tuple[str, int]] = [] + + for abs_path in included: + try: + rel = abs_path.relative_to(project_dir).as_posix() + size = abs_path.stat().st_size + except (ValueError, OSError): + continue + + if _is_binary(abs_path): + continue + + tree.append((rel, size)) + + return tree + + +def _build_prompt_txt( + prompt: str, + selected_rel: list[str], + ecosystem: str, + method: str, +) -> str: + """ + Build the prompt.txt to include inside the ZIP. + + Any AI tool that receives the ZIP immediately sees the task description, + the framework, and exactly which files were selected and why. + """ + from contextzip import __version__ + + selector_label = ( + "contextzip AI (Gemini)" + if method == USED_GEMINI + else "contextzip (keyword heuristic — Gemini was rate limited)" + ) + + lines: list[str] = [ + f"Task: {prompt}", + "", + f"Framework: {ecosystem}", + "", + ] + + if selected_rel: + lines.append(f"Files selected by {selector_label}:") + for rel in selected_rel: + lines.append(f" - {rel}") + else: + lines.append("(No files selected)") + + return "\n".join(lines) + "\n" + + +def _is_binary(path: Path, peek: int = 512) -> bool: + """Return True if the file appears to be binary (null bytes in first peek bytes).""" + try: + with path.open("rb") as fh: + return b"\x00" in fh.read(peek) + except OSError: + return False diff --git a/contextzip/cli.py b/contextzip/cli.py index 34d88ec..1ef30a4 100644 --- a/contextzip/cli.py +++ b/contextzip/cli.py @@ -27,6 +27,7 @@ the flags naturally follow the verb, matching the git / docker / cargo UX: from __future__ import annotations import os +import webbrowser from pathlib import Path import click @@ -36,6 +37,7 @@ from rich.table import Table from rich import box from contextzip import __version__ +from contextzip.config import get_api_key, save_api_key, delete_api_key, config_path, diagnose_api_key from contextzip.detector import detect from contextzip.filters import ( build_spec, @@ -66,6 +68,15 @@ def _modifier_options(f): contextzip exclude CHANGELOG.md --dry-run --verbose """ decorators = [ + click.option( + "--prompt", "-p", + default=None, metavar="TEXT", + help=( + "Describe your task in natural language. contextzip uses Gemini AI " + "to select only the files relevant to that task. " + "Requires a free Gemini API key (you will be guided on first use)." + ), + ), click.option( "--dry-run", "-n", is_flag=True, default=False, @@ -136,6 +147,7 @@ def main( ctx: click.Context, include: tuple[str, ...], exclude: tuple[str, ...], + prompt: str | None, dry_run: bool, output: str | None, no_clipboard: bool, @@ -165,6 +177,7 @@ def main( _run( extra_exclude=list(exclude), include_only=list(include), + prompt=prompt, dry_run=dry_run, output=output, no_clipboard=no_clipboard, @@ -183,6 +196,7 @@ def main( @_modifier_options def cmd_exclude( patterns: tuple[str, ...], + prompt: str | None, dry_run: bool, output: str | None, no_clipboard: bool, @@ -206,6 +220,7 @@ def cmd_exclude( _run( extra_exclude=list(patterns), include_only=None, + prompt=prompt, dry_run=dry_run, output=output, no_clipboard=no_clipboard, @@ -224,6 +239,7 @@ def cmd_exclude( @_modifier_options def cmd_include( paths: tuple[str, ...], + prompt: str | None, dry_run: bool, output: str | None, no_clipboard: bool, @@ -246,6 +262,7 @@ def cmd_include( _run( extra_exclude=None, include_only=list(paths), + prompt=prompt, dry_run=dry_run, output=output, no_clipboard=no_clipboard, @@ -263,6 +280,7 @@ def _run( *, extra_exclude: list[str] | None, include_only: list[str] | None, + prompt: str | None, dry_run: bool, output: str | None, no_clipboard: bool, @@ -404,14 +422,33 @@ def _run( # ── Dry run ────────────────────────────────────────────────────────────── if dry_run: - console.print() - console.print( - Panel.fit( - "[yellow]Dry run — no ZIP created.[/]\n" - "[dim]Remove --dry-run to produce the archive.[/]", - border_style="yellow", padding=(0, 2), + # If --prompt is set, still run AI selection so the user can preview + # which files would be chosen — but don't create the ZIP. + if prompt: + diagnosis = diagnose_api_key() + if diagnosis: + console.print(f"\n [yellow]⚠[/] {diagnosis}\n") + api_key = get_api_key() + if not api_key: + api_key = _onboard_api_key() + if not api_key: + raise SystemExit(0) + _run_ai_selection_preview( + resolved=resolved, + project_dir=project_dir, + prompt=prompt, + ecosystem=detection.display_name, + api_key=api_key, + ) + else: + console.print() + console.print( + Panel.fit( + "[yellow]Dry run — no ZIP created.[/]\n" + "[dim]Remove --dry-run to produce the archive.[/]", + border_style="yellow", padding=(0, 2), + ) ) - ) return if not resolved.included: @@ -421,6 +458,38 @@ def _run( ) return + # ── AI-powered file selection (--prompt mode) ──────────────────────────── + prompt_txt: str | None = None + + if prompt: + diagnosis = diagnose_api_key() + if diagnosis: + console.print(f"\n [yellow]⚠[/] {diagnosis}\n") + api_key = get_api_key() + if not api_key: + api_key = _onboard_api_key() + if not api_key: + raise SystemExit(0) + + selected_paths, prompt_txt = _run_ai_selection( + resolved=resolved, + project_dir=project_dir, + prompt=prompt, + ecosystem=detection.display_name, + api_key=api_key, + ) + + if not selected_paths: + console.print( + "\n[red]AI selection returned no files.[/] " + "Try a more specific prompt, or run without [cyan]--prompt[/] " + "to package the full project." + ) + raise SystemExit(1) + + # Replace the resolved file list with only the AI-selected files + resolved.included = selected_paths + # ── Create ZIP ─────────────────────────────────────────────────────────── console.print() output_path = Path(output).resolve() if output else None @@ -432,6 +501,7 @@ def _run( output_path=output_path, console=console, git_changes=git_changes, + prompt_txt=prompt_txt, ) except Exception as exc: console.print(f"\n[red]Failed to create ZIP:[/] {exc}") @@ -458,6 +528,295 @@ def _run( _print_clipboard_result(cb) + +# --------------------------------------------------------------------------- +# API key onboarding +# --------------------------------------------------------------------------- + +_GEMINI_KEY_URL = "https://aistudio.google.com/apikey" + + +def _open_browser_silent(url: str) -> None: + """ + Open *url* in the default browser, suppressing all stderr output. + + On Linux, xdg-open frequently emits KDE/DBus/MIME warnings to stderr + that are completely unrelated to contextzip and appear right on top of + the API key prompt, confusing users. We redirect both stdout and stderr + to /dev/null to keep the terminal clean. + """ + import subprocess + try: + subprocess.Popen( + ["xdg-open", url], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return + except (FileNotFoundError, OSError): + pass + # macOS / Windows fallback + try: + webbrowser.open(url) + except Exception: + pass + + +def _onboard_api_key() -> str | None: + """ + Guide the user through getting and saving a Gemini API key. + + Returns the key string if successfully configured, or None if the + user declined or provided nothing. + """ + console.print() + console.print( + Panel( + " [bold]--prompt[/] requires a free Gemini API key.\n\n" + " contextzip uses [bold cyan]Google AI Studio[/] — " + "no credit card needed.\n\n" + f" [dim]Get your free key at:[/]\n" + f" [bold cyan link={_GEMINI_KEY_URL}]{_GEMINI_KEY_URL}[/]", + title="[bold yellow]Gemini API Key Required[/]", + border_style="yellow", + padding=(0, 2), + ) + ) + console.print() + + open_browser = click.confirm( + " Open Google AI Studio in your browser now?", + default=True, + ) + if open_browser: + _open_browser_silent(_GEMINI_KEY_URL) + console.print( + " [dim]Browser opened. Generate a key, then come back here.[/]" + ) + + console.print() + console.print(" [dim]─────────────────────────────────────────────[/]") + console.print(" Once you have your key, paste it below and press [bold]Enter[/].") + console.print(" [dim]─────────────────────────────────────────────[/]") + console.print() + key = click.prompt( + " Paste your API key", + hide_input=True, + prompt_suffix=" › ", + ).strip() + + if not key: + console.print("\n [red]No key provided. Exiting.[/]") + return None + + # Basic sanity check — Gemini keys start with "AIza" + if not key.startswith("AIza"): + console.print( + "\n [yellow]⚠[/] That doesn't look like a Gemini key " + "(expected it to start with [cyan]AIza[/]).\n" + " [dim]Double-check you copied the full key from AI Studio.[/]" + ) + if not click.confirm("\n Save it anyway?", default=False): + return None + + try: + save_api_key(key) + console.print( + f"\n [green]✓[/] Key saved to " + f"[dim]{config_path()}[/]\n" + f" [dim]You won't be asked again. " + f"Run [cyan]contextzip config --reset-key[/] to change it.[/]" + ) + except OSError as exc: + console.print( + f"\n [yellow]⚠[/] Could not save key to disk ([dim]{exc}[/]).\n" + f" [dim]Set [cyan]GEMINI_API_KEY[/] as an environment variable " + f"to avoid this prompt next time.[/]" + ) + + console.print() + return key + + +# --------------------------------------------------------------------------- +# AI selection helpers +# --------------------------------------------------------------------------- + +def _run_ai_selection( + *, + resolved, + project_dir: Path, + prompt: str, + ecosystem: str, + api_key: str, +) -> tuple[list[Path], str]: + """ + Call the AI selector and return (selected_paths, prompt_txt). + Handles display of progress, heuristic fallback warning, and errors. + Raises SystemExit on hard failure. + """ + from contextzip.ai.selector import ai_select, USED_HEURISTIC + from contextzip.ai.gemini import GeminiError + + console.print() + with console.status( + "[cyan]Asking Gemini to select relevant files…[/]", spinner="dots" + ): + try: + selected_paths, prompt_txt, method = ai_select( + resolved=resolved, + project_dir=project_dir, + prompt=prompt, + ecosystem=ecosystem, + api_key=api_key, + ) + except GeminiError as exc: + console.print( + Panel.fit( + f"[red]Gemini error:[/] {exc}", + border_style="red", padding=(0, 2), + ) + ) + raise SystemExit(1) + + if method == USED_HEURISTIC: + console.print( + " [yellow]⚠[/] [dim]Gemini rate limited — used keyword heuristic instead. " + "Results may be less precise. Try again in a moment for AI selection.[/]" + ) + + _print_ai_selection(selected_paths, project_dir, prompt) + return selected_paths, prompt_txt + + +def _run_ai_selection_preview( + *, + resolved, + project_dir: Path, + prompt: str, + ecosystem: str, + api_key: str, +) -> None: + """ + Run AI selection and show the result without creating a ZIP. + Used when --dry-run and --prompt are combined. + """ + from contextzip.ai.selector import ai_select, USED_HEURISTIC + from contextzip.ai.gemini import GeminiError + + console.print() + with console.status( + "[cyan]Asking Gemini to select relevant files…[/]", spinner="dots" + ): + try: + selected_paths, _, method = ai_select( + resolved=resolved, + project_dir=project_dir, + prompt=prompt, + ecosystem=ecosystem, + api_key=api_key, + ) + except GeminiError as exc: + console.print( + Panel.fit( + f"[red]Gemini error:[/] {exc}", + border_style="red", padding=(0, 2), + ) + ) + raise SystemExit(1) + + if method == USED_HEURISTIC: + console.print( + " [yellow]⚠[/] [dim]Gemini rate limited — used keyword heuristic instead. " + "Results may be less precise. Try again in a moment for AI selection.[/]" + ) + + _print_ai_selection(selected_paths, project_dir, prompt) + console.print() + console.print( + Panel.fit( + "[yellow]Dry run — no ZIP created.[/]\n" + "[dim]Remove --dry-run to package these files.[/]", + border_style="yellow", padding=(0, 2), + ) + ) + + +# --------------------------------------------------------------------------- +# Subcommand: config +# --------------------------------------------------------------------------- + +@main.command("config") +@click.option( + "--reset-key", + is_flag=True, default=False, + help="Clear the stored Gemini API key and re-run the setup prompt.", +) +@click.option( + "--show-key-path", + is_flag=True, default=False, + help="Print the path to the config file and exit.", +) +def cmd_config(reset_key: bool, show_key_path: bool) -> None: + """ + Manage contextzip configuration. + + \b + EXAMPLES + contextzip config --reset-key # clear stored API key, re-onboard + contextzip config --show-key-path # print config file location + """ + if show_key_path: + console.print(f"\n [dim]Config file:[/] [cyan]{config_path()}[/]\n") + return + + if reset_key: + removed = delete_api_key() + if removed: + console.print( + "\n [green]✓[/] API key removed from " + f"[dim]{config_path()}[/]" + ) + else: + console.print("\n [dim]No API key was stored.[/]") + + console.print() + new_key = _onboard_api_key() + if not new_key: + console.print( + " [dim]You can set [cyan]GEMINI_API_KEY[/] as an environment " + "variable instead.[/]\n" + ) + return + + # Default: show current config status + key = get_api_key() + from_env = bool(os.environ.get("GEMINI_API_KEY", "").strip()) + + if key: + masked = key[:8] + "…" + key[-4:] if len(key) > 12 else "****" + source = "[dim](from environment variable)[/]" if from_env else f"[dim]({config_path()})[/]" + console.print( + Panel( + f" [dim]Gemini API key:[/] [green]{masked}[/] {source}\n\n" + " [dim]Run [cyan]contextzip config --reset-key[/] to change it.[/]", + title="[bold]contextzip config[/]", + border_style="cyan", padding=(0, 2), + ) + ) + else: + console.print( + Panel( + " No Gemini API key configured.\n\n" + " Run [cyan]contextzip --prompt \"your task\"[/] to set one up,\n" + " or [cyan]contextzip config --reset-key[/] to go through setup now.", + title="[bold]contextzip config[/]", + border_style="yellow", padding=(0, 2), + ) + ) + console.print() + + # --------------------------------------------------------------------------- # Normalisation # --------------------------------------------------------------------------- @@ -628,6 +987,39 @@ def _print_package_result(result) -> None: ) +def _print_ai_selection( + selected_paths: list[Path], + project_dir: Path, + prompt: str, +) -> None: + """Display a panel showing the files Gemini selected.""" + table = Table(box=box.ROUNDED, show_header=False, padding=(0, 2)) + table.add_column(style="cyan") + table.add_column(style="dim") + + for p in selected_paths: + try: + rel = p.relative_to(project_dir).as_posix() + size = _human_size(p.stat().st_size) + except (ValueError, OSError): + rel = str(p) + size = "?" + table.add_row(rel, size) + + console.print( + Panel( + table, + title=f"[bold cyan]AI Selected — [dim]\"{prompt}\"[/][/]", + border_style="cyan", padding=(0, 1), + ) + ) + console.print( + f" [dim]↳ {len(selected_paths)} file" + f"{'s' if len(selected_paths) != 1 else ''} selected by Gemini " + f"· prompt.txt will be included in the ZIP[/]" + ) + + def _print_clipboard_result(cb) -> None: tier_style = { Tier.FILE_ON_CLIPBOARD: ("green", "✓ Ready to paste"), @@ -649,4 +1041,4 @@ def _human_size(n: int) -> str: if n < 1024: return f"{n:.0f} {unit}" n /= 1024 - return f"{n:.1f} TB" \ No newline at end of file + return f"{n:.1f} TB" diff --git a/contextzip/config.py b/contextzip/config.py new file mode 100644 index 0000000..baf1b52 --- /dev/null +++ b/contextzip/config.py @@ -0,0 +1,158 @@ +""" +config.py — Persistent configuration storage for contextzip. + +Stores user settings (e.g. Gemini API key) in the platform-appropriate +config directory: + + Linux/macOS : ~/.config/contextzip/config.json + Windows : C:\\Users\\\\AppData\\Roaming\\contextzip\\config.json + +This follows the XDG standard and mirrors what tools like gh, gcloud, +and aws configure use. Never uses temp — config persists across sessions. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + + +# --------------------------------------------------------------------------- +# Config directory resolution (no platformdirs dependency) +# --------------------------------------------------------------------------- + +def _config_dir() -> Path: + """ + Return the platform-appropriate config directory for contextzip. + + Mirrors what platformdirs.user_config_dir("contextzip") would return, + but without the extra dependency. + """ + if os.name == "nt": + # Windows: %APPDATA%\\contextzip + base = Path(os.environ.get("APPDATA", Path.home() / "AppData" / "Roaming")) + else: + # Linux / macOS: $XDG_CONFIG_HOME/contextzip or ~/.config/contextzip + xdg = os.environ.get("XDG_CONFIG_HOME", "") + base = Path(xdg) if xdg else Path.home() / ".config" + + return base / "contextzip" + + +_CONFIG_FILE = _config_dir() / "config.json" + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def get_api_key() -> str | None: + """ + Return the stored Gemini API key, or None if not configured. + + Checks in order: + 1. GEMINI_API_KEY environment variable (must look like a real key) + 2. Persisted config file + + A key is considered valid only if it starts with "AIza" — this prevents + a garbage or placeholder env var from bypassing the onboarding flow and + causing a cryptic API error downstream. + """ + env_key = os.environ.get("GEMINI_API_KEY", "").strip() + if env_key and env_key.startswith("AIza"): + return env_key + + # Env var present but malformed — don't silently use it, fall through + # to config file. If config file also has nothing, onboarding triggers. + + try: + data = _read_config() + key = data.get("gemini_api_key", "").strip() + # Same validation as env var — must look like a real Gemini key + return key if (key and key.startswith("AIza")) else None + except Exception: + return None + + +def save_api_key(key: str) -> None: + """ + Persist *key* to the config file. + + Creates the config directory if it doesn't exist. + Raises OSError if the directory or file cannot be written. + """ + _CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True) + data = _read_config() + data["gemini_api_key"] = key.strip() + _CONFIG_FILE.write_text( + json.dumps(data, indent=2) + "\n", + encoding="utf-8", + ) + + +def delete_api_key() -> bool: + """ + Remove the stored API key from config. + + Returns True if a key was present and removed, False if there was + nothing to remove. + """ + try: + data = _read_config() + if "gemini_api_key" not in data: + return False + del data["gemini_api_key"] + _CONFIG_FILE.write_text( + json.dumps(data, indent=2) + "\n", + encoding="utf-8", + ) + return True + except Exception: + return False + + +def diagnose_api_key() -> str | None: + """ + Return a human-readable explanation if a key exists but is invalid, + or None if everything is fine (or nothing is configured at all). + + Used by the CLI to surface "you have a bad key" before onboarding. + """ + env_key = os.environ.get("GEMINI_API_KEY", "").strip() + if env_key and not env_key.startswith("AIza"): + return ( + f"GEMINI_API_KEY is set in your environment but doesn't look like " + f"a valid Gemini key (got: {env_key[:12]}…). " + f"Unset it or replace it with a key from https://aistudio.google.com/apikey" + ) + try: + data = _read_config() + key = data.get("gemini_api_key", "").strip() + if key and not key.startswith("AIza"): + return ( + f"Saved API key in {_CONFIG_FILE} doesn't look valid. " + f"Run [cyan]contextzip config --reset-key[/] to replace it." + ) + except Exception: + pass + return None + + +def config_path() -> Path: + """Return the path to the config file (may or may not exist yet).""" + return _CONFIG_FILE + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + +def _read_config() -> dict: + """Read and parse the config file. Returns {} if missing or malformed.""" + if not _CONFIG_FILE.is_file(): + return {} + try: + return json.loads(_CONFIG_FILE.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + return {} diff --git a/contextzip/packager.py b/contextzip/packager.py index 2c2e1cf..6d80c31 100644 --- a/contextzip/packager.py +++ b/contextzip/packager.py @@ -76,6 +76,7 @@ def create_zip( output_path: Path | None, console: Console, git_changes: bool = False, + prompt_txt: str | None = None, ) -> PackageResult: """ Write the included files from *resolve_result* into a ZIP archive. @@ -87,6 +88,10 @@ def create_zip( directory (created automatically) at the git root, or the CWD if no git repository is detected. + If *prompt_txt* is provided (set when --prompt is used), a ``prompt.txt`` + file is written as the first entry in the ZIP. Any AI tool that receives + the ZIP will immediately see the task description and selected file list. + Returns a :class:`PackageResult` with compression stats and any files that had to be skipped during writing (e.g. permission denied). """ @@ -125,6 +130,10 @@ def create_zip( compression=zipfile.ZIP_DEFLATED, compresslevel=6, ) as zf: + # Write prompt.txt first so it's the first thing AI tools see + if prompt_txt is not None: + zf.writestr("prompt.txt", prompt_txt.encode("utf-8")) + for abs_path in included: if not abs_path.is_file(): continue @@ -272,4 +281,4 @@ def _workspace_output_path( def _safe_name(name: str) -> str: safe = "".join(c if c.isalnum() or c in "-_." else "_" for c in name) - return safe[:48] or "project" + return safe[:48] or "project" \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 3decd35..1479b0d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "click>=8.1", "rich>=13.0", "pathspec>=0.12", + "httpx>=0.27", ] [project.urls]