core/chat/engine.py

from typing import List
from llama_index.core.vector_stores import (
    MetadataFilter,
    MetadataFilters,
)

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.agent.openai import OpenAIAgent
from llama_index.llms.openai import OpenAI
from llama_index.core.query_engine import CitationQueryEngine
from llama_index.core import Settings
from core.chat.chatstore import ChatStore

from config import GPTBOT_CONFIG
from core.prompt import SYSTEM_BOT_TEMPLATE, ADDITIONAL_INFORMATIONS
from core.parser import join_list


class Engine:
    def __init__(self):
        self.llm = OpenAI(
            temperature=GPTBOT_CONFIG.temperature,
            model=GPTBOT_CONFIG.model,
            max_tokens=GPTBOT_CONFIG.max_tokens,
            api_key=GPTBOT_CONFIG.api_key,
        )

        self.chat_store = ChatStore()
        Settings.llm = self.llm

    def get_citation_engine(self, titles:List, index):
        filters = [
            MetadataFilter(
                key="title",
                value=title,
                operator="==",
            )
            for title in titles
        ]

        filters = MetadataFilters(filters=filters, condition="or")

        # Create the QueryEngineTool with the index and filters
        kwargs = {"similarity_top_k": 5, "filters": filters}

        retriever = index.as_retriever(**kwargs)

        # citation_engine = CitationQueryEngine(retriever=retriever)

        return CitationQueryEngine.from_args(index, retriever=retriever)

    def get_chat_engine(self, session_id, index, titles=None, type_bot="general"):
        # Create the QueryEngineTool based on the type
        if type_bot == "general":
            # query_engine = index.as_query_engine(similarity_top_k=3)
            citation_engine = CitationQueryEngine.from_args(index, similarity_top_k=5)
            description = "A book containing information about medicine"
        else:
            citation_engine = self.get_citation_engine(titles, index)
            description = "A book containing information about medicine"

        metadata = ToolMetadata(name="bot-belajar", description=description)
        print(metadata)

        vector_query_engine = QueryEngineTool(
            query_engine=citation_engine, metadata=metadata
        )
        print(vector_query_engine)

        # Initialize the OpenAI agent with the tools
        
        if type_bot == "general":
            system_prompt = SYSTEM_BOT_TEMPLATE.format(additional_information="")
        else:
            additional_information = ADDITIONAL_INFORMATIONS.format(titles=join_list(titles))
            system_prompt = SYSTEM_BOT_TEMPLATE.format(additional_information=additional_information)
        chat_engine = OpenAIAgent.from_tools(
            tools=[vector_query_engine],
            llm=self.llm,
            memory=self.chat_store.initialize_memory_bot(session_id),
            system_prompt=system_prompt,
        )

        return chat_engine