Экзамен: RAG-агент с ChromaDB и веб-поиском: vectorstore.py

vectorstore.py

@@ -0,0 +1,96 @@
+# vectorstore.py
+"""
+Модуль для работы с ChromaDB и OllamaEmbeddings.
+"""
+
+import os
+from pathlib import Path
+from typing import Iterable
+
+from langchain_ollama.embeddings import OllamaEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_chroma import Chroma
+from langchain.docstore.document import Document
+
+
+def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
+    """
+    Создаёт (или загружает) коллекцию ChromaDB с OllamaEmbeddings.
+
+    Args:
+        persist_directory: Путь, где будет храниться база данных.
+                            Если директория не существует – создаётся автоматически.
+
+    Returns:
+        Chroma объект, готовый к работе.
+    """
+    # Создаём путь и при необходимости создаём папку
+    path = Path(persist_directory)
+    path.mkdir(parents=True, exist_ok=True)
+
+    # Инициализируем эмбеддер Ollama
+    embeddings = OllamaEmbeddings(model="nomic-embed-text")
+
+    # Создаём/загружаем коллекцию Chroma
+    vectorstore = Chroma(
+        persist_directory=str(path),
+        embedding_function=embeddings,
+    )
+    return vectorstore
+
+
+def load_documents(directory: str | Path, vectorstore: Chroma) -> None:
+    """
+    Читает все .txt и .md файлы из указанной директории, разбивает их на чанки
+    и добавляет в коллекцию ChromaDB.
+
+    Args:
+        directory: Путь к папке с документами.
+        vectorstore: Объект Chroma, куда будут загружены документы.
+    """
+    # Подготовка текстового разделителя
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1000,
+        chunk_overlap=200,
+    )
+
+    docs_to_add: list[Document] = []
+
+    # Ищем файлы .txt и .md рекурсивно
+    for file_path in Path(directory).rglob("*"):
+        if not file_path.is_file():
+            continue
+        if file_path.suffix.lower() not in {".txt", ".md"}:
+            continue
+
+        try:
+            text = file_path.read_text(encoding="utf-8")
+        except Exception as exc:
+            # Если файл не читается – пропускаем и выводим предупреждение
+            print(f"⚠️ Не удалось прочитать {file_path}: {exc}")
+            continue
+
+        # Разбиваем текст на чанки
+        chunks = splitter.split_text(text)
+
+        # Создаём Document объекты с метаданными
+        for i, chunk in enumerate(chunks):
+            docs_to_add.append(
+                Document(
+                    page_content=chunk,
+                    metadata={
+                        "source": str(file_path),
+                        "page": i + 1,
+                    },
+                )
+            )
+
+    if not docs_to_add:
+        print("⚠️ В указанной папке не найдено подходящих файлов.")
+        return
+
+    # Добавляем документы в коллекцию
+    vectorstore.add_documents(docs_to_add)
+
+    # Сохраняем изменения (persist)
+    vectorstore.persist()