Compare commits
13
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
68d51896ef | ||
|
|
a1429c225b | ||
|
|
6d2769415e | ||
|
|
c5d0ae0077 | ||
|
|
fcc1f4d249 | ||
|
|
bda123d4c1 | ||
|
|
c88e683c6d | ||
|
|
0f8aa0a25d | ||
|
|
e97404a0a1 | ||
|
|
fd37422de2 | ||
|
|
b76a993c75 | ||
|
|
685dbe61f4 | ||
|
|
7be729cb22 |
No files matched your search
+147
-14
@@ -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
|
||||
@@ -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__':
|
||||
|
||||
+149
-23
@@ -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()
|
||||
if value:
|
||||
yield value
|
||||
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 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)
|
||||
|
||||
|
||||
Reference in new issue
Block a user