49 lines
1.8 KiB
Python
49 lines
1.8 KiB
Python
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) |