Compare commits
12
Commits
4408f23c68
...
ec81d8db9b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ec81d8db9b | ||
|
|
9ad793289d | ||
|
|
208213c1c7 | ||
|
|
96a1175e20 | ||
|
|
c651b19893 | ||
|
|
244f641088 | ||
|
|
92d98d74cf | ||
|
|
5b768e0aed | ||
|
|
99c345f551 | ||
|
|
4da6950062 | ||
|
|
0f387b7df2 | ||
|
|
4374c69caa |
No files matched your search
@@ -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__
|
||||
|
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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' }
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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 @@
|
||||
sentence-transformers>=5.4,<6
|
||||
sentence-transformers==5.7.0
|
||||
+203
-28
@@ -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))
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user