diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..a4d3a71 --- /dev/null +++ b/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() \ No newline at end of file