feat: apply ai preferences to file selection

contextzip/ai/gemini.py

@@ -70,6 +70,7 @@ def select_files(
     file_tree: list[tuple[str, int]],  # [(rel_path, size_bytes), ...]
     ecosystem: str,
     model: str = DEFAULT_MODEL,
+    max_files: int | None = None,
 ) -> list[str]:
     """
     Ask Gemini which files are relevant to *prompt* and return their paths.
@@ -89,6 +90,11 @@ def select_files(
         Gives the model important context for relevance scoring.
     model:
         Gemini model identifier. Defaults to gemini-2.5-flash-lite.
+    max_files:
+        Hard cap on how many files may be returned, regardless of what the
+        model outputs. Defaults to _MAX_FILES (12) when omitted — typically
+        overridden by a project's `ai.max_files` preference
+        (.contextzip/config.json).
 
     Returns
     -------
@@ -110,7 +116,8 @@ def select_files(
             "Install it with: pip install httpx"
         )
 
-    system_prompt = _build_system_prompt()
+    effective_max_files = max_files if max_files and max_files > 0 else _MAX_FILES
+    system_prompt = _build_system_prompt(effective_max_files)
     user_message = _build_user_message(prompt, file_tree, ecosystem)
 
     url = _API_URL.format(model=model, key=api_key)
@@ -153,7 +160,7 @@ def select_files(
             f"Gemini API returned HTTP {response.status_code}: {response.text[:200]}"
         )
 
-    return _parse_response(response.json(), file_tree)
+    return _parse_response(response.json(), file_tree, effective_max_files)
 
 
 # ---------------------------------------------------------------------------
@@ -161,8 +168,8 @@ def select_files(
 # ---------------------------------------------------------------------------
 
 
-def _build_system_prompt() -> str:
-    return """\
+def _build_system_prompt(max_files: int = _MAX_FILES) -> str:
+    return f"""\
 You are a precise file relevance assistant for software projects.
 
 Your only job: given a developer's task description and a project file tree,
@@ -178,7 +185,7 @@ Rules you must follow:
 - 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.
+- Never return more than {max_files} files. For most tasks 2–5 files is correct.
 - Order by relevance: most directly relevant file first.\
 """
 
@@ -209,13 +216,14 @@ Return only a JSON array of the most relevant file paths for this task.\
 def _parse_response(
     data: dict,
     file_tree: list[tuple[str, int]],
+    max_files: int = _MAX_FILES,
 ) -> 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
+    - Enforces the *max_files* hard cap
     - Warns (via exception) if the model returned mostly invalid paths
     """
     # Navigate the Gemini response structure
@@ -266,7 +274,7 @@ def _parse_response(
         )
 
     # Enforce hard cap
-    return validated[:_MAX_FILES]
+    return validated[:max_files]
 
 
 # ---------------------------------------------------------------------------

contextzip/ai/heuristic.py

@@ -201,6 +201,7 @@ def select_files(
     *,
     prompt: str,
     file_tree: list[tuple[str, int]],
+    max_files: int | None = None,
 ) -> list[str]:
     """
     Score and rank *file_tree* entries by relevance to *prompt*.
@@ -212,18 +213,24 @@ def select_files(
     file_tree:
         Candidate files as (relative_posix_path, size_bytes) tuples.
         Standard exclusions must already have been applied by the caller.
+    max_files:
+        Cap on how many files may be returned. Defaults to MAX_FILES (8)
+        when omitted — typically overridden by a project's `ai.max_files`
+        preference (.contextzip/config.json).
 
     Returns
     -------
     list[str]
         Relative POSIX paths of the top-scoring files, best first.
-        Files scoring zero are excluded. Result is capped at MAX_FILES.
+        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 []
 
+    effective_max_files = max_files if max_files and max_files > 0 else MAX_FILES
+
     scored: list[tuple[float, str]] = []
 
     for rel_path, size_bytes in file_tree:
@@ -234,7 +241,7 @@ def select_files(
     # 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]]
+    return [path for _, path in scored[:effective_max_files]]
 
 
 # ---------------------------------------------------------------------------

contextzip/ai/selector.py

@@ -32,10 +32,15 @@ def ai_select(
     prompt: str,
     ecosystem: str,
     api_key: str,
+    max_files: int | None = None,
 ) -> tuple[list[Path], str, str]:
     """
     Select the minimum relevant files for *prompt* from *resolved.included*.
 
+    *max_files* caps how many files either backend may return — typically a
+    project's `ai.max_files` preference (.contextzip/config.json). None
+    falls back to each backend's own built-in default cap.
+
     Returns
     -------
     (selected_paths, prompt_txt, method)
@@ -54,12 +59,14 @@ def ai_select(
             prompt=prompt,
             file_tree=file_tree,
             ecosystem=ecosystem,
+            max_files=max_files,
         )
     except GeminiRateLimitError:
         # Confirmed HTTP 429 only — fall back to keyword heuristic
         selected_rel = _heuristic.select_files(
             prompt=prompt,
             file_tree=file_tree,
+            max_files=max_files,
         )
         method = USED_HEURISTIC
     # All other GeminiErrors (bad key, network, unexpected status) propagate