diff --git a/scripts/test-retrieval-server.py b/scripts/test-retrieval-server.py index 6af0a93..ca83dca 100644 --- a/scripts/test-retrieval-server.py +++ b/scripts/test-retrieval-server.py @@ -175,6 +175,30 @@ class RetrievalServerTests(unittest.TestCase): 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', '', '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_old_cache_is_preserved_but_not_reused(self): with TemporaryDirectory() as temporary: root = Path(temporary) diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index 985327a..b7a4358 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -4,11 +4,13 @@ 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 @@ -20,13 +22,77 @@ CHUNK_TOKENS = 400 OVERLAP_TOKENS = 64 RERANK_CANDIDATES = 20 FINAL_RESULTS = 20 -INDEX_VERSION = 0 +INDEX_VERSION = 1 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 sections(text: str, path: str): + lines = text.splitlines(keepends=True) + boundaries = {0: "", len(lines): ""} + 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 != "" 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, "") + 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 [("", text)]: + 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):