From 2fd48076a418fba4ac333143d197e5cf5bef85da Mon Sep 17 00:00:00 2001
From: lonpatovaadelina <lonpatovaadelina@mail.ru>
Date: Mon, 1 Jun 2026 07:49:56 +0000
Subject: [PATCH] =?UTF-8?q?=D0=AD=D0=BA=D0=B7=D0=B0=D0=BC=D0=B5=D0=BD:=20R?=
 =?UTF-8?q?AG-=D0=B0=D0=B3=D0=B5=D0=BD=D1=82=20=D1=81=20ChromaDB=20=D0=B8?=
 =?UTF-8?q?=20=D0=B2=D0=B5=D0=B1-=D0=BF=D0=BE=D0=B8=D1=81=D0=BA=D0=BE?=
 =?UTF-8?q?=D0=BC:=20vectorstore.py?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 vectorstore.py | 56 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 vectorstore.py

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