Spaces:
Sleeping
Sleeping
| 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.embeddings.openai import OpenAIEmbedding | |
| from llama_index.multi_modal_llms.openai import OpenAIMultiModal | |
| from llama_index.core import Settings | |
| from core.chat.chatstore import ChatStore | |
| from core.multimodal import MultimodalQueryEngine | |
| from config import GPTBOT_CONFIG | |
| from core.prompt import SYSTEM_BOT_TEMPLATE, ADDITIONAL_INFORMATIONS,SYSTEM_BOT_GENERAL_TEMPLATE, SYSTEM_BOT_IMAGE_TEMPLATE | |
| 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 | |
| embed_model = OpenAIEmbedding(model="text-embedding-3-large") | |
| Settings.embed_model = embed_model | |
| def get_citation_engine(self, titles:List, index): | |
| model_multimodal = OpenAIMultiModal(model="gpt-4o-mini", max_new_tokens=4096) | |
| 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": 10, "filters": filters} | |
| retriever = index.as_retriever(**kwargs) | |
| # citation_engine = CitationQueryEngine(retriever=retriever) | |
| # return CitationQueryEngine.from_args(index, retriever=retriever) | |
| return MultimodalQueryEngine(retriever=retriever, multi_modal_llm=model_multimodal) | |
| 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) | |
| model_multimodal = OpenAIMultiModal(model="gpt-4o-mini", max_new_tokens=4096) | |
| retriever = index.as_retriever(similarity_top_k=10) | |
| citation_engine = MultimodalQueryEngine(retriever=retriever, multi_modal_llm=model_multimodal) | |
| # 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) | |
| # vector_query_engine = QueryEngineTool( | |
| # query_engine=citation_engine, metadata=metadata | |
| # ) | |
| vector_tool = QueryEngineTool.from_defaults( | |
| query_engine=citation_engine, | |
| name="vector_tool", | |
| description=( | |
| "Useful for retrieving specific context from the data from a book containing information about medicine" | |
| ), | |
| ) | |
| # Initialize the OpenAI agent with the tools | |
| # if type_bot == "general": | |
| # system_prompt = SYSTEM_BOT_GENERAL_TEMPLATE | |
| # 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, | |
| # ) | |
| if type_bot == "general": | |
| system_prompt = SYSTEM_BOT_IMAGE_TEMPLATE | |
| else: | |
| additional_information = ADDITIONAL_INFORMATIONS.format(titles=join_list(titles)) | |
| system_prompt = SYSTEM_BOT_IMAGE_TEMPLATE.format(additional_information=additional_information) | |
| chat_engine = OpenAIAgent.from_tools( | |
| tools=[vector_tool], | |
| llm=self.llm, | |
| memory=self.chat_store.initialize_memory_bot(session_id), | |
| system_prompt=system_prompt, | |
| ) | |
| return chat_engine | |