96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
# 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() |