diff --git a/load_docs.py b/load_docs.py index 9f98af6..d690a62 100644 --- a/load_docs.py +++ b/load_docs.py @@ -1,49 +1,19 @@ -import os -import argparse -from dotenv import load_dotenv -from langchain_openai import OpenAIEmbeddings -from langchain_qdrant import QdrantVectorStore -from langchain_text_splitters import RecursiveCharacterTextSplitter +import sys +from pathlib import Path + +from utils import splitter, vector_store from langchain_core.documents import Document -from qdrant_client import QdrantClient -# Загрузка переменных окружения -load_dotenv() - -# Инициализация эмбеддингов -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), -) - -# Инициализация Qdrant -client = QdrantClient(url="http://localhost:6333") -collection_name = "knowledge_base" -vector_store = QdrantVectorStore( - client=client, - collection_name=collection_name, - embeddings=embeddings, -) - -# Чанкинг -splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - -def load_directory(dir_path: str): - """Загружает все .txt файлы из директории в векторную базу.""" - for root, dirs, files in os.walk(dir_path): - for file in files: - if file.lower().endswith(".txt"): - path = os.path.join(root, file) - with open(path, "r", encoding="utf-8") as f: - text = f.read() - chunks = splitter.split_text(text) - docs = [Document(page_content=chunk, metadata={"title": file}) for chunk in chunks] - vector_store.add_documents(docs) - print(f"Added {len(docs)} chunks from {file}") +def load_documents(dir_path: str): + for file_path in Path(dir_path).rglob("*.txt"): + content = file_path.read_text(encoding="utf-8") + chunks = splitter.split_text(content) + docs = [Document(page_content=chunk, metadata={"title": file_path.name}) for chunk in chunks] + vector_store.add_documents(docs) + print(f"Added {len(docs)} chunks from {file_path}") if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Load documents into Qdrant.") - parser.add_argument("directory", help="Path to directory with .txt files") - args = parser.parse_args() - load_directory(args.directory) \ No newline at end of file + if len(sys.argv) != 2: + print("Usage: python load_docs.py ") + sys.exit(1) + load_documents(sys.argv[1]) \ No newline at end of file