106 lines
4.0 KiB
Python
106 lines
4.0 KiB
Python
<|channel|>final code<|message|>import os
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_qdrant import QdrantVectorStore
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain.document_loaders import TextLoader, DirectoryLoader
|
||
|
||
|
||
def _get_documents_from_dir(directory: str) -> List[str]:
|
||
"""
|
||
Возвращает список путей к текстовым файлам в указанной директории.
|
||
"""
|
||
dir_path = Path(directory)
|
||
if not dir_path.is_dir():
|
||
raise ValueError(f"Путь {directory} не является директорией.")
|
||
# Поддерживаем только .txt, .md и .pdf (если понадобится)
|
||
files = list(dir_path.rglob("*"))
|
||
return [str(p) for p in files if p.suffix.lower() in {".txt", ".md"}]
|
||
|
||
|
||
def _load_documents(file_paths: List[str]) -> List[dict]:
|
||
"""
|
||
Загружает содержимое файлов в список словарей с полями 'content' и 'metadata'.
|
||
"""
|
||
documents = []
|
||
for path in file_paths:
|
||
loader = TextLoader(path, encoding="utf-8")
|
||
docs = loader.load()
|
||
# LangChain возвращает объекты Document; преобразуем их в dict
|
||
for doc in docs:
|
||
documents.append(
|
||
{
|
||
"content": doc.page_content,
|
||
"metadata": {"source": path},
|
||
}
|
||
)
|
||
return documents
|
||
|
||
|
||
def _chunk_documents(documents: List[dict], chunk_size: int = 1000, chunk_overlap: int = 200) -> List[str]:
|
||
"""
|
||
Разбивает документы на чанки с помощью RecursiveCharacterTextSplitter.
|
||
Возвращает список строк-чанков.
|
||
"""
|
||
splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=chunk_size,
|
||
chunk_overlap=chunk_overlap,
|
||
)
|
||
chunks = []
|
||
for doc in documents:
|
||
splits = splitter.split_text(doc["content"])
|
||
# Добавляем метаданные в конец чанка для удобства поиска
|
||
for s in splits:
|
||
chunks.append(s)
|
||
return chunks
|
||
|
||
|
||
def _embed_and_store(chunks: List[str], collection_name: str, host: str = "localhost", port: int = 6333):
|
||
"""
|
||
Создаёт клиент Qdrant и сохраняет чанки с эмбеддингами Ollama.
|
||
"""
|
||
# Инициализация векторного хранилища
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
vectorstore = QdrantVectorStore(
|
||
client_kwargs={"host": host, "port": port},
|
||
collection_name=collection_name,
|
||
embedding_function=embeddings,
|
||
)
|
||
# Добавляем чанки в коллекцию
|
||
vectorstore.add_texts(chunks)
|
||
|
||
|
||
def load_directory_to_qdrant(directory: str, collection_name: str = "knowledge_base"):
|
||
"""
|
||
Полный пайплайн загрузки документов из директории в Qdrant.
|
||
"""
|
||
file_paths = _get_documents_from_dir(directory)
|
||
if not file_paths:
|
||
print(f"В каталоге {directory} не найдено текстовых файлов.")
|
||
return
|
||
|
||
documents = _load_documents(file_paths)
|
||
chunks = _chunk_documents(documents)
|
||
|
||
_embed_and_store(chunks, collection_name)
|
||
|
||
print(
|
||
f"Загружено {len(chunks)} чанков из {len(file_paths)} документов в коллекцию '{collection_name}'."
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="Загрузка документов в Qdrant.")
|
||
parser.add_argument("directory", help="Путь к директории с документами")
|
||
parser.add_argument(
|
||
"--collection",
|
||
default="knowledge_base",
|
||
help="Имя коллекции в Qdrant (по умолчанию: knowledge_base)",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
load_directory_to_qdrant(args.directory, args.collection) |