diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/init_loader.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/init_loader.py new file mode 100644 index 0000000..1baa6eb --- /dev/null +++ b/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) \ No newline at end of file