Code RAG: индексируем код с помощью Tree-sitter и AST

Столкнулись с ситуацией: в monorepo на 500 000 строк нужно найти функцию обработки платежей, но grep выдает сотни совпадений. RAG по кодовой базе решает эту проблему, но только если чанкинг сохраняет структуру кода. Мы в таких проектах используем комбинацию Tree-sitter и AST для синтаксического разбора и разбивки на логические единицы: функции, классы, модули. Каждый чанк обогащается метаданными — именем, сигнатурой, docstring, импортами и полным путем в модульной нотации. Это позволяет семантическому поиску находить именно ту единицу кода, которая нужна, а не случайный кусок текста.

Зачем preserve структуру кода при чанкинге?

Обычный документный RAG режет текст на абзацы. Для кода это не работает: разрыв между сигнатурой и телом функции убивает контекст. Код имеет иерархию — функция внутри класса, класс внутри модуля. Мы сохраняем эту иерархию в метаданных: модульный путь, строки начала и конца, список методов для класса, декораторы для функции. Это позволяет при поиске по запросу «как реализован X» получить именно ту единицу кода, где X определён.

Как мы реализуем code-aware парсинг?

Мы построили индексер на основе Tree-sitter. Он парсит код на 50+ языков и даёт синтаксическое дерево. Для каждого узла (функция, класс, метод) извлекаем:

• имя и сигнатуру,
• docstring (если есть),
• тело функции/класса,
• декораторы и аннотации,
• список импортов (до 10).

Например, для Python используем ast для точного выделения:

import ast from tree_sitter import Language, Parser class CodebaseIndexer: def __init__(self): # Tree-sitter для syntax-aware парсинга PY_LANGUAGE = Language('build/languages.so', 'python') self.parser = Parser() self.parser.set_language(PY_LANGUAGE) def extract_python_units(self, file_path: str) -> list[dict]: """Извлечение функций и классов как отдельных единиц индексации""" with open(file_path, 'r', encoding='utf-8') as f: source = f.read() try: tree = ast.parse(source) except SyntaxError: return [{'text': source, 'type': 'file', 'file': file_path}] units = [] for node in ast.walk(tree): if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): # Получение исходного кода функции func_source = ast.get_source_segment(source, node) docstring = ast.get_docstring(node) units.append({ 'type': 'function', 'name': node.name, 'file': file_path, 'line_start': node.lineno, 'line_end': node.end_lineno, 'text': func_source, 'docstring': docstring or '', 'decorators': [ast.unparse(d) for d in node.decorator_list], 'signature': self._get_signature(node) }) elif isinstance(node, ast.ClassDef): class_source = ast.get_source_segment(source, node) docstring = ast.get_docstring(node) units.append({ 'type': 'class', 'name': node.name, 'file': file_path, 'line_start': node.lineno, 'line_end': node.end_lineno, 'text': class_source, 'docstring': docstring or '', 'methods': [m.name for m in ast.walk(node) if isinstance(m, ast.FunctionDef)] }) return units def _get_signature(self, func_node: ast.FunctionDef) -> str: args = [] for arg in func_node.args.args: annotation = f": {ast.unparse(arg.annotation)}" \ if arg.annotation else "" args.append(f"{arg.arg}{annotation}") return_type = f" -> {ast.unparse(func_node.returns)}" \ if func_node.returns else "" return f"def {func_node.name}({', '.join(args)}){return_type}" 

Обогащение метаданными: почему это важно?

Просто разбить код на чанки недостаточно. Для качественного поиска каждый чанк надо обогатить: добавить имя, сигнатуру, docstring, импорты и полный путь в модульной нотации. Это превращает плоский текст в структурированный объект, который при векторизации даёт более точные эмбеддинги. Мы формируем rich_text — комбинацию всех метаданных, которая подаётся на вход модели эмбеддингов.

