diff --git a/load_docs.py b/load_docs.py new file mode 100644 index 0000000..9f98af6 --- /dev/null +++ b/load_docs.py @@ -0,0 +1,49 @@ +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 +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}") + +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