Агент с RAG‑памятью: init_loader.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/init_loader.py

@@ -0,0 +1,106 @@
+<|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)