import os from pathlib import Path from qdrant_client import QdrantClient from qdrant_client.models import Distance, VectorParams, PointStruct from langchain_community.document_loaders import ( TextLoader, PyPDFLoader, DirectoryLoader, ) from langchain_text_splitters import Language, RecursiveCharacterTextSplitter from langchain_community.embeddings import OllamaEmbeddings def should_skip(path: str) -> bool: """Проверяет, нужно ли пропустить путь (системные папки)""" skip_dirs = {'.git', 'build', 'vendor', 'node_modules', '__pycache__', '.venv', 'venv'} return any(part in skip_dirs for part in Path(path).parts) class LocalKnowledgeBase: def __init__(self, project_name: str = "default", persist_dir="./qdrant_data"): # Инициализируем клиент Qdrant self.client = QdrantClient( path=persist_dir, prefer_grpc=True # Локальное хранение ) self.persist_dir = persist_dir self.set_project(project_name) # Инициализация эмбеддингов self.embeddings = OllamaEmbeddings( model="nomic-embed-text", # Хорошие локальные эмбеддинги ) def set_project(self, project_name: str): """Переключиться на проект (создать/выбрать коллекцию)""" # Очищаем имя от недопустимых символов safe_name = "".join(c for c in project_name if c.isalnum() or c in "-_") self.collection_name = safe_name self._create_collection() print(f"Переключено на проект: {self.collection_name}") def get_current_project(self) -> str: """Получить имя текущего проекта""" return self.collection_name def list_projects(self) -> list: """Список всех проектов (коллекций)""" try: collections = self.client.get_collections() return [c.name for c in collections.collections] except: return [] def delete_project(self, project_name: str = None): """Удалить проект (коллекцию)""" name = project_name or self.collection_name try: self.client.delete_collection(collection_name=name) print(f"Удалена коллекция: {name}") except Exception as e: print(f"Не удалось удалить коллекцию {name}: {e}") def _create_collection(self): try: self.client.get_collection(self.collection_name) print(f"Коллекция {self.collection_name} уже существует") except: # Создаем новую коллекцию self.client.create_collection( collection_name=self.collection_name, vectors_config=VectorParams( size=768, # Размерность эмбеддингов nomic-embed-text distance=Distance.COSINE, ), ) print(f"Создана коллекция {self.collection_name}") def load_documents(self, directory_path: str): # Кастомные разделители для C++ - в порядке приоритета cpp_separators = [ # Конец блока - самый важный разделитель "\n}", # Конец statement ";", # Объявления - должны быть отдельными чанками "\nclass ", "\nstruct ", "\nenum ", "\nnamespace ", "\ntemplate<", "\npublic:", "\nprivate:", "\nprotected:", "\nvirtual ", "\noverride ", # Функции "\nvoid ", "\nint ", "\nfloat ", "\ndouble ", "\nbool ", "\nauto ", "\nstd::", "\nconst ", "\nstatic ", "\nexplicit ", # Комментарии - тоже разделитель "\n//", "\n/*", "\n*", # Перенос строки "\n", ] loaders = { ".pdf": PyPDFLoader, ".txt": lambda path: TextLoader(path, encoding="utf-8"), ".docx": lambda path: DirectoryLoader(path, glob="**/*.docx"), ".hpp": lambda path: TextLoader(path, encoding="utf-8"), ".cpp": lambda path: TextLoader(path, encoding="utf-8"), ".h": lambda path: TextLoader(path, encoding="utf-8"), ".cc": lambda path: TextLoader(path, encoding="utf-8"), } all_documents = [] files_processed = 0 # Рекурсивный обход с фильтрацией папок for root, dirs, files in os.walk(directory_path): # Фильтруем директории на месте dirs[:] = [d for d in dirs if not should_skip(os.path.join(root, d))] for file_path in files: ext = os.path.splitext(file_path)[1] if ext in loaders: full_path = os.path.join(root, file_path) try: loader = loaders[ext](full_path) documents = loader.load() # Добавляем метаданные о пути к файлу for doc in documents: doc.metadata["source"] = full_path doc.metadata["file_name"] = file_path doc.metadata["file_ext"] = ext all_documents.extend(documents) files_processed += 1 print(f"✓ Загружен {file_path}") except Exception as e: print(f"✗ Ошибка загрузки {file_path}: {e}") print(f"\n📁 Загружено файлов: {files_processed}") print(f"📄 Всего документов до сплиттинга: {len(all_documents)}") # Для C++ файлов используем РАЗНЫЕ стратегии для .cpp и .h/.hpp cpp_source_docs = [d for d in all_documents if d.metadata.get("file_ext") in [".cpp", ".cc"]] cpp_header_docs = [d for d in all_documents if d.metadata.get("file_ext") in [".hpp", ".h"]] other_docs = [d for d in all_documents if d.metadata.get("file_ext") not in [".cpp", ".hpp", ".h", ".cc"]] chunks = [] # === .cpp файлы: приоритетные, сохраняем контекст === if cpp_source_docs: print("\n🔧 Сплиттинг .cpp файлов (приоритетные)...") cpp_source_separators = [ # Только логические блоки - функции, классы, namespaces "\nclass ", "\nstruct ", "\nnamespace ", "\nvoid ", "\nint ", "\nfloat ", "\ndouble ", "\nbool ", "\nauto ", "\ntemplate<", "\nconst ", "\nstatic ", "\nexplicit ", ] cpp_source_splitter = RecursiveCharacterTextSplitter( separators=cpp_source_separators, chunk_size=1500, # Увеличили для сохранения контекста chunk_overlap=200, # Больше overlap length_function=len, keep_separator=True, ) source_chunks = cpp_source_splitter.split_documents(cpp_source_docs) # Добавляем приоритет в метаданные for chunk in source_chunks: chunk.metadata["priority"] = 1 # Высший приоритет chunks.extend(source_chunks) print(f" .cpp чанков: {len(source_chunks)}") # === .h/.hpp файлы: менее агрессивный сплиттинг === if cpp_header_docs: print("\n🔧 Сплиттинг header файлов...") cpp_header_separators = [ # Только большие логические блоки "\nclass ", "\nstruct ", "\nnamespace ", "\ntemplate<", "\npublic:", "\nprivate:", "\nprotected:", "\nvirtual ", ] cpp_header_splitter = RecursiveCharacterTextSplitter( separators=cpp_header_separators, chunk_size=1200, chunk_overlap=150, length_function=len, keep_separator=True, ) header_chunks = cpp_header_splitter.split_documents(cpp_header_docs) # Добавляем пониженный приоритет for chunk in header_chunks: chunk.metadata["priority"] = 2 # Низший приоритет chunks.extend(header_chunks) print(f" Header чанков: {len(header_chunks)}") # Для остальных файлов - стандартный сплиттер if other_docs: print("\n📄 Сплиттинг остальных файлов...") text_splitter = RecursiveCharacterTextSplitter( chunk_size=1000, chunk_overlap=100, length_function=len ) other_chunks = text_splitter.split_documents(other_docs) for chunk in other_chunks: chunk.metadata["priority"] = 3 chunks.extend(other_chunks) print(f" Остальных чанков: {len(other_chunks)}") print(f"\n📊 Всего чанков после сплиттинга: {len(chunks)}") # Создаем эмбеддинги и сохраняем в Qdrant self._index_documents(chunks) return chunks def _index_documents(self, documents): """Индексация документов в Qdrant""" points = [] for i, doc in enumerate(documents): # Создаем эмбеддинг для каждого чанка embedding = self.embeddings.embed_query(doc.page_content) # Извлекаем название функции/класса/строки из контекста metadata = doc.metadata # Простая эвристика: ищем объявление функции/класса в начале чанка import re # Ищем паттерны: "void funcName(...)", "class ClassName", "int funcName(...)" func_match = re.search(r'(?:class|struct|enum)\s+(\w+)', doc.page_content) if not func_match: func_match = re.search(r'(?:void|int|float|double|bool|auto)\s+(\w+)\s*\(', doc.page_content) function_name = func_match.group(1) if func_match else "unknown" point = PointStruct( id=i, vector=embedding, payload={ "text": doc.page_content, "source": doc.metadata.get("source", "unknown"), "page": doc.metadata.get("page", 0), "file_type": doc.metadata.get("file_ext", ""), "priority": doc.metadata.get("priority", 3), # Добавить "function_name": function_name, # ✅ Добавил "class_name": func_match.group(1) if func_match else None, # ✅ Добавил "line_start": doc.metadata.get("line", 0), # Если есть в метаданных }, # payload={ # "text": doc.page_content, # "source": doc.metadata.get("source", "unknown"), # "page": doc.metadata.get("page", 0), # "file_type": "cpp", # }, ) points.append(point) # Пакетная загрузка каждые 100 точек if len(points) >= 100: self.client.upsert(collection_name=self.collection_name, points=points) points = [] print(f"Индексировано {i+1} документов") # Загружаем оставшиеся if points: self.client.upsert(collection_name=self.collection_name, points=points) print(f"Индексация завершена. Всего документов: {len(documents)}") def search(self, query: str, top_k: int = 5): """Поиск в базе знаний""" # Создаем эмбеддинг запроса query_embedding = self.embeddings.embed_query(query) # Ищем в Qdrant search_result = self.client.query_points( collection_name=self.collection_name, query=query_embedding, limit=top_k, ) # Форматируем результаты results = [] for hit in search_result.points: # Добавляем "штраф" для header-файлов score = hit.score if hit.payload.get("file_ext") in [".h", ".hpp"]: score *= 0.7 # Снижаем приоритет на 30% results.append({ "text": hit.payload["text"], "score": score, "source": hit.payload.get("source", "unknown"), "file_ext": hit.payload.get("file_ext", ""), }) # Пересортируем по новому score results.sort(key=lambda x: x["score"], reverse=True) return results[:top_k]