External/contextzip
1
0
Fork 0
mirror of https://github.com/akadeepesh/contextzip.git synced 2026-10-20 08:36:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: AI-powered file selection via --prompt flag (Gemini 2.5 Flash Lite) (#8)

Browse Source
This commit is contained in:
Deepesh Kumar
2026-05-24 01:31:11 +05:30
committed by GitHub
parent 4f31b18697 43ec8fce27
commit 30568d15a1
10 changed files with 1317 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3
View File
@@ -35,3 +35,6 @@ desktop.ini
# contextzip output files (don't commit test zips)
*_context_*.zip
# contextzip workspace
.contextzip/
README.md
+85 -3
View File
@@ -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.
---
contextzip/ai/__init__.py
+1
View File
@@ -0,0 +1 @@
"""contextzip.ai — AI-powered file selection for --prompt mode."""
contextzip/ai/gemini.py
+283
View File
@@ -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"
contextzip/ai/heuristic.py
+214
View File
@@ -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]
contextzip/ai/selector.py
+162
View File
@@ -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
contextzip/cli.py
+400 -8
View File
@@ -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"
return f"{n:.1f} TB"
contextzip/config.py
+158
View File
@@ -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\\<user>\\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 {}
contextzip/packager.py
+10 -1
View File
@@ -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"
pyproject.toml
+1
View File
@@ -30,6 +30,7 @@ dependencies = [
"click>=8.1",
"rich>=13.0",
"pathspec>=0.12",
"httpx>=0.27",
]
[project.urls]