diff --git a/vectorstore.py b/vectorstore.py index 4827b73..533469f 100644 --- a/vectorstore.py +++ b/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}") \ No newline at end of file + print(f"Загружено {len(all_chunks)} чанков из {directory} в ChromaDB.") \ No newline at end of file