92 lines
3.8 KiB
Python
92 lines
3.8 KiB
Python
import os
|
||
from pathlib import Path
|
||
from typing import Iterable
|
||
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain_chroma import Chroma
|
||
|
||
|
||
def create_vectorstore(persist_directory: str = "./chroma_db") -> 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")
|
||
|
||
# Инициализируем Chroma с указанным каталогом хранения
|
||
vectorstore = Chroma(
|
||
persist_directory=persist_directory,
|
||
embedding_function=embeddings,
|
||
)
|
||
return vectorstore
|
||
|
||
|
||
def _load_text_files(directory: str) -> Iterable[str]:
|
||
"""
|
||
Генератор, возвращающий содержимое всех .txt и .md файлов из указанной папки.
|
||
Папка может быть вложенной – рекурсивно ищем файлы.
|
||
|
||
Args:
|
||
directory (str): Корневая директория для поиска файлов.
|
||
|
||
Yields:
|
||
str: Текстовый контент файла.
|
||
"""
|
||
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
|
||
и добавляется в vectorstore.
|
||
|
||
Args:
|
||
directory (str): Путь к каталогу, содержащему .txt/.md файлы.
|
||
vectorstore (Chroma): Экземпляр vectorstore, куда будут добавлены документы.
|
||
"""
|
||
# Читаем все текстовые файлы из директории
|
||
texts = list(_load_text_files(directory))
|
||
if not texts:
|
||
print(f"В каталоге {directory} не найдено .txt или .md файлов.")
|
||
return
|
||
|
||
# Разбиваем каждый документ на чанки
|
||
splitter = RecursiveCharacterTextSplitter(
|
||
chunk_size=1000, # размер чанка в символах (можно настроить)
|
||
chunk_overlap=200, # перекрытие для сохранения контекста
|
||
)
|
||
all_chunks = []
|
||
for text in texts:
|
||
chunks = splitter.split_text(text)
|
||
all_chunks.extend(chunks)
|
||
|
||
if not all_chunks:
|
||
print("После чанкинга не осталось текста.")
|
||
return
|
||
|
||
# Добавляем чанки в vectorstore
|
||
vectorstore.add_texts(all_chunks)
|
||
|
||
# Сохраняем изменения (persist)
|
||
vectorstore.persist()
|
||
print(f"Загружено {len(all_chunks)} чанков из {directory} в ChromaDB.") |