Multimedika/Bot_Development

1import logging2import re3import os4 5from typing import List6from datetime import datetime7from fastapi.responses import JSONResponse8from script.vector_db import IndexManager9from llama_index.core.llms import MessageRole10 11from core.chat.engine import Engine12from core.chat.chatstore import ChatStore13from core.parser import (14    update_response,15    sort_and_renumber_sources,16    redesign_structure_message,17)18 19from service.dto import ChatMessage20from pymongo.mongo_client import MongoClient21 22 23class ChatCompletionService:24    def __init__(25        self,26        session_id: str,27        user_request: str,28        titles: List = None,29        type_bot: str = "general",30    ):31        self.session_id = session_id32        self.user_request = user_request33        self.titles = titles34        self.type_bot = type_bot35        self.client = MongoClient(os.getenv("MONGO_URI"))36        self.engine = Engine()37        self.index_manager = IndexManager()38        self.chatstore = ChatStore()39 40    def generate_completion(self):41        if not self._ping_mongo():42            return JSONResponse(43                status_code=500, content="Database Error: Unable to connect to MongoDB"44            )45 46        try:47            # Load and retrieve chat engine with appropriate index48            index = self.index_manager.load_existing_indexes()49            chat_engine = self._get_chat_engine(index)50 51            # Generate chat response52            response = chat_engine.chat(self.user_request)53 54            sources = response.sources55            number_reference_sorted = self._extract_sorted_references(response)56 57            contents, metadata_collection, scores = self._process_sources(58                sources, number_reference_sorted59            )60 61            # Update response and renumber sources62            response = update_response(str(response))63            contents = sort_and_renumber_sources(contents)64 65            # Add contents to metadata66            metadata_collection = self._attach_contents_to_metadata(67                contents, metadata_collection68            )69 70            if self.type_bot == "general":71                response = redesign_structure_message(response, metadata_collection)72 73            # Save the message to chat store74            self._store_message_in_chatstore(response, metadata_collection)75 76        except Exception as e:77            logging.error(f"An error occurred in generate text: {e}")78            return JSONResponse(79                status_code=500, content=f"An internal server error occurred: {e}"80            )81 82        try:83            if self.type_bot == "specific":84                self._save_chat_history_to_db(response, metadata_collection)85 86            return str(response), metadata_collection, scores87 88        except Exception as e:89            logging.error(f"An error occurred while saving chat history: {e}")90            return JSONResponse(91                status_code=500,92                content=f"An internal server error occurred while saving chat history: {e}",93            )94 95    def _ping_mongo(self):96        try:97            self.client.admin.command("ping")98            print("Pinged your deployment. Successfully connected to MongoDB!")99            return True100        except Exception as e:101            logging.error(f"MongoDB connection failed: {e}")102            return False103 104    def _get_chat_engine(self, index):105        if self.type_bot == "general":106            return self.engine.get_chat_engine(self.session_id, index)107        return self.engine.get_chat_engine(108            self.session_id, index, self.titles, self.type_bot109        )110 111    def _extract_sorted_references(self, response):112        number_reference = list(set(re.findall(r"\[(\d+)\]", str(response))))113        return sorted(number_reference)114 115    def _process_sources(self, sources, number_reference_sorted):116        contents, metadata_collection, scores = [], [], []117        if not number_reference_sorted:118            print("There are no references")119            return contents, metadata_collection, scores120 121        for number in number_reference_sorted:122            number = int(number)123            if sources and len(sources) > 0:124                node = dict(sources[0])["raw_output"].source_nodes125                if 0 <= number - 1 < len(node):126                    content = node[number - 1].node.get_text()127                    contents.append(content)128                    metadata = dict(node[number - 1].node.metadata)129                    metadata_collection.append(metadata)130                    score = node[number - 1].score131                    scores.append(score)132                else:133                    print(f"Invalid reference number: {number}")134            else:135                print("No sources available")136 137        return contents, metadata_collection, scores138    139    # def _process_sources_images(self, sources, number_reference_sorted):140    #     contents, metadata_collection, scores = [], [], []141    #     if not number_reference_sorted:142    #         print("There are no references")143    #         return contents, metadata_collection, scores144 145    #     for number in range (len(sources)):146    #         number = int(number)147    #         if sources and len(sources) > 0:148    #             node = dict(sources[0])["raw_output"].source_nodes149    #             if 0 <= number - 1 < len(node):150    #                 content = node[number - 1].node.get_text()151    #                 contents.append(content)152    #                 metadata = dict(node[number - 1].node.metadata)153    #                 metadata_collection.append(metadata)154    #                 score = node[number - 1].score155    #                 scores.append(score)156    #             else:157    #                 print(f"Invalid reference number: {number}")158    #         else:159    #             print("No sources available")160 161    #     return contents, metadata_collection, scores162 163    def _attach_contents_to_metadata(self, contents, metadata_collection):164        for i in range(min(len(contents), len(metadata_collection))):165            metadata_collection[i]["content"] = re.sub(r"source \d+:", "", contents[i])166        return metadata_collection167 168    def _store_message_in_chatstore(self, response, metadata_collection):169        message = ChatMessage(170            role=MessageRole.ASSISTANT, content=response, metadata=metadata_collection171        )172        self.chatstore.delete_last_message(self.session_id)173        self.chatstore.add_message(self.session_id, message)174        self.chatstore.clean_message(self.session_id)175 176    def _save_chat_history_to_db(self, response, metadata_collection):177        chat_history_db = [178            ChatMessage(179                role=MessageRole.USER,180                content=self.user_request,181                timestamp=datetime.now(),182                payment="free" if self.type_bot == "general" else None,183            ),184            ChatMessage(185                role=MessageRole.ASSISTANT,186                content=response,187                metadata=metadata_collection,188                timestamp=datetime.now(),189                payment="free" if self.type_bot == "general" else None,190            ),191        ]192 193        chat_history_json = [message.model_dump() for message in chat_history_db]194 195        db = self.client["bot_database"]  # Replace with your database name196        collection = db[self.session_id]  # Replace with your collection name197        collection.insert_many(chat_history_json)198