| |
| """RAG по своей кодовой базе (№1): индексация + поиск + ответ по проекту. |
| |
| Запуск: |
| python rag.py index <папка> — проиндексировать проект |
| python rag.py ask "где у меня логика авторизации?" [backend] |
| |
| Работает без тяжёлых зависимостей (пересечение слов). Для лучшего качества можно |
| поставить sentence-transformers и заменить search() на эмбеддинги с косинусной близостью. |
| """ |
| import json |
| import os |
| import re |
| import sys |
|
|
| from brain import pick, generate |
|
|
| INDEX = "rag_index.json" |
| SKIP = {"venv", ".venv", ".git", "__pycache__", "node_modules", ".idea"} |
|
|
|
|
| def chunk_file(path): |
| code = open(path, encoding="utf-8").read() |
| lines = code.splitlines() |
| chunks = [] |
| for i in range(0, len(lines), 50): |
| chunks.append((path, i + 1, "\n".join(lines[i:i + 50]))) |
| return chunks |
|
|
|
|
| def index(folder): |
| chunks = [] |
| for root, dirs, files in os.walk(folder): |
| dirs[:] = [d for d in dirs if d not in SKIP] |
| for f in files: |
| if f.endswith(".py"): |
| chunks.extend(chunk_file(os.path.join(root, f))) |
| json.dump(chunks, open(INDEX, "w", encoding="utf-8"), ensure_ascii=False) |
| print(f"проиндексировано чанков: {len(chunks)} -> {INDEX}") |
|
|
|
|
| def search(query, top=5): |
| if not os.path.exists(INDEX): |
| print(f"сначала проиндексируй: rag.py index <папка>") |
| return [] |
| chunks = json.load(open(INDEX, encoding="utf-8")) |
| qwords = set(re.findall(r"\w+", query.lower())) |
| scored = [] |
| for path, line, text in chunks: |
| twords = set(re.findall(r"\w+", text.lower())) |
| s = len(qwords & twords) |
| if s: |
| scored.append((s, path, line, text)) |
| scored.sort(reverse=True) |
| return scored[:top] |
|
|
|
|
| def ask(query, backend): |
| hits = search(query) |
| if not hits: |
| print("ничего релевантного не найдено") |
| return |
| context = "\n\n".join(f"--- {p}:{l} ---\n{t[:800]}" for _, p, l, t in hits) |
| msgs = [{"role": "system", |
| "content": "Ты анализируешь кодовую базу. Отвечай по предоставленному " |
| "контексту, указывай файлы и строки."}, |
| {"role": "user", |
| "content": f"Контекст из кодовой базы:\n{context}\n\nВопрос: {query}"}] |
| print(generate(backend, msgs, 1024)) |
|
|
|
|
| def main(): |
| args = sys.argv[1:] |
| if len(args) < 2: |
| print("usage: rag.py index <папка> | rag.py ask <вопрос> [backend]") |
| sys.exit(2) |
| if args[0] == "index": |
| index(args[1]) |
| elif args[0] == "ask": |
| ask(" ".join(args[1:]), pick()) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
