Multimedika/Bot_Development

1import logging2import re3import os4import pytz5 6from typing import List7from datetime import datetime8from datetime import timedelta9from fastapi.responses import JSONResponse10from script.vector_db import IndexManager11from llama_index.core.llms import MessageRole12 13from core.chat.engine import Engine14from core.chat.chatstore import ChatStore15from core.parser import (16    filter_metadata_by_pages,17    extract_sorted_page_numbers18 19)20 21from service.dto import ChatMessage22from pymongo.mongo_client import MongoClient23 24 25class ChatCompletionService:26    def __init__(27        self,28        session_id: str,29        user_request: str,30        titles: List = None,31        type_bot: str = "general",32    ):33        self.session_id = session_id34        self.user_request = user_request35        self.titles = titles36        self.type_bot = type_bot37        self.client = MongoClient(os.getenv("MONGO_URI"))38        self.engine = Engine()39        self.index_manager = IndexManager()40        self.chatstore = ChatStore()41 42    def generate_completion(self):43        if not self._ping_mongo():44            return JSONResponse(45                status_code=500, content="Database Error: Unable to connect to MongoDB"46            )47 48        try:49            # Load and retrieve chat engine with appropriate index50            index = self.index_manager.load_existing_indexes()51            chat_engine = self._get_chat_engine(index)52 53            # Generate chat response54            response = chat_engine.chat(self.user_request)55 56            sources = response.source_nodes57 58            contents, metadata_collection, scores = self._process_sources_images(sources)59            60            # Update response and renumber sources61            response = str(response)62 63            # Add contents to metadata64            metadata_collection = self._attach_contents_to_metadata(65                contents, metadata_collection66            )67            68            page_sources = extract_sorted_page_numbers(response)69            metadata_collection = filter_metadata_by_pages(metadata_collection, page_sources)70 71            # Save the message to chat store72            self._store_message_in_chatstore(response, metadata_collection)73 74        except Exception as e:75            logging.error(f"An error occurred in generate text: {e}")76            return JSONResponse(77                status_code=500, content=f"An internal server error occurred: {e}"78            )79 80        try:81            if self.type_bot == "specific":82                self._save_chat_history_to_db(response, metadata_collection)83 84            return str(response), metadata_collection, scores85 86        except Exception as e:87            logging.error(f"An error occurred while saving chat history: {e}")88            return JSONResponse(89                status_code=500,90                content=f"An internal server error occurred while saving chat history: {e}",91            )92 93    def _ping_mongo(self):94        try:95            self.client.admin.command("ping")96            print("Pinged your deployment. Successfully connected to MongoDB!")97            return True98        except Exception as e:99            logging.error(f"MongoDB connection failed: {e}")100            return False101 102    def _get_chat_engine(self, index):103        if self.type_bot == "general":104            return self.engine.get_chat_engine(self.session_id, index)105        return self.engine.get_chat_engine(106            self.session_id, index, self.titles, self.type_bot107        )108 109    def _extract_sorted_references(self, response):110        number_reference = list(set(re.findall(r"\[(\d+)\]", str(response))))111        return sorted(number_reference)112 113    114    def _process_sources_images(self, sources):115        contents, metadata_collection, scores = [], [], []116 117        for number in range (len(sources)):118            if sources and len(sources) > 0:119                content = sources[number - 1].node.get_text()120                contents.append(content)121                metadata = dict(sources[number - 1].node.metadata)122                metadata_collection.append(metadata)123                score = sources[number - 1].score124                scores.append(score)125            else:126                print("No sources available")127 128        return contents, metadata_collection, scores129 130    def _attach_contents_to_metadata(self, contents, metadata_collection):131        for i in range(min(len(contents), len(metadata_collection))):132            metadata_collection[i]["content"] = contents[i]133        return metadata_collection134 135    def _store_message_in_chatstore(self, response, metadata_collection):136        message = ChatMessage(137            role=MessageRole.ASSISTANT, 138            content=response, 139            metadata=metadata_collection,140        )141        self.chatstore.delete_last_message(self.session_id)142        self.chatstore.add_message(self.session_id, message)143        self.chatstore.clean_message(self.session_id)144 145    def _save_chat_history_to_db(self, response, metadata_collection):146        jakarta_tz = pytz.timezone("Asia/Jakarta")147        time_now = datetime.now(jakarta_tz)148        user_timestamp = time_now - timedelta(seconds=0.2)149        chat_history_db = [150            ChatMessage(151                role=MessageRole.USER,152                content=self.user_request,153                timestamp=user_timestamp,154                payment="free" if self.type_bot == "general" else None,155            ),156            ChatMessage(157                role=MessageRole.ASSISTANT,158                content=response,159                metadata=metadata_collection,160                timestamp=time_now,161                payment="free" if self.type_bot == "general" else None,162            ),163        ]164 165        chat_history_json = [message.model_dump() for message in chat_history_db]166 167        db = self.client["bot_database"]  # Replace with your database name168        collection = db[self.session_id]  # Replace with your collection name169        collection.insert_many(chat_history_json)170