Індексація кодової бази для RAG

Індексація кодової бази для RAG

RAG за кодовою базою — основа для AI помічників коду, автоматичної документації та пошуку архітектурних рішень. Ключова відмінність від документного RAG: код має структуру (функції, класи, імпорти), яку потрібно зберігати при чанкуванні.

Аналіз коду з урахуванням синтаксису

import ast
from tree_sitter import Language, Parser

class CodebaseIndexer:
    def __init__(self):
        # Tree-sitter для синтаксично-усвідомленого аналізу
        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}"

Збагачення метаданими для пошуку

class CodeMetadataEnricher:
    def enrich(self, unit: dict) -> dict:
        unit = unit.copy()

        # Створення насиченого тексту для вбудовування
        # Комбінування імені, сигнатури, 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]

        # Шлях у вигляді хлібних крошок
        parts = unit['file'].replace('\\', '/').split('/')
        unit['module_path'] = '.'.join(
            p.replace('.py', '') for p in parts if not p.startswith('.')
        )

        return unit

Індексація історії Git

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

Оцінка якості RAG для коду

Хорошою метрикою є: при запитанні "Як реалізовано X?" система повинна повернути функцію або клас, який реалізує X, а не просто файл із подібною назвою. Для оцінки: створіть золотий набір із 50-100 запитань до вашої кодової бази з відомими відповідями (конкретними функціями). Precision@3 > 0.8 — це хороший результат.