Multimedika/Bot_Development

1from typing import List2from llama_index.core.vector_stores import (3    MetadataFilter,4    MetadataFilters,5)6 7from llama_index.core.tools import QueryEngineTool, ToolMetadata8from llama_index.agent.openai import OpenAIAgent9from llama_index.llms.openai import OpenAI10from llama_index.core.query_engine import CitationQueryEngine11from llama_index.embeddings.openai import OpenAIEmbedding12from llama_index.multi_modal_llms.openai import OpenAIMultiModal13from llama_index.core import Settings14from core.chat.chatstore import ChatStore15from core.multimodal import MultimodalQueryEngine16 17from config import GPTBOT_CONFIG18from core.prompt import SYSTEM_BOT_TEMPLATE, ADDITIONAL_INFORMATIONS,SYSTEM_BOT_GENERAL_TEMPLATE, SYSTEM_BOT_IMAGE_TEMPLATE19from core.parser import join_list20 21 22class Engine:23    def __init__(self):24        self.llm = OpenAI(25            temperature=GPTBOT_CONFIG.temperature,26            model=GPTBOT_CONFIG.model,27            max_tokens=GPTBOT_CONFIG.max_tokens,28            api_key=GPTBOT_CONFIG.api_key,29        )30 31        self.chat_store = ChatStore()32        Settings.llm = self.llm33        embed_model = OpenAIEmbedding(model="text-embedding-3-large")34        Settings.embed_model = embed_model35 36    def get_citation_engine(self, titles:List, index):37        model_multimodal = OpenAIMultiModal(model="gpt-4o-mini", max_new_tokens=4096)38        filters = [39            MetadataFilter(40                key="title",41                value=title,42                operator="==",43            )44            for title in titles45        ]46 47        filters = MetadataFilters(filters=filters, condition="or")48 49        # Create the QueryEngineTool with the index and filters50        kwargs = {"similarity_top_k": 10, "filters": filters}51 52        retriever = index.as_retriever(**kwargs)53 54        # citation_engine = CitationQueryEngine(retriever=retriever)55 56        # return CitationQueryEngine.from_args(index, retriever=retriever)57        return MultimodalQueryEngine(retriever=retriever, multi_modal_llm=model_multimodal)58    59 60    def get_chat_engine(self, session_id, index, titles=None, type_bot="general"):61        # Create the QueryEngineTool based on the type62        if type_bot == "general":63            # query_engine = index.as_query_engine(similarity_top_k=3)64            # citation_engine = CitationQueryEngine.from_args(index, similarity_top_k=5)65            model_multimodal = OpenAIMultiModal(model="gpt-4o-mini", max_new_tokens=4096)66            retriever = index.as_retriever(similarity_top_k=10)67            citation_engine = MultimodalQueryEngine(retriever=retriever, multi_modal_llm=model_multimodal)68            # description = "A book containing information about medicine"69        else:70            citation_engine = self.get_citation_engine(titles, index)71            # description = "A book containing information about medicine"72 73        # metadata = ToolMetadata(name="bot-belajar", description=description)74 75        # vector_query_engine = QueryEngineTool(76        #     query_engine=citation_engine, metadata=metadata77        # )78        79        vector_tool = QueryEngineTool.from_defaults(80            query_engine=citation_engine,81            name="vector_tool",82            description=(83                "Useful for retrieving specific context from the data from a book containing information about medicine"84            ),85        )86        87 88        # Initialize the OpenAI agent with the tools89        90        # if type_bot == "general":91        #     system_prompt = SYSTEM_BOT_GENERAL_TEMPLATE92        # else:93        #     additional_information = ADDITIONAL_INFORMATIONS.format(titles=join_list(titles))94        #     system_prompt = SYSTEM_BOT_TEMPLATE.format(additional_information=additional_information)95        # chat_engine = OpenAIAgent.from_tools(96        #     tools=[vector_query_engine],97        #     llm=self.llm,98        #     memory=self.chat_store.initialize_memory_bot(session_id),99        #     system_prompt=system_prompt,100        # )101        102        if type_bot == "general":103            system_prompt = SYSTEM_BOT_IMAGE_TEMPLATE104        else:105            additional_information = ADDITIONAL_INFORMATIONS.format(titles=join_list(titles))106            system_prompt = SYSTEM_BOT_IMAGE_TEMPLATE.format(additional_information=additional_information)107            108        chat_engine = OpenAIAgent.from_tools(109            tools=[vector_tool],110            llm=self.llm,111            memory=self.chat_store.initialize_memory_bot(session_id),112            system_prompt=system_prompt,113        )114 115        return chat_engine116