minimal change

app/data/generated_docs_db.py

@@ -0,0 +1,115 @@
+import json
+import os
+import sys
+from typing import Any, Dict, List, Optional
+
+import torch
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from loguru import logger
+
+# Добавляем путь к родительской директории для импорта config
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import settings
+
+
+def load_json_files(directory: str) -> List[Dict[str, Any]]:
+    """Загрузка всех JSON файлов из указанной директории."""
+    documents = []
+
+    try:
+        if not os.path.exists(directory):
+            logger.error(f"Директория {directory} не существует")
+            return documents
+
+        for filename in os.listdir(directory):
+            if filename.endswith(".json"):
+                file_path = os.path.join(directory, filename)
+                try:
+                    with open(file_path, "r", encoding="utf-8") as file:
+                        data = json.load(file)
+                        documents.append(
+                            {"text": data["text"], "metadata": data["metadata"]}
+                        )
+                        logger.info(f"Загружен файл: {filename}")
+                except Exception as e:
+                    logger.error(f"Ошибка при чтении файла {filename}: {e}")
+
+        logger.success(f"Загружено {len(documents)} JSON файлов")
+        return documents
+    except Exception as e:
+        logger.error(f"Ошибка при загрузке JSON файлов: {e}")
+        return documents
+
+
+def split_text_into_chunks(text: str, metadata: Dict[str, Any]) -> List[Any]:
+    """Разделение текста на чанки с сохранением метаданных."""
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=settings.MAX_CHUNK_SIZE,
+        chunk_overlap=settings.CHUNK_OVERLAP,
+        length_function=len,
+        is_separator_regex=False,
+    )
+
+    chunks = text_splitter.create_documents(texts=[text], metadatas=[metadata])
+    return chunks
+
+
+def generate_chroma_db() -> Optional[Chroma]:
+    """Инициализация ChromaDB с данными из JSON файлов."""
+    try:
+        # Создаем директорию для хранения базы данных, если она не существует
+        os.makedirs(settings.DOCS_CHROMA_PATH, exist_ok=True)
+
+        # Загружаем JSON файлы
+        documents = load_json_files(settings.PARSED_JSON_PATH)
+
+        if not documents:
+            logger.warning("Нет документов для добавления в базу данных")
+            return None
+
+        # Инициализируем модель эмбеддингов
+        embeddings = HuggingFaceEmbeddings(
+            model_name=settings.LM_MODEL_NAME,
+            model_kwargs={"device": "cuda" if torch.cuda.is_available() else "cpu"},
+            encode_kwargs={"normalize_embeddings": True},
+        )
+
+        # Подготавливаем данные для Chroma
+        all_chunks = []
+        for i, doc in enumerate(documents):
+            chunks = split_text_into_chunks(doc["text"], doc["metadata"])
+            all_chunks.extend(chunks)
+            logger.info(
+                f"Документ {i+1}/{len(documents)} разбит на {len(chunks)} чанков"
+            )
+
+        # Создаем векторное хранилище
+        texts = [chunk.page_content for chunk in all_chunks]
+        metadatas = [chunk.metadata for chunk in all_chunks]
+        ids = [f"doc_{i}" for i in range(len(all_chunks))]
+
+        chroma_db = Chroma.from_texts(
+            texts=texts,
+            embedding=embeddings,
+            ids=ids,
+            metadatas=metadatas,
+            persist_directory=settings.DOCS_CHROMA_PATH,
+            collection_name=settings.DOCS_COLLECTION_NAME,
+            collection_metadata={
+                "hnsw:space": "cosine",
+            },
+        )
+
+        logger.success(
+            f"База Chroma инициализирована, добавлено {len(all_chunks)} чанков из {len(documents)} документов"
+        )
+        return chroma_db
+    except Exception as e:
+        logger.error(f"Ошибка инициализации Chroma: {e}")
+        raise
+
+
+if __name__ == "__main__":
+    generate_chroma_db()