diff --git a/vectorstore.py b/vectorstore.py new file mode 100644 index 0000000..4827b73 --- /dev/null +++ b/vectorstore.py @@ -0,0 +1,56 @@ +import os +from pathlib import Path +from typing import List + +from langchain_ollama import OllamaEmbeddings +from langchain_chroma import Chroma +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.docstore.document import Document + + +def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: + """ + Создаёт и возвращает объект Chroma, подключённый к указанной директории. + Если директория не существует, она будет создана. + """ + Path(persist_directory).mkdir(parents=True, exist_ok=True) + embeddings = OllamaEmbeddings(model="nomic-embed-text") + return Chroma( + persist_directory=persist_directory, + embedding_function=embeddings, + ) + + +def _load_text_files(directory: str) -> List[Document]: + """ + Читает все .txt и .md файлы из указанной директории и возвращает список Document. + """ + 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 + + +def load_documents(directory: str, vectorstore: Chroma) -> None: + """ + Загружает документы из указанной директории в ChromaDB. + Делает чанкинг с помощью RecursiveCharacterTextSplitter и добавляет в коллекцию. + """ + # Читаем файлы + raw_docs = _load_text_files(directory) + + # Разбиваем на чанки + splitter = RecursiveCharacterTextSplitter( + chunk_size=1000, + chunk_overlap=200, + separators=["\n\n", "\n", " ", ""], + ) + chunks = splitter.split_documents(raw_docs) + + # Добавляем в коллекцию + vectorstore.add_documents(chunks) + # Сохраняем изменения + vectorstore.persist() + print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}") \ No newline at end of file