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

vectorstore.py

@@ -1,56 +1,92 @@
 import os
 from pathlib import Path
-from typing import List
+from typing import Iterable
 
 from langchain_ollama import OllamaEmbeddings
-from langchain_chroma import Chroma
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.docstore.document import Document
+from langchain_chroma import Chroma
 
 
 def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
     """
-    Создаёт и возвращает объект Chroma, подключённый к указанной директории.
-    Если директория не существует, она будет создана.
+    Создаёт или загружает коллекцию ChromaDB с эмбеддингами Ollama.
+
+    Args:
+        persist_directory (str): Путь к каталогу, где хранится база данных.
+            Если каталог не существует – будет создан новый.
+
+    Returns:
+        Chroma: объект vectorstore, готовый к добавлению документов и поиску.
     """
+    # Убедимся, что директория для хранения существует
     Path(persist_directory).mkdir(parents=True, exist_ok=True)
+
+    # Создаём эмбеддер Ollama (модель nomic-embed-text)
     embeddings = OllamaEmbeddings(model="nomic-embed-text")
-    return Chroma(
+
+    # Инициализируем Chroma с указанным каталогом хранения
+    vectorstore = Chroma(
         persist_directory=persist_directory,
         embedding_function=embeddings,
     )
+    return vectorstore
 
 
-def _load_text_files(directory: str) -> List[Document]:
+def _load_text_files(directory: str) -> Iterable[str]:
     """
-    Читает все .txt и .md файлы из указанной директории и возвращает список Document.
+    Генератор, возвращающий содержимое всех .txt и .md файлов из указанной папки.
+    Папка может быть вложенной – рекурсивно ищем файлы.
+
+    Args:
+        directory (str): Корневая директория для поиска файлов.
+
+    Yields:
+        str: Текстовый контент файла.
     """
-    docs: List[Document] = []
-    for file_path in Path(directory).rglob("*"):
-        if file_path.suffix.lower() in {".txt", ".md"}:
-            text = file_path.read_text(encoding="utf-8")
-            docs.append(Document(page_content=text, metadata={"source": str(file_path)}))
-    return docs
+    for root, _, files in os.walk(directory):
+        for file_name in files:
+            if file_name.lower().endswith((".txt", ".md")):
+                path = Path(root) / file_name
+                try:
+                    with open(path, "r", encoding="utf-8") as f:
+                        yield f.read()
+                except Exception as exc:  # pragma: no cover
+                    print(f"Не удалось прочитать {path}: {exc}")
 
 
 def load_documents(directory: str, vectorstore: Chroma) -> None:
     """
-    Загружает документы из указанной директории в ChromaDB.
-    Делает чанкинг с помощью RecursiveCharacterTextSplitter и добавляет в коллекцию.
+    Загружает документы из указанной папки в коллекцию ChromaDB.
+    Каждый документ разбивается на чанки с помощью RecursiveCharacterTextSplitter
+    и добавляется в vectorstore.
+
+    Args:
+        directory (str): Путь к каталогу, содержащему .txt/.md файлы.
+        vectorstore (Chroma): Экземпляр vectorstore, куда будут добавлены документы.
     """
-    # Читаем файлы
-    raw_docs = _load_text_files(directory)
+    # Читаем все текстовые файлы из директории
+    texts = list(_load_text_files(directory))
+    if not texts:
+        print(f"В каталоге {directory} не найдено .txt или .md файлов.")
+        return
 
-    # Разбиваем на чанки
+    # Разбиваем каждый документ на чанки
     splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200,
-        separators=["\n\n", "\n", " ", ""],
+        chunk_size=1000,  # размер чанка в символах (можно настроить)
+        chunk_overlap=200,  # перекрытие для сохранения контекста
     )
-    chunks = splitter.split_documents(raw_docs)
+    all_chunks = []
+    for text in texts:
+        chunks = splitter.split_text(text)
+        all_chunks.extend(chunks)
 
-    # Добавляем в коллекцию
-    vectorstore.add_documents(chunks)
-    # Сохраняем изменения
+    if not all_chunks:
+        print("После чанкинга не осталось текста.")
+        return
+
+    # Добавляем чанки в vectorstore
+    vectorstore.add_texts(all_chunks)
+
+    # Сохраняем изменения (persist)
     vectorstore.persist()
-    print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}")
+    print(f"Загружено {len(all_chunks)} чанков из {directory} в ChromaDB.")