315 lines
14 KiB
Python
315 lines
14 KiB
Python
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]
|
||
|