Compare commits

...
12 Commits
Author SHA1 Message Date
Julian lechnerandClaude Sonnet 5 ec81d8db9b docs(retrieval): document the background daemon and the max-files limit
CI / linux (push) Successful in 6s
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:36:08 +02:00
Julian lechnerandClaude Sonnet 5 9ad793289d feat(retrieval): add a background daemon to avoid reloading models per query
Each server.py search/rebuild invocation was a fresh process that reloaded
the embedding and reranker models from disk before answering a single
query. Add an optional daemon (server.py daemon) that keeps both models
resident in memory and serves repeated queries, across repository roots,
over a local socket/pipe.

search and rebuild now try the daemon first, spawn one on a cache miss
(guarded by an exclusive lock file to avoid duplicate spawns), and fall
back to the original one-shot in-process worker if the daemon cannot be
reached or started, so behavior is unchanged when the daemon is
unavailable. The daemon exits on its own after 30 minutes of inactivity.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:36:02 +02:00
Julian lechnerandClaude Sonnet 5 208213c1c7 feat(retrieval): pin model revisions and share encoder/reranker instances
Pin the Qwen embedding/reranker models to specific commit revisions (and
sentence-transformers to the exact version already in use) so index
builds are reproducible instead of floating on each model's main branch.

Also move the encoder and reranker instances to module-level singletons
shared by every Index, instead of one pair per instance, so a process
serving multiple repository roots only loads each model once.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:35:46 +02:00
Julian lechnerandClaude Sonnet 5 96a1175e20 feat(retrieval): make the file-scan limit configurable and warn on it
Repositories with more than 2000 indexable files silently lost the
remainder during indexing. Make the cap configurable via
SEMANTIC_RETRIEVAL_MAX_FILES (default unchanged) and print an explicit
stderr warning naming the repository when it is reached, instead of
truncating without any signal.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:35:38 +02:00
Julian lechnerandClaude Sonnet 5 c651b19893 test(build): update reviewer tool assertion for shell-specific tools (Bash)
Mirror the PowerShell test's shell-specific reviewer tools assertion.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:32:14 +02:00
Julian lechnerandClaude Sonnet 5 244f641088 test(build): update reviewer tool assertion for shell-specific tools
The reviewer agent's tools field now resolves to Bash or PowerShell
depending on the package, so assert against the shell-specific value
instead of the old invalid Read,Diff,Search profile.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:32:08 +02:00
Julian lechnerandClaude Sonnet 5 92d98d74cf test(build): cover unknown Claude tool name in Bash build validation
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:32:00 +02:00
Julian lechnerandClaude Sonnet 5 5b768e0aed test(build): cover unknown Claude tool name in PowerShell build validation
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:31:54 +02:00
Julian lechnerandClaude Sonnet 5 99c345f551 feat(build): fail configuration validation on unknown Claude tool names
Check every generated agent's tools frontmatter against the set of real
Claude Code tool names, so an unrecognized name fails validate-config.py
(and therefore the build) instead of silently shipping a broken agent.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:31:48 +02:00
Julian lechnerandClaude Sonnet 5 4da6950062 feat(build): resolve shell-specific agent tool placeholder in Bash build
Mirror the PowerShell build's __SHELL__ resolution so both entry points
stay behaviorally equivalent.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:31:41 +02:00
Julian lechnerandClaude Sonnet 5 0f387b7df2 feat(build): resolve shell-specific agent tool placeholder in PowerShell build
Replace __SHELL__ in each generated Claude agent's tools field with
PowerShell or Bash depending on the package being built, so a role that
needs shell execution only gets the one matching its target shell.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:31:34 +02:00
Julian lechnerandClaude Sonnet 5 4374c69caa fix(claude-agents): map agent capabilities to real Claude Code tool names
Search, Diff, Shell, and Tests are not valid Claude Code tools, so every
generated subagent silently launched with a broken or empty tool set. Map
each role to the tools its own prompt actually needs, and use a __SHELL__
placeholder so build resolves it to the target package's own shell tool
instead of granting both Bash and PowerShell unconditionally.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-09-18 09:31:27 +02:00
12 changed files with 297 additions and 40 deletions

No files matched your search

