Добавил возможность индексации директории, указываемой из интерфейса

This commit is contained in:
ki.sagidullin
2026-04-01 09:04:19 +05:00
parent b31e1377b3
commit 4816950382
2 changed files with 67 additions and 11 deletions

View File

@@ -1,4 +1,5 @@
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 (
@@ -10,6 +11,12 @@ 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, collection_name="documents", persist_dir="./qdrant_data"):
# Инициализируем клиент Qdrant
@@ -43,28 +50,33 @@ class LocalKnowledgeBase:
print(f"Создана коллекция {self.collection_name}")
def load_documents(self, directory_path: str):
"""Загрузка документов из директории"""
"""Загрузка документов из директории (рекурсивно)"""
loaders = {
".pdf": PyPDFLoader,
".txt": lambda path: DirectoryLoader(path, glob="**/*.txt"),
".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"), # добавить
".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 = []
for ext, loader_class in loaders.items():
for file_path in os.listdir(directory_path):
if file_path.endswith(ext):
full_path = os.path.join(directory_path, file_path)
# Рекурсивный обход с фильтрацией папок
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 = loader_class(full_path)
loader = loaders[ext](full_path)
documents = loader.load()
all_documents.extend(documents)
print(f"Загружен {file_path}: {len(documents)} страниц")
print(f"Загружен {file_path}: {len(documents)} документов")
except Exception as e:
print(f"Ошибка загрузки {file_path}: {e}")