Compare commits

..
13 Commits
3 changed files with 427 additions and 47 deletions

No files matched your search

+147 -14
View File
@@ -6,30 +6,163 @@ on:
pull_request:
schedule:
- cron: '23 4 * * 1'
workflow_dispatch:
permissions:
contents: read
concurrency:
group: ci-${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
env:
PYTHONUNBUFFERED: '1'
PIP_DISABLE_PIP_VERSION_CHECK: '1'
jobs:
ubuntu:
capabilities:
name: Capabilities and shell quality
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
- name: Install ShellCheck
run: sudo apt-get update && sudo apt-get install --yes shellcheck
- name: Shell syntax and lint
shell: bash
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Install shell tools
run: |
git ls-files -z '*.sh' | xargs -0 -r -n1 bash -n
shellcheck --severity=warning $(git ls-files '*.sh')
- name: Python tests
run: for file in scripts/test-*.py; do python3 "$file"; done
- name: Build
sudo apt-get update
sudo apt-get install --yes --no-install-recommends shellcheck jq
- name: Check capabilities
run: bash scripts/check-capabilities.sh
- name: Validate Bash syntax
run: git ls-files -z '*.sh' | xargs -0 -r -n1 bash -n
- name: Lint Bash
run: shellcheck --severity=warning $(git ls-files '*.sh')
python:
name: Python tests and compilation
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- uses: actions/setup-python@82c7e631bb3cdc910f68e0081d67478d79c6982d # v5.6.0
with:
python-version: '3.12'
- name: Compile Python sources
run: python -m compileall -q scripts shared
- name: Run Python tests
run: |
set -euo pipefail
for file in scripts/test-*.py; do
echo "==> $file"
python "$file"
done
- name: Validate generated configuration
run: python scripts/validate-config.py generated/codex-powershell/config.toml generated/claude-powershell/settings.json
powershell-on-ubuntu:
name: PowerShell tests on Ubuntu
runs-on: ubuntu-latest
timeout-minutes: 20
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Check PowerShell runtime
run: pwsh -NoLogo -NoProfile -NonInteractive -Command '$PSVersionTable.PSVersion.ToString()'
- name: Run PowerShell tests
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
./scripts/check-capabilities.ps1
Get-ChildItem scripts/test-*.ps1 | ForEach-Object {
Write-Host "==> $($_.Name)"
& $_.FullName -Summary
}
build:
name: Build and inspect generated packages
runs-on: ubuntu-latest
timeout-minutes: 15
needs: [capabilities, python]
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Build packages
env:
REQUIRE_CODEX_SCHEMA: 'false'
run: bash scripts/build.sh --summary
- name: Bash tests
- name: Run build validation
run: bash scripts/test-build-validation.sh --summary
- name: Check generated files are reproducible
run: |
bash scripts/test-build-validation.sh --summary
before=$(find generated -type f -print0 | sort -z | xargs -0 sha256sum)
bash scripts/build.sh --summary
after=$(find generated -type f -print0 | sort -z | xargs -0 sha256sum)
test "$before" = "$after"
shell-integration:
name: Bash integration and installation guards
runs-on: ubuntu-latest
timeout-minutes: 25
needs: [build]
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Run integration tests
run: |
set -euo pipefail
bash scripts/test-install-guards.sh --summary
bash scripts/test-plugin-selection.sh --summary
bash scripts/test-quick-update.sh --summary
bash scripts/test-shell-packages.sh --summary
bash scripts/test-managed-install.sh --summary
bash scripts/test-shell-packages.sh --summary
- name: Run doctor and prompt budget checks
run: |
bash scripts/doctor.sh --summary
bash scripts/prompt-budget.sh --summary
- name: Verify dry-run installation safety
run: |
bash scripts/install.sh --summary --dry-run --client both --shell bash
bash scripts/update.sh --summary --dry-run --client both --shell bash
retrieval:
name: Qwen retrieval integration
runs-on: ubuntu-latest
timeout-minutes: 30
needs: [python]
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- uses: actions/setup-python@82c7e631bb3cdc910f68e0081d67478d79c6982d # v5.6.0
with:
python-version: '3.12'
cache: pip
cache-dependency-path: shared/retrieval/requirements.txt
- name: Install retrieval dependencies
run: python -m pip install -r shared/retrieval/requirements.txt
- uses: actions/cache@6849a6489940f00c2f30c0fb92c6274307ccb58a # v4.2.3
with:
path: .cache/huggingface
key: qwen-retrieval-${{ runner.os }}-${{ hashFiles('shared/retrieval/requirements.txt') }}
- name: Run retrieval tests with real models
env:
QWEN_TEST_MODEL_CACHE: ${{ github.workspace }}/.cache/huggingface
run: python scripts/test-retrieval-server.py
security:
name: CI configuration security checks
runs-on: ubuntu-latest
timeout-minutes: 10
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
- name: Reject unsafe workflow settings
run: |
set -euo pipefail
! grep -RInE 'runs-on:[[:space:]]*windows|pull_request_target|permissions:[[:space:]]*write-all|actions/checkout@v|actions/setup-python@v|curl[[:space:]].*\|[[:space:]]*(bash|sh)|wget[[:space:]].*\|[[:space:]]*(bash|sh)' .github/workflows
test "$(grep -Roh 'runs-on: ubuntu-latest' .github/workflows | wc -l)" -ge 1
test "$(grep -Roh 'uses: actions/[^@]*@[0-9a-f]\{40\}' .github/workflows | wc -l)" -ge 1
- name: Check for tracked secrets
run: |
set -euo pipefail
if git ls-files | grep -E '(^|/)(\.env($|\.)|.*\.(pem|key))$'; then
echo 'Potential secret file or private key tracked' >&2
exit 1
fi
if git ls-files -z | xargs -0 -r grep -IlE 'BEGIN (RSA|OPENSSH|EC) PRIVATE KEY'; then
echo 'Potential private key material tracked' >&2
exit 1
fi
+132 -11
View File
@@ -3,6 +3,7 @@ from array import array
from concurrent.futures import ThreadPoolExecutor
import os
from pathlib import Path
import re
from tempfile import TemporaryDirectory
from types import SimpleNamespace
import time
@@ -16,6 +17,15 @@ MODULE = importlib.util.module_from_spec(SPEC)
SPEC.loader.exec_module(MODULE)
def tokenize(text, **kwargs):
offsets = [match.span() for match in re.finditer(r'\S+', text)]
return {'offset_mapping': offsets, 'input_ids': list(range(len(offsets)))}
def embedding():
return array('f', [1.0] + [0.0] * (MODULE.EMBEDDING_DIMENSIONS - 1))
class RetrievalServerTests(unittest.TestCase):
@unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.')
def test_real_models(self):
@@ -26,6 +36,13 @@ class RetrievalServerTests(unittest.TestCase):
(root / 'weather.md').write_text('The weather service fetches temperature and rainfall forecasts for a city.', encoding='utf-8')
index = MODULE.Index(root, root / 'data')
try:
tokenizer = index.encoder().tokenizer
for source in ['alpha beta gamma ' * 450, 'Grüße 世界 🙂 ' * 300]:
windows = list(MODULE.chunks(source, 'sample.txt', tokenizer))
self.assertGreater(len(windows), 1)
self.assertTrue(all(len(tokenizer(value, add_special_tokens=False)['input_ids']) <= 400
for _, value in windows))
self.assertTrue(all('\ufffd' not in value for _, value in windows))
for query, expected in [('Where are password reset tokens created?', 'authentication.md'),
('Which service provides rainfall forecasts?', 'weather.md')]:
started = time.monotonic()
@@ -52,10 +69,10 @@ class RetrievalServerTests(unittest.TestCase):
def encode(values, **kwargs):
self.assertEqual(kwargs.get('prompt_name'), 'query' if values == ['query'] else None)
encoded.extend(values)
return [array('f', [1.0]) for _ in values]
return [embedding() for _ in values]
for index in indexes:
index.encoder = lambda: SimpleNamespace(encode=encode)
index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
index.rerank = lambda query, rows, limit: rows[:limit]
numpy = SimpleNamespace(dot=lambda a, b: sum(x * y for x, y in zip(a, b)),
frombuffer=lambda blob, **kwargs: array('f', blob), float32=None)
@@ -88,21 +105,23 @@ class RetrievalServerTests(unittest.TestCase):
root = Path(temporary)
(root / 'a.md').write_text('before', encoding='utf-8')
index = MODULE.Index(root, root / 'data')
index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [array('f', [1.0]) for _ in values])
index.encoder = lambda: SimpleNamespace(encode=lambda values, **kwargs: [embedding() for _ in values], tokenizer=tokenize)
try:
index.rebuild()
(root / 'a.md').write_text('after', encoding='utf-8')
(root / 'b.md').write_text('fail', encoding='utf-8')
def encode(values, **kwargs):
if values == ['fail']:
if any(value.endswith('\n\nfail') for value in values):
raise RuntimeError('encoder failed')
return [array('f', [1.0]) for _ in values]
return [embedding() for _ in values]
index.encoder = lambda: SimpleNamespace(encode=encode)
index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
with self.assertRaisesRegex(RuntimeError, 'encoder failed'):
index.rebuild()
self.assertEqual(index.db.execute('select path, text from chunks').fetchall(), [('a.md', 'before')])
self.assertEqual(index.db.execute('select document from lexical').fetchall(),
[('Path: a.md\nSymbol: <module>\n\nbefore',)])
finally:
index.db.close()
@@ -127,22 +146,124 @@ class RetrievalServerTests(unittest.TestCase):
index = MODULE.Index.__new__(MODULE.Index)
class Reranker:
def predict(self, pairs):
def predict(self, pairs, **kwargs):
self.pairs = pairs
return [0.1, 0.8, 0.5]
reranker = Reranker()
index.reranker = lambda: reranker
rows = [
{'path': 'a.md', 'score': 0.9, 'text': 'first'},
{'path': 'b.md', 'score': 0.8, 'text': 'second'},
{'path': 'c.md', 'score': 0.7, 'text': 'third'},
{'path': 'a.md', 'symbol': 'A', 'score': 0.9, 'text': 'first'},
{'path': 'b.md', 'symbol': 'B', 'score': 0.8, 'text': 'second'},
{'path': 'c.md', 'symbol': 'C', 'score': 0.7, 'text': 'third'},
]
result = index.rerank('query', rows, 2)
self.assertEqual([item['path'] for item in result], ['b.md', 'c.md'])
self.assertEqual(reranker.pairs, [('query', 'first'), ('query', 'second'), ('query', 'third')])
self.assertEqual(reranker.pairs, [('query', 'Path: a.md\nSymbol: A\n\nfirst'),
('query', 'Path: b.md\nSymbol: B\n\nsecond'),
('query', 'Path: c.md\nSymbol: C\n\nthird')])
def test_token_windows_have_exact_overlap_and_no_redundant_tail(self):
for size, lengths in [(0, []), (400, [400]), (401, [400, 65]), (736, [400, 400]), (800, [400, 400, 128])]:
words = [f'token{number}' for number in range(size)]
result = list(MODULE.chunks(' '.join(words), 'plain.txt', tokenize))
windows = [content.split() for _, content in result]
self.assertEqual([len(window) for window in windows], lengths)
for first, second in zip(windows, windows[1:]):
self.assertEqual(first[-64:], second[:64])
if windows:
reconstructed = windows[0] + [word for window in windows[1:] for word in window[64:]]
self.assertEqual(reconstructed, words)
def test_python_symbols_include_decorators_and_restore_parent(self):
source = 'class Service:\n @cached\n def load(self):\n return 1\n\nvalue = 2\n\nasync def refresh():\n return value\n'
result = list(MODULE.sections(source, 'service.py'))
self.assertEqual([symbol for symbol, _ in result], ['Service', 'Service.load', '<module>', 'refresh'])
self.assertIn('@cached', result[1][1])
self.assertEqual(''.join(content for _, content in result), source)
broken = 'def incomplete(\n' + 'word ' * 450
self.assertEqual(len(list(MODULE.chunks(broken, 'broken.py', tokenize))), 2)
def test_markdown_sections_ignore_fenced_headings(self):
source = '# Setup\nIntro\n## Install\n```sh\n# not a heading\n```\nSteps\n## Run\nStart\n'
result = list(MODULE.chunks(source, 'guide.md', tokenize))
self.assertEqual([symbol for symbol, _ in result], ['Setup', 'Setup / Install', 'Setup / Run'])
self.assertIn('# not a heading', result[1][1])
def test_common_declarations_and_config_sections(self):
for path, source, symbols in [
('service.ts', 'export function load() {\n return 1;\n}\nexport const save = () => 2;\n', ['load', 'save']),
('service.cs', 'public class Service\n{\n public void Load()\n {\n }\n}\n', ['Service', 'Load']),
('setup.ps1', 'function Get-Result {\n return 1\n}\n', ['Get-Result']),
('config.toml', '[server]\nport = 1\n[client]\nport = 2\n', ['server', 'client']),
]:
self.assertEqual([symbol for symbol, _ in MODULE.sections(source, path)], symbols)
def test_rrf_rewards_agreement_and_breaks_ties_deterministically(self):
self.assertEqual(MODULE.reciprocal_rank_fusion([1, 2, 3], [4, 2, 3]), [2, 3, 1, 4])
self.assertEqual(MODULE.reciprocal_rank_fusion([], [4, 2]), [4, 2])
def test_hybrid_candidates_and_document_metadata(self):
with TemporaryDirectory() as temporary:
root = Path(temporary)
for number in range(110):
(root / f'{number:03}.md').write_text(f'# Symbol{number}\ncontent {number}', encoding='utf-8', newline='\n')
index = MODULE.Index(root, root / 'data')
documents = []
def encode(values, **kwargs):
if kwargs.get('prompt_name') != 'query':
documents.extend(values)
return [embedding() for _ in values]
pairs = []
index.encoder = lambda: SimpleNamespace(encode=encode, tokenizer=tokenize)
index.reranker = lambda: SimpleNamespace(predict=lambda values, **kwargs:
pairs.extend(values) or [float('109.md' in document) for _, document in values])
numpy = SimpleNamespace(dot=lambda a, b: 1.0,
frombuffer=lambda blob, **kwargs: array('f', blob), float32=None)
original_fusion = MODULE.reciprocal_rank_fusion
rankings = []
def fuse(*values):
rankings.extend(values)
return original_fusion(*values)
try:
with patch.dict('sys.modules', numpy=numpy), patch.object(MODULE, 'reciprocal_rank_fusion', fuse):
results = index.search('109.md Symbol109 content', 20)
self.assertEqual([len(ranking) for ranking in rankings], [50, 50])
self.assertEqual(len(pairs), 50)
self.assertEqual(len(results), 5)
self.assertEqual(results[0]['path'], '109.md')
self.assertEqual(results[0]['symbol'], 'Symbol109')
self.assertIn('Path: 109.md\nSymbol: Symbol109\n\n# Symbol109\ncontent 109', documents)
self.assertTrue(all(document in documents for _, document in pairs))
with patch.dict('sys.modules', numpy=numpy):
self.assertEqual(len(index.search('" OR * () : -')), 5)
self.assertEqual(len(index.search('***')), 5)
(root / '109.md').unlink()
index.search('Symbol109')
self.assertEqual(index.db.execute('select count(*) from lexical where lexical match ?', ('"Symbol109"',)).fetchone()[0], 0)
finally:
index.db.close()
def test_old_cache_is_preserved_but_not_reused(self):
with TemporaryDirectory() as temporary:
root = Path(temporary)
data = root / 'data'
data.mkdir()
identity = os.path.normcase(str(root.resolve())) + '\n' + MODULE.MODEL
old = data / (MODULE.hashlib.sha256(identity.encode()).hexdigest() + '.sqlite3')
old.write_bytes(b'old cache')
index = MODULE.Index(root, data)
try:
self.assertNotEqual(Path(index.db.execute('pragma database_list').fetchone()[2]), old)
self.assertEqual(old.read_bytes(), b'old cache')
finally:
index.db.close()
if __name__ == '__main__':
+148 -22
View File
@@ -4,26 +4,126 @@
from __future__ import annotations
import argparse
import ast
import hashlib
import heapq
import json
import os
from pathlib import Path
import re
import sqlite3
import threading
MODEL = "Qwen/Qwen3-Embedding-0.6B"
RERANKER_MODEL = "Qwen/Qwen3-Reranker-0.6B"
EMBEDDING_DIMENSIONS = 1024
CHUNK_TOKENS = 400
OVERLAP_TOKENS = 64
DENSE_CANDIDATES = 50
LEXICAL_CANDIDATES = 50
RERANK_CANDIDATES = 50
FINAL_RESULTS = 5
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"}
def chunks(text: str):
for start in range(0, len(text), 1050):
value = text[start:start + 1200].strip()
def sections(text: str, path: str):
lines = text.splitlines(keepends=True)
boundaries = {0: "<module>", len(lines): "<module>"}
suffix = Path(path).suffix.lower()
if suffix == ".py":
try:
tree = ast.parse(text)
except (SyntaxError, ValueError):
tree = None
def visit(node, parent):
symbol = parent
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
symbol = f"{parent}.{node.name}" if parent != "<module>" else node.name
start = min([node.lineno, *(item.lineno for item in node.decorator_list)]) - 1
boundaries[start] = symbol
boundaries.setdefault(node.end_lineno, parent)
for child in ast.iter_child_nodes(node):
visit(child, symbol)
if tree is not None:
visit(tree, "<module>")
else:
fence = None
headings = []
for number, line in enumerate(lines):
if suffix == ".md":
marker = re.match(r"^\s{0,3}(`{3,}|~{3,})", line)
if marker:
value = marker[1]
if fence is None:
fence = value
elif value[0] == fence[0] and len(value) >= len(fence):
fence = None
continue
heading = re.match(r"^\s{0,3}(#{1,6})\s+(.+?)\s*#*\s*$", line)
if heading and fence is None:
level = len(heading[1])
headings = [(depth, title) for depth, title in headings if depth < level]
headings.append((level, heading[2]))
boundaries[number] = " / ".join(title for _, title in headings)
elif suffix == ".toml":
heading = re.match(r"^\s*\[\[?([^\]]+)\]\]?\s*(?:#.*)?$", line)
if heading:
boundaries[number] = heading[1]
elif suffix in {".js", ".ts", ".tsx", ".cs", ".java", ".go", ".rs", ".c", ".cpp", ".ps1", ".sh"}:
symbol = re.match(
r"^\s*(?:(?:export|default|public|private|protected|internal|static|abstract|async|sealed|partial|pub)\s+)*"
r"(?:(?:class|interface|struct|enum|record|function|fn|def)\s+([\w$-]+)"
r"|func\s+(?:\([^)]*\)\s*)?([\w]+)"
r"|(?:const|let|var)\s+([\w$]+)\s*=.*=>"
r"|([\w$-]+)\s*\([^;]*\)\s*\{"
r"|[\w<>\[\],?]+\s+([\w]+)\s*\([^;]*\)\s*(?:\{|=>|$))", line)
if symbol:
name = next(value for value in symbol.groups() if value)
if name not in {"if", "for", "foreach", "while", "switch", "catch", "using", "lock"}:
boundaries[number] = name
positions = sorted(boundaries)
for start, end in zip(positions, positions[1:]):
content = "".join(lines[start:end])
if content.strip():
yield boundaries[start], content
def chunks(text: str, path: str, tokenizer):
for symbol, content in sections(text, path):
offsets = tokenizer(content, add_special_tokens=False, return_offsets_mapping=True)["offset_mapping"]
start = 0
while start < len(offsets):
end = min(start + CHUNK_TOKENS, len(offsets))
first = 0 if start == 0 else offsets[start][0]
while end > start:
last = len(content) if end == len(offsets) else offsets[end][0]
value = content[first:last].strip()
if len(tokenizer(value, add_special_tokens=False)["input_ids"]) <= CHUNK_TOKENS:
break
end -= 1
if value:
yield value
yield symbol, value
if end == len(offsets):
break
start = max(start + 1, end - OVERLAP_TOKENS)
def document_text(path: str, symbol: str, content: str):
return f"Path: {path}\nSymbol: {symbol}\n\n{content}"
def reciprocal_rank_fusion(*rankings):
scores = {}
for ranking in rankings:
for rank, identifier in enumerate(ranking, 1):
scores[identifier] = scores.get(identifier, 0.0) + 1.0 / (RRF_K + rank)
return sorted(scores, key=lambda identifier: (-scores[identifier], identifier))
class Index:
@@ -31,11 +131,13 @@ class Index:
self.root = root.resolve()
self.data = data.resolve()
self.data.mkdir(parents=True, exist_ok=True)
identity = os.path.normcase(str(self.root)) + "\n" + MODEL
identity = json.dumps([os.path.normcase(str(self.root)), MODEL, EMBEDDING_DIMENSIONS,
INDEX_VERSION, CHUNK_TOKENS, OVERLAP_TOKENS])
database = hashlib.sha256(identity.encode()).hexdigest() + ".sqlite3"
self.db = sqlite3.connect(self.data / database, check_same_thread=False, timeout=30)
self.db.execute("create table if not exists chunks (path text, text text, vector blob)")
self.db.execute("create table if not exists chunks (path text, symbol text, text text, vector blob)")
self.db.execute("create index if not exists chunks_path on chunks (path)")
self.db.execute("create virtual table if not exists lexical using fts5(document)")
self.db.execute("create table if not exists files (path text primary key, digest text)")
self.db.commit()
self.lock = threading.RLock()
@@ -62,7 +164,8 @@ class Index:
def encoder(self):
if self.model is None:
from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(MODEL, prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "})
self.model = SentenceTransformer(MODEL, truncate_dim=EMBEDDING_DIMENSIONS,
prompts={"query": f"Instruct: {RETRIEVAL_INSTRUCTION}\nQuery: "})
return self.model
def reranker(self):
@@ -72,11 +175,16 @@ class Index:
return self.reranker_model
def rerank(self, query: str, rows: list[dict], limit: int):
candidates = rows[:min(len(rows), max(limit, 20))]
scores = self.reranker().predict([(query, item["text"]) for item in candidates])
candidates = rows[:RERANK_CANDIDATES]
scores = self.reranker().predict([(query, document_text(item["path"], item["symbol"], item["text"]))
for item in candidates], batch_size=8)
for item, score in zip(candidates, scores):
item["score"] = round(float(score), 6)
return sorted(candidates, key=lambda item: item["score"], reverse=True)[:limit]
return sorted(candidates, key=lambda item: item["score"], reverse=True)[:min(limit, FINAL_RESULTS)]
def delete_path(self, path):
self.db.execute("delete from lexical where rowid in (select rowid from chunks where path = ?)", (path,))
self.db.execute("delete from chunks where path = ?", (path,))
def rebuild(self):
with self.lock:
@@ -93,17 +201,22 @@ class Index:
old_digest = previous.pop(relative, None)
if not force and digest == old_digest:
continue
values = list(chunks(content.decode("utf-8", errors="ignore")))
vectors = self.encoder().encode(values, normalize_embeddings=True) if values else []
self.db.execute("delete from chunks where path = ?", (relative,))
self.db.executemany("insert into chunks values (?, ?, ?)",
((relative, value, vector.tobytes()) for value, vector in zip(values, vectors)))
values = list(chunks(content.decode("utf-8", errors="ignore"), relative, self.encoder().tokenizer))
documents = [document_text(relative, symbol, value) for symbol, value in values]
vectors = self.encoder().encode(documents, normalize_embeddings=True) if documents else []
self.delete_path(relative)
for (symbol, value), document, vector in zip(values, documents, vectors):
if len(vector) != EMBEDDING_DIMENSIONS:
raise ValueError("Embedding must have 1024 dimensions")
cursor = self.db.execute("insert into chunks values (?, ?, ?, ?)",
(relative, symbol, value, vector.tobytes()))
self.db.execute("insert into lexical(rowid, document) values (?, ?)", (cursor.lastrowid, document))
self.db.execute("insert or replace into files values (?, ?)", (relative, digest))
for relative in previous:
self.db.execute("delete from chunks where path = ?", (relative,))
self.delete_path(relative)
self.db.execute("delete from files where path = ?", (relative,))
def search(self, query: str, top_k: int):
def search(self, query: str, top_k: int = FINAL_RESULTS):
if not query.strip():
raise ValueError("query must not be empty")
with self.lock:
@@ -112,10 +225,23 @@ class Index:
return []
import numpy as np
vector = self.encoder().encode([query], prompt_name="query", normalize_embeddings=True)[0]
rows = ({"path": path, "score": float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))), "text": text}
for path, text, blob in self.db.execute("select path, text, vector from chunks"))
limit = max(1, min(top_k, 20))
return self.rerank(query, heapq.nlargest(20, rows, key=lambda item: item["score"]), limit)
if len(vector) != EMBEDDING_DIMENSIONS:
raise ValueError("Embedding must have 1024 dimensions")
scores = ((identifier, float(np.dot(vector, np.frombuffer(blob, dtype=np.float32))))
for identifier, blob in self.db.execute("select rowid, vector from chunks"))
dense = [identifier for identifier, _ in heapq.nlargest(DENSE_CANDIDATES, scores, key=lambda item: item[1])]
terms = list(dict.fromkeys(re.findall(r"[^\W_]+", query)))[:128]
expression = " OR ".join(f'"{term}"' for term in terms)
lexical = [row[0] for row in self.db.execute(
"select rowid from lexical where lexical match ? order by bm25(lexical), rowid limit ?",
(expression, LEXICAL_CANDIDATES))] if terms else []
identifiers = reciprocal_rank_fusion(dense, lexical)[:RERANK_CANDIDATES]
placeholders = ",".join("?" for _ in identifiers)
rows = {identifier: {"path": path, "symbol": symbol, "text": text}
for identifier, path, symbol, text in self.db.execute(
f"select rowid, path, symbol, text from chunks where rowid in ({placeholders})", identifiers)}
candidates = [rows[identifier] for identifier in identifiers]
return self.rerank(query, candidates, max(1, min(top_k, FINAL_RESULTS)))
def main():
@@ -132,7 +258,7 @@ def main():
mcp = FastMCP("my-ai-qwen3-retrieval")
@mcp.tool()
def semantic_search(query: str, top_k: int = 5) -> list[dict]:
def semantic_search(query: str, top_k: int = FINAL_RESULTS) -> list[dict]:
"""Search repository concepts, behavior, and implementation patterns first with Qwen3 embedding and reranking. Verify returned passages against current files; use direct reads for known paths."""
return index.search(query, top_k)