fix: load_docs.py — Агент с RAG-памятью

This commit is contained in:
2026-07-02 09:38:31 +00:00
parent 7f446cd9c0
commit 49e6e09d63
+14 -44
View File
@@ -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]
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}")
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)
if len(sys.argv) != 2:
print("Usage: python load_docs.py <directory>")
sys.exit(1)
load_documents(sys.argv[1])