File size: 3,350 Bytes
63971d2
 
 
 
 
 
 
 
3ef184a
63971d2
90898ed
3ef184a
63971d2
 
 
3ef184a
63971d2
 
 
 
 
 
3ef184a
63971d2
 
 
 
 
 
 
 
 
3ef184a
63971d2
 
3ef184a
2651354
 
 
87dbdac
2651354
 
 
 
 
 
 
 
87dbdac
3ef184a
 
87dbdac
3ef184a
 
 
 
 
 
87dbdac
3ef184a
 
87dbdac
3ef184a
 
 
87dbdac
3ef184a
 
 
 
63971d2
662abb5
 
 
2d01c4b
662abb5
63971d2
3ef184a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import json
import os
from typing import List, Optional
from loguru import logger
from app.model.chat_model import ChatCompletion
from app.repository.chat_repository import ChatRepository
from app.config.db import db_config


class InitialSetup:
    """Initial setup manager for the application when database type is embedded"""

    def __init__(self):
        self._chat_repository: Optional[ChatRepository] = None
        self.data_dir = os.path.join(os.path.dirname(__file__), "data")

    @property
    def chat_repository(self) -> ChatRepository:
        """Lazy loading of ChatRepository"""
        if self._chat_repository is None:
            self._chat_repository = ChatRepository()
        return self._chat_repository

    def _load_initial_data(self) -> List[ChatCompletion]:
        """Load initial data from JSON files"""
        try:
            with open(os.path.join(self.data_dir, "initial_chat_completions.json"), "r") as f:
                data = json.load(f)
                return [ChatCompletion(**item) for item in data["chat_completions"]]
        except Exception as e:
            logger.error(f"Error loading initial data: {e}")
            return []

    async def setup(self) -> None:
        """Setup initial data if database type is embedded"""
        try:
            if db_config.DATABASE_TYPE != "embedded":
                logger.info("Skipping initial setup as database type is not embedded")
                return

            # if MONGO_URI is not set, it means we are using embedded database
            # last check is for the case of using mongomock-motor for database_type=embedded
            # so last exit before the bridge :) turkish joke
            if db_config.MONGO_URI is None:
                # delete all chat completions in the embedded database
                logger.warning("Deleting all chat completions in the embedded database")
                await self.chat_repository.db.chat_completion.delete_many({})
                logger.warning("Deleting all chat completions in the embedded database done")

            chat_completions = self._load_initial_data()
            logger.info(f"Loaded {len(chat_completions)} initial chat completions")

            for completion in chat_completions:
                try:
                    found_id = await self.chat_repository.find_by_id(completion.completion_id)
                    if found_id:
                        logger.debug(f"Chat completion already exists: {found_id}")
                        continue

                    saved = await self.chat_repository.save(completion)
                    logger.info(f"Successfully saved chat completion: {saved.completion_id}")

                except Exception as e:
                    logger.error(f"Error saving chat completion {completion.completion_id}: {e}")
                    raise

        except Exception as e:
            logger.error(f"Setup failed: {e}")
            raise

        # query of the saved chat completions
        saved_chat_completions = await self.chat_repository.find()
        logger.debug("********** Begin of Saved chat completions**********")
        logger.trace(f"{saved_chat_completions}")

        logger.debug("********** End of Saved chat completions**********")

        logger.info("Initial setup completed successfully for embedded database")