class CodeMetadataEnricher: def enrich(self, unit: dict) -> dict: unit = unit.copy() # Создание rich text для эмбеддинга # Комбинирование имени, сигнатуры, docstring и кода rich_text_parts = [] if unit.get('name'): rich_text_parts.append(f"# {unit['name']}") if unit.get('signature'): rich_text_parts.append(f"Signature: {unit['signature']}") if unit.get('docstring'): rich_text_parts.append(f"Description: {unit['docstring']}") rich_text_parts.append(unit['text']) unit['rich_text'] = '\n\n'.join(rich_text_parts) # Извлечение импортов для контекста imports = re.findall(r'^(?:import|from)\s+\S+', unit['text'], re.MULTILINE) unit['imports'] = imports[:10] # Путь в виде breadcrumb parts = unit['file'].replace('\\', '/').split('/') unit['module_path'] = '.'.join( p.replace('.py', '') for p in parts if not p.startswith('.') ) return unit 

Индексация Git истории: что изменилось?

RAG по коду может отвечать не только на вопросы о структуре, но и об истории изменений. Мы индексируем последние 100 коммитов с diff и метаданными: автор, дата, сообщение, файлы. Это позволяет найти, когда и кем была изменена конкретная функция. Например, запрос «Кто правил calculate_total в прошлом месяце?» вернёт коммиты с этой функцией в diff.

import subprocess class GitHistoryIndexer: def get_recent_changes(self, repo_path: str, n: int = 100) -> list[dict]: """Индексация последних коммитов с diff""" result = subprocess.run( ['git', 'log', f'-{n}', '--format=%H|%an|%ae|%ad|%s'], cwd=repo_path, capture_output=True, text=True ) commits = [] for line in result.stdout.strip().split('\n'): if not line: continue hash_, author, email, date, subject = line.split('|', 4) # Получение diff для этого коммита diff_result = subprocess.run( ['git', 'diff', f'{hash_}^', hash_, '--stat'], cwd=repo_path, capture_output=True, text=True ) commits.append({ 'hash': hash_, 'author': author, 'date': date, 'message': subject, 'changes_summary': diff_result.stdout[:500], 'text': f"Commit: {subject}\nAuthor: {author}\nDate: {date}\n\nChanges: {diff_result.stdout[:500]}" }) return commits 

Как оценить качество code RAG?

Хорошая метрика: при вопросе «Как реализован X?» система должна вернуть функцию или класс, который реализует X, а не просто файл с похожим названием. Для оценки мы используем golden set из 50–100 вопросов с известными ответами (конкретными функциями). Precision@3 > 0.8 — хороший результат. Ниже — сравнение стратегий чанкинга:

Смешанный чанкинг даёт выигрыш в точности в 2 раза по сравнению с файловым. Мы используем именно этот подход: каждый чанк — функция или класс, а файл становится метаданным.

Какая модель эмбеддингов подходит для кода?

Для кода лучше использовать модели, обученные на программном коде, а не на общих текстах. Ниже — сравнение популярных вариантов:

Типичные ошибки при индексации кода

• Игнорирование docstring — без них модель не понимает назначение функции, recall падает на 30%.
• Чанкинг по строкам — разрывает логические блоки, precision снижается вдвое.
• Отсутствие метаданных — только код без имени и сигнатуры даёт эмбеддинг, похожий на случайный кусок текста.
• Пропуск Git-истории — теряется информация об авторстве и контексте изменений.
• Выбор не той модели эмбеддингов — модель для документов плохо работает на коде.

Что входит в работу?

• Аудит кодовой базы: оценка размера, языков, структуры репозитория.
• Проектирование пайплайна: выбор инструментов (Tree-sitter, векторная БД, модель эмбеддингов), настройка метаданных.
• Реализация индексации: написание парсера, обогащение, векторизация, загрузка в векторную БД.
• Тестирование: проверка на golden set, итеративное улучшение чанкинга и метаданных.
• Интеграция: настройка API для поиска, интеграция с IDE, чат-ботами или внутренними инструментами.
• Деплой и мониторинг: развёртывание, логирование, метрики качества (precision, recall, latency p99).

Сроки и результаты

Ориентировочные сроки — от 2 до 4 недель в зависимости от размера кодовой базы и сложности интеграции. Результаты: полностью индексированная кодовая база с code-aware чанкингом, API для семантического поиска, документация и обучение команды (1–2 часа), поддержка в течение месяца после сдачи.

Наш опыт — 5 лет на рынке, более 20 реализованных RAG-проектов для fintech, edtech и e-commerce. Гарантируем качество: precision@3 не ниже 0.8 на вашем golden set. Свяжитесь с нами — оценим проект за 1 день и предложим архитектуру вашего code RAG. Получите консультацию по оптимизации уже на первом созвоне.