Добавил возможность индексации директории, указываемой из интерфейса
This commit is contained in:
@@ -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}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user