+5 -5
View File
@@ -1,6 +1,6 @@
role tools
architect Read,Search,WebFetch
researcher Read,Search,WebFetch
reviewer Read,Diff,Search
verifier Read,Shell
implementer Read,Write,Shell,Tests
architect Read,Grep,Glob,WebFetch,WebSearch
researcher Read,Grep,Glob,WebFetch,WebSearch
reviewer Read,Grep,Glob,__SHELL__
verifier Read,Grep,Glob,__SHELL__
implementer Read,Grep,Glob,Edit,Write,__SHELL__
1 role tools
2 architect Read,Search,WebFetch Read,Grep,Glob,WebFetch,WebSearch
3 researcher Read,Search,WebFetch Read,Grep,Glob,WebFetch,WebSearch
4 reviewer Read,Diff,Search Read,Grep,Glob,__SHELL__
5 verifier Read,Shell Read,Grep,Glob,__SHELL__
6 implementer Read,Write,Shell,Tests Read,Grep,Glob,Edit,Write,__SHELL__
+2
View File
@@ -16,6 +16,8 @@ Codex receives Caveman through its native plugin adapter. Humanizer, Impeccable,
Normal installation and update also ask whether to update the selected client CLIs, enable Flashbang, and enable the optional Qwen3-Embedding-0.6B semantic retrieval layer. Retrieval uses a per-user virtual environment and index and is invoked through the installed `server.py search` CLI, disabled by default. Quickupdate reuses the saved shell, client, CLI update, Flashbang, retrieval, and extension choices without opening menus. Disabling retrieval removes its setup-owned runtime and index. Only successful runs save selections; dry runs do not update CLIs, extensions, or retrieval.
The `server.py search`/`rebuild` CLI transparently reuses a background daemon (spawned automatically on first use per `--data-dir`) that keeps the embedding and reranking models loaded in memory, so only the first query after a cold start or an idle period pays the model-loading cost; every call falls back to loading the models directly in-process if the daemon cannot be reached or started. The daemon exits on its own after 30 minutes of inactivity. Indexing skips files larger than 512 KB and, per repository, stops after 2000 indexable files by default; set `SEMANTIC_RETRIEVAL_MAX_FILES` to raise that cap for larger monorepos. When the cap is reached, the CLI prints an explicit warning to stderr rather than silently omitting files.
Plugin installation requires the corresponding client CLI, network access, and any authentication required by the plugin. Context7 may open an OAuth flow. Codex or plugin hooks may require a trust review in `/hooks`. The scripts never install plugins during a build or doctor run and never install them when `--dry-run` or `-DryRun` is used.
Marketplace source conflicts stop installation; the setup does not replace a foreign marketplace registration. Failed extension operations retain the ownership records of earlier successful operations so the next run can reconcile them.
+2
View File
@@ -177,7 +177,9 @@ foreach ($shell in @('powershell','bash')) {
if ($client -eq 'codex') {
Set-Content (Join-Path $agentsOutput "$name.toml") "name = $(Quote-Toml $name)`r`ndescription = $(Quote-Toml $description)`r`ndeveloper_instructions = $(Quote-Toml $instructions)" -Encoding UTF8
} else {
$shellTool = if ($shell -eq 'powershell') { 'PowerShell' } else { 'Bash' }
$tools = $agentTools[$name]
if ($tools) { $tools = $tools.Replace('__SHELL__', $shellTool) }
$frontmatter = "---`r`nname: $name`r`ndescription: $description`r`n"
if ($tools) { $frontmatter += "tools: $tools`r`n" }
Set-Content (Join-Path $agentsOutput "$name.md") ($frontmatter + "---`r`n`r`n$instructions") -Encoding UTF8
+3
View File
@@ -175,6 +175,9 @@ for client in codex claude; do
printf 'name = %s\ndescription = %s\ndeveloper_instructions = %s\n' "$(quote_toml "$name")" "$(quote_toml "$description")" "$(quote_toml "$instructions")" > "$agents_output/$name.toml"
else
tools=${agent_tools[$name]:-}
shell_tool=Bash
[ "$shell" != powershell ] || shell_tool=PowerShell
tools=${tools//__SHELL__/$shell_tool}
if [ -n "$tools" ]; then printf '%s\nname: %s\ndescription: %s\ntools: %s\n%s\n\n%s\n' '---' "$name" "$description" "$tools" '---' "$instructions" > "$agents_output/$name.md"; else printf '%s\nname: %s\ndescription: %s\n%s\n\n%s\n' '---' "$name" "$description" '---' "$instructions" > "$agents_output/$name.md"; fi
fi
done
+8 -1
View File
@@ -38,6 +38,13 @@ Test-BuildFails 'duplicate agent name' {
(Get-Content $architect -Raw) -replace 'name:\s*\S+', 'name: implementer' | Set-Content $architect -NoNewline
}
Test-BuildFails 'unknown Claude tool name in capability manifest' {
param($copy)
$manifest = Join-Path $copy 'adapters/claude/capabilities.tsv'
$lines = Get-Content $manifest | ForEach-Object { if ($_ -match '^architect\t') { "architect`tShell" } else { $_ } }
Set-Content $manifest $lines
}
Test-BuildFails 'duplicate plugin name' {
param($copy)
$manifest = Join-Path $copy 'adapters/plugins.tsv'
@@ -57,5 +64,5 @@ Test-BuildFails 'unknown placeholder beside an allowed installer placeholder' {
}
Remove-Item $work -Recurse -Force -ErrorAction SilentlyContinue
if ($Summary) { Write-Output 'Tests: PASS | build validation' } else { Write-Output 'PASS build validation: invalid JSON, duplicate agent, duplicate plugin, and leftover placeholders are all rejected' }
if ($Summary) { Write-Output 'Tests: PASS | build validation' } else { Write-Output 'PASS build validation: invalid JSON, duplicate agent, unknown Claude tool name, duplicate plugin, and leftover placeholders are all rejected' }
exit 0
+6 -1
View File
@@ -44,6 +44,11 @@ mutate_duplicate_agent() {
}
test_build_fails 'duplicate agent name' mutate_duplicate_agent
mutate_invalid_tool_name() {
sed -i 's/^architect\t.*/architect\tShell/' "$1/adapters/claude/capabilities.tsv"
}
test_build_fails 'unknown Claude tool name in capability manifest' mutate_invalid_tool_name
mutate_duplicate_plugin() {
local manifest="$1/adapters/plugins.tsv"
sed -n '2p' "$manifest" >> "$manifest"
@@ -60,4 +65,4 @@ mutate_mixed_placeholders() {
}
test_build_fails 'unknown placeholder beside an allowed installer placeholder' mutate_mixed_placeholders
if [ "$summary" = true ]; then printf 'Tests: PASS | build validation\n'; else printf 'PASS build validation: invalid JSON, duplicate agent, duplicate plugin, and leftover placeholders are all rejected\n'; fi
if [ "$summary" = true ]; then printf 'Tests: PASS | build validation\n'; else printf 'PASS build validation: invalid JSON, duplicate agent, unknown Claude tool name, duplicate plugin, and leftover placeholders are all rejected\n'; fi
+48 -2
View File
@@ -160,7 +160,7 @@ class RetrievalServerTests(unittest.TestCase):
finally:
index.db.close()
def test_scan_prunes_ignored_directories_and_stops_at_limit(self):
def test_scan_prunes_ignored_directories_and_warns_at_limit(self):
with TemporaryDirectory() as temporary:
root = Path(temporary)
index = MODULE.Index(root, root / 'data')
@@ -172,11 +172,57 @@ class RetrievalServerTests(unittest.TestCase):
try:
self.assertEqual(index.files(), [(root / 'visible.md').resolve()])
filenames = ['visible.md'] * 2001
with patch.object(MODULE.os, 'walk', return_value=iter([(str(root), [], filenames)])):
with patch.object(MODULE.os, 'walk', return_value=iter([(str(root), [], filenames)])), \
patch.object(MODULE.sys, 'stderr') as stderr:
self.assertEqual(len(index.files()), 2000)
self.assertIn('stopped scanning after 2000', stderr.write.call_args_list[0].args[0])
finally:
index.db.close()
def test_max_files_is_configurable_via_environment(self):
with TemporaryDirectory() as temporary:
root = Path(temporary)
index = MODULE.Index(root, root / 'data')
try:
with patch.dict(os.environ, {MODULE.MAX_FILES_ENV: '3'}):
self.assertEqual(MODULE.read_max_files(), 3)
with patch.dict(os.environ, {MODULE.MAX_FILES_ENV: 'not-a-number'}):
with self.assertRaises(ValueError):
MODULE.read_max_files()
with patch.dict(os.environ, {}, clear=True):
self.assertEqual(MODULE.read_max_files(), MODULE.DEFAULT_MAX_FILES)
(root / 'visible.md').write_text('visible', encoding='utf-8')
index.max_files = 3
filenames = ['visible.md'] * 5
with patch.object(MODULE.os, 'walk', return_value=iter([(str(root), [], filenames)])), \
patch.object(MODULE.sys, 'stderr'):
self.assertEqual(len(index.files()), 3)
finally:
index.db.close()
def test_daemon_address_is_deterministic_and_platform_specific(self):
with TemporaryDirectory() as temporary:
data_dir = Path(temporary)
first = MODULE.daemon_address(data_dir)
second = MODULE.daemon_address(data_dir)
self.assertEqual(first, second)
if os.name == 'nt':
self.assertTrue(first.startswith(r'\\.\pipe\ai-config-retrieval-'))
else:
self.assertTrue(first.startswith(str(data_dir)))
def test_shared_encoder_and_reranker_are_cached_module_singletons(self):
with patch.object(MODULE, '_shared_encoder', None), patch.object(MODULE, '_shared_reranker', None), \
patch.object(MODULE, 'preferred_device', lambda: 'cpu'):
sentinel = SimpleNamespace()
with patch.dict('sys.modules', sentence_transformers=SimpleNamespace(
SentenceTransformer=lambda *args, **kwargs: sentinel,
CrossEncoder=lambda *args, **kwargs: sentinel)):
self.assertIs(MODULE.shared_encoder(), sentinel)
self.assertIs(MODULE.shared_encoder(), MODULE.shared_encoder())
self.assertIs(MODULE.shared_reranker(), sentinel)
self.assertIs(MODULE.shared_reranker(), MODULE.shared_reranker())
def test_rerank_uses_cross_encoder_scores(self):
index = MODULE.Index.__new__(MODULE.Index)
+2 -1
View File
@@ -31,7 +31,8 @@ foreach ($shell in @('powershell','bash')) {
}
if ($client -eq 'claude') { $settings = $content | ConvertFrom-Json }
if (@(Get-ChildItem "$package/agents" -File).Count -ne $sourceAgentCount) { throw "Missing agents in $package" }
if ($client -eq 'claude' -and (Get-Content (Join-Path $package 'agents/reviewer.md') -Raw) -notmatch '(?m)^tools: Read,Diff,Search\r?$') { throw "Reviewer capability profile is missing in $package" }
$expectedReviewerShellTool = if ($shell -eq 'powershell') { 'PowerShell' } else { 'Bash' }
if ($client -eq 'claude' -and (Get-Content (Join-Path $package 'agents/reviewer.md') -Raw) -notmatch "(?m)^tools: Read,Grep,Glob,$expectedReviewerShellTool\r?`$") { throw "Reviewer capability profile is missing in $package" }
$expectedSkillCount = if ($client -eq 'claude') { $sourceSkillCount + $ruleSkillCount } else { $sourceSkillCount }
if (@(Get-ChildItem "$package/skills" -Filter SKILL.md -Recurse).Count -ne $expectedSkillCount) { throw "Missing skills in $package" }
$expected = if ($shell -eq 'powershell') { '__POWERSHELL_COMMAND__ .*flashbang.ps1' } else { 'bash .*flashbang.sh' }
+3 -1
View File
@@ -41,7 +41,9 @@ for shell in powershell bash; do
done < "$root/adapters/rule-skills.tsv"
fi
test "$(find "$package/agents" -type f | wc -l)" -eq "$source_agent_count"
if [ "$client" = claude ]; then grep -q '^tools: Read,Diff,Search$' "$package/agents/reviewer.md"; fi
expected_reviewer_shell_tool=Bash
[ "$shell" != powershell ] || expected_reviewer_shell_tool=PowerShell
if [ "$client" = claude ]; then grep -q "^tools: Read,Grep,Glob,$expected_reviewer_shell_tool\$" "$package/agents/reviewer.md"; fi
expected_skill_count=$source_skill_count
[ "$client" != claude ] || expected_skill_count=$((source_skill_count + rule_skill_count))
test "$(find "$package/skills" -name SKILL.md | wc -l)" -eq "$expected_skill_count"
+14
View File
@@ -15,6 +15,7 @@ CLAUDE_SANDBOX_KEYS = {"enabled", "allowUnsandboxedCommands"}
CLAUDE_STATUS_LINE_KEYS = {"type", "command"}
CLAUDE_HOOK_KEYS = {"type", "command", "timeout"}
CLAUDE_PERMISSION_MODES = {"acceptEdits", "auto", "bypassPermissions", "default", "dontAsk", "plan"}
CLAUDE_VALID_TOOLS = {"Read", "Grep", "Glob", "Edit", "Write", "Bash", "PowerShell", "WebFetch", "WebSearch"}
def fail(message: str) -> None:
@@ -72,6 +73,18 @@ def validate_rule_skills(package: Path) -> None:
fail(f"{path}: skill description must be a quoted, nonempty string")
def validate_agent_tools(package: Path) -> None:
for path in (package / "agents").glob("*.md"):
lines = path.read_text(encoding="utf-8-sig").splitlines()
tools_line = next((line for line in lines if line.startswith("tools: ")), None)
if tools_line is None:
continue
tools = [item.strip() for item in tools_line.removeprefix("tools: ").split(",")]
unknown = [tool for tool in tools if tool not in CLAUDE_VALID_TOOLS]
if unknown:
fail(f"{path}: unknown Claude tool name(s): {', '.join(unknown)}")
def main() -> int:
if len(sys.argv) != 3:
print("usage: validate-config.py <codex-toml> <claude-json>", file=sys.stderr)
@@ -79,6 +92,7 @@ def main() -> int:
tomllib.loads(Path(sys.argv[1]).read_text(encoding="utf-8-sig"))
validate_claude(Path(sys.argv[2]))
validate_rule_skills(Path(sys.argv[2]).parent)
validate_agent_tools(Path(sys.argv[2]).parent)
return 0
+1 -1
View File
@@ -1 +1 @@
sentence-transformers>=5.4,<6
sentence-transformers==5.7.0
+203 -28
View File
@@ -9,16 +9,21 @@ import hashlib
import heapq
import json
import multiprocessing
from multiprocessing.connection import Client, Listener
import os
from pathlib import Path
import re
import sqlite3
import subprocess
import sys
import threading
import time
MODEL = "Qwen/Qwen3-Embedding-0.6B"
MODEL_REVISION = "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3"
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
RERANKER_MODEL_REVISION = "e61197ed45024b0ed8a2d74b80b4d909f1255473"
EMBEDDING_DIMENSIONS = 1024
CHUNK_TOKENS = 400
OVERLAP_TOKENS = 64
@@ -31,6 +36,15 @@ RRF_K = 60
INDEX_VERSION = 2
RETRIEVAL_INSTRUCTION = "Given a codebase question, retrieve relevant code and documentation that answer the question."
TEXT_EXTENSIONS = {".c", ".cpp", ".cs", ".go", ".java", ".js", ".json", ".md", ".py", ".ps1", ".rs", ".sh", ".toml", ".ts", ".tsx", ".txt", ".yaml", ".yml"}
DEFAULT_MAX_FILES = 2000
MAX_FILES_ENV = "SEMANTIC_RETRIEVAL_MAX_FILES"
# A daemon keeps the embedding/reranker models resident in memory across CLI
# invocations, since each `server.py search` call is otherwise a fresh process
# that would reload both models from disk before it can answer a single query.
DAEMON_IDLE_SECONDS = 1800
DAEMON_READY_TIMEOUT_SECONDS = 120
DAEMON_AUTHKEY = b"ai-config-semantic-retrieval"
def runtime_python(script=__file__, executable=sys.executable, system=os.name):
@@ -148,6 +162,45 @@ def preferred_device():
return "cuda" if torch.cuda.is_available() else "cpu"
def read_max_files():
value = os.environ.get(MAX_FILES_ENV)
if value is None:
return DEFAULT_MAX_FILES
try:
parsed = int(value)
except ValueError:
parsed = 0
if parsed <= 0:
raise ValueError(f"{MAX_FILES_ENV} must be a positive integer, got {value!r}")
return parsed
_shared_model_lock = threading.Lock()
_shared_encoder = None
_shared_reranker = None
def shared_encoder():
global _shared_encoder
with _shared_model_lock:
if _shared_encoder is None:
from sentence_transformers import SentenceTransformer
_shared_encoder = SentenceTransformer(MODEL, revision=MODEL_REVISION, truncate_dim=EMBEDDING_DIMENSIONS,
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
device=preferred_device())
return _shared_encoder
def shared_reranker():
global _shared_reranker
with _shared_model_lock:
if _shared_reranker is None:
from sentence_transformers import CrossEncoder
_shared_reranker = CrossEncoder(RERANKER_MODEL, revision=RERANKER_MODEL_REVISION, prompts={"query": RETRIEVAL_INSTRUCTION},
default_prompt_name="query", device=preferred_device())
return _shared_reranker
def reciprocal_rank_fusion(*rankings):
scores = {}
for ranking in rankings:
@@ -156,6 +209,122 @@ def reciprocal_rank_fusion(*rankings):
return sorted(scores, key=lambda identifier: (-scores[identifier], identifier))
def daemon_lock_path(data_dir: Path) -> Path:
return data_dir / "daemon.starting"
def daemon_address(data_dir: Path) -> str:
digest = hashlib.sha256(os.path.normcase(str(data_dir.resolve())).encode()).hexdigest()[:16]
if os.name == "nt":
return r"\\.\pipe\ai-config-retrieval-" + digest
return str(data_dir / f"daemon-{digest}.sock")
def request_daemon(address: str, payload: dict):
with Client(address, authkey=DAEMON_AUTHKEY) as connection:
connection.send(payload)
if not connection.poll(CLI_TIMEOUT_SECONDS):
raise TimeoutError(f"Semantic retrieval daemon exceeded the {CLI_TIMEOUT_SECONDS} second timeout.")
success, result = connection.recv()
if not success:
raise RuntimeError(str(result))
return result
def spawn_daemon(data_dir: Path, model_cache: Path) -> bool:
lock_path = daemon_lock_path(data_dir)
try:
os.close(os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY))
except FileExistsError:
try:
if time.time() - lock_path.stat().st_mtime > DAEMON_READY_TIMEOUT_SECONDS:
lock_path.unlink(missing_ok=True)
except OSError:
pass
return True
try:
subprocess.Popen(
[sys.executable, str(Path(__file__).resolve()), "daemon", "--data-dir", str(data_dir), "--model-cache", str(model_cache)],
stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
start_new_session=(os.name != "nt"),
creationflags=(subprocess.DETACHED_PROCESS | subprocess.CREATE_NEW_PROCESS_GROUP) if os.name == "nt" else 0)
return True
except OSError:
lock_path.unlink(missing_ok=True)
return False
def try_daemon(root: Path, data_dir: Path, model_cache: Path, action: str, query: str | None, top_k: int):
address = daemon_address(data_dir)
payload = {"action": action, "root": str(root), "query": query, "top_k": top_k}
try:
return request_daemon(address, payload)
except (OSError, EOFError):
pass
if not spawn_daemon(data_dir, model_cache):
return None
deadline = time.monotonic() + DAEMON_READY_TIMEOUT_SECONDS
while time.monotonic() < deadline:
try:
return request_daemon(address, payload)
except (OSError, EOFError):
time.sleep(0.5)
return None
def run_daemon(data_dir: Path, model_cache: Path):
address = daemon_address(data_dir)
if os.name != "nt":
Path(address).unlink(missing_ok=True)
model_cache.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("HF_HOME", str(model_cache.resolve()))
indexes: dict[str, Index] = {}
activity_lock = threading.Lock()
last_activity = time.monotonic()
def watchdog():
while True:
time.sleep(5)
with activity_lock:
idle_for = time.monotonic() - last_activity
if idle_for > DAEMON_IDLE_SECONDS:
os._exit(0)
threading.Thread(target=watchdog, daemon=True).start()
listener = Listener(address, family="AF_PIPE" if os.name == "nt" else "AF_UNIX", authkey=DAEMON_AUTHKEY)
daemon_lock_path(data_dir).unlink(missing_ok=True)
try:
while True:
connection = listener.accept()
try:
response = (True, None)
try:
request = connection.recv()
index = indexes.get(request["root"])
if index is None:
if len(indexes) >= 8:
indexes.pop(next(iter(indexes))).db.close()
index = Index(Path(request["root"]), data_dir)
indexes[request["root"]] = index
if request["action"] == "search":
response = (True, index.search(request["query"], request["top_k"]))
else:
response = (True, {"chunks": index.rebuild()})
except EOFError:
continue
except Exception as error:
response = (False, str(error))
connection.send(response)
except (OSError, EOFError):
pass
finally:
connection.close()
with activity_lock:
last_activity = time.monotonic()
finally:
listener.close()
class Index:
def __init__(self, root: Path, data: Path):
self.root = root.resolve()
@@ -172,8 +341,7 @@ class Index:
self.db.execute("create table if not exists query_cache (cache_key text primary key, result text not null)")
self.db.commit()
self.lock = threading.RLock()
self.model = None
self.reranker_model = None
self.max_files = read_max_files()
def files(self):
ignored = {".git", ".venv", "generated", "node_modules", ".my-ai-configuration"}
@@ -188,24 +356,18 @@ class Index:
if path.is_symlink() or path.suffix.lower() not in TEXT_EXTENSIONS or not path.is_file() or path.stat().st_size > 524288:
continue
result.append(path)
if len(result) == 2000:
if len(result) == self.max_files:
print(f"Semantic retrieval: stopped scanning after {self.max_files} indexable files; "
f"some files under {self.root} were not indexed. Set {MAX_FILES_ENV} to raise this limit.",
file=sys.stderr)
return result
return result
def encoder(self):
if self.model is None:
from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(MODEL, truncate_dim=EMBEDDING_DIMENSIONS,
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "},
device=preferred_device())
return self.model
return shared_encoder()
def reranker(self):
if self.reranker_model is None:
from sentence_transformers import CrossEncoder
self.reranker_model = CrossEncoder(RERANKER_MODEL, prompts={"query": RETRIEVAL_INSTRUCTION},
default_prompt_name="query", device=preferred_device())
return self.reranker_model
return shared_reranker()
def rerank(self, query: str, rows: list[dict], limit: int):
candidates = rows[:RERANK_CANDIDATES]
@@ -308,21 +470,9 @@ class Index:
return result
def main():
parser = argparse.ArgumentParser()
parser.add_argument("action", choices=("search", "rebuild"))
parser.add_argument("--query")
parser.add_argument("--top-k", type=int, default=FINAL_RESULTS)
parser.add_argument("--root", default=os.getcwd(), type=Path)
parser.add_argument("--data-dir", required=True, type=Path)
parser.add_argument("--model-cache", required=True, type=Path)
args = parser.parse_args()
if args.action == "search" and not args.query:
parser.error("search requires --query")
def run_direct(action: str, root: Path, data_dir: Path, model_cache: Path, query: str | None, top_k: int):
output = multiprocessing.Queue(1)
worker = multiprocessing.Process(target=run_action,
args=(args.action, args.root.resolve(), args.data_dir.resolve(),
args.model_cache.resolve(), args.query, args.top_k, output))
worker = multiprocessing.Process(target=run_action, args=(action, root, data_dir, model_cache, query, top_k, output))
worker.start()
worker.join(CLI_TIMEOUT_SECONDS)
if worker.is_alive():
@@ -334,6 +484,31 @@ def main():
success, result = output.get()
if not success:
raise RuntimeError(str(result))
return result
def main():
parser = argparse.ArgumentParser()
parser.add_argument("action", choices=("search", "rebuild", "daemon"))
parser.add_argument("--query")
parser.add_argument("--top-k", type=int, default=FINAL_RESULTS)
parser.add_argument("--root", default=os.getcwd(), type=Path)
parser.add_argument("--data-dir", required=True, type=Path)
parser.add_argument("--model-cache", required=True, type=Path)
args = parser.parse_args()
data_dir = args.data_dir.resolve()
model_cache = args.model_cache.resolve()
if args.action == "daemon":
data_dir.mkdir(parents=True, exist_ok=True)
run_daemon(data_dir, model_cache)
return
if args.action == "search" and not args.query:
parser.error("search requires --query")
root = args.root.resolve()
data_dir.mkdir(parents=True, exist_ok=True)
result = try_daemon(root, data_dir, model_cache, args.action, args.query, args.top_k)
if result is None:
result = run_direct(args.action, root, data_dir, model_cache, args.query, args.top_k)
print(json.dumps(result, ensure_ascii=False))