chat/model_manage.py

# # my_app/model_manager.py
# import google.generativeai as genai
# import chat.arxiv_bot.arxiv_bot_utils as utils
# import json

# model = None

# model_retrieval = None

# model_answer = None

# RETRIEVAL_INSTRUCT = """You are an auto chatbot that response with only one action below based on user question.
#    1. If the guest question is asking about a science topic, you need to respond the information in JSON schema below:
#         {
#             "keywords": [a list of string keywords about the topic],
#             "description": "a paragraph describing the topic in about 50 to 100 words"
#         }
#    2. If the guest is not asking for any informations or documents, you need to respond in JSON schema below:
#         {
#             "answer": "your answer to the user question"
#         }"""

# ANSWER_INSTRUCT = """You are a library assistant that help answering customer question based on the information given.
#         You always answer in a conversational form naturally and politely.
#         You must introduce all the records given, each must contain title, authors and the link to the pdf file."""

# def create_model():
#     with open("apikey.txt","r") as apikey:
#         key = apikey.readline()
#         genai.configure(api_key=key)
#     for m in genai.list_models():
#         if 'generateContent' in m.supported_generation_methods:
#             print(m.name)
#     print("He was there")
#     config = genai.GenerationConfig(max_output_tokens=2048,
#                                 temperature=1.0)
#     safety_settings = [
#         {
#             "category": "HARM_CATEGORY_DANGEROUS",
#             "threshold": "BLOCK_NONE",
#         },
#         {
#             "category": "HARM_CATEGORY_HARASSMENT",
#             "threshold": "BLOCK_NONE",
#         },
#         {
#             "category": "HARM_CATEGORY_HATE_SPEECH",
#             "threshold": "BLOCK_NONE",
#         },
#         {
#             "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
#             "threshold": "BLOCK_NONE",
#         },
#         {
#             "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
#             "threshold": "BLOCK_NONE",
#         },
#     ]
#     global model, model_retrieval, model_answer
#     model = genai.GenerativeModel("gemini-1.5-pro-latest",
#                                 generation_config=config,
#                                 safety_settings=safety_settings)
#     model_retrieval = genai.GenerativeModel("gemini-1.5-pro-latest",
#                                             generation_config=config,
#                                             safety_settings=safety_settings,
#                                             system_instruction=RETRIEVAL_INSTRUCT)
#     model_answer = genai.GenerativeModel("gemini-1.5-pro-latest",
#                                          generation_config=config,
#                                          safety_settings=safety_settings,
#                                          system_instruction=ANSWER_INSTRUCT)
#     return model, model_answer, model_retrieval

# def get_model():
#     global model, model_answer, model_retrieval
#     if model is None:
#         # Khởi tạo model ở đây
#         model, model_answer, model_retrieval = create_model()  # Giả sử create_model là hàm tạo model của bạn
#     return model, model_answer, model_retrieval

# def extract_keyword_prompt(query):
#     """A prompt that return a JSON block as arguments for querying database"""

#     prompt = """[INST] SYSTEM: You are an auto chatbot that response with only one action below based on user question.
#    1. If the guest question is asking about a science topic, you need to respond the information in JSON schema below:
#         {
#             "keywords": [a list of string keywords about the topic],
#             "description": "a paragraph describing the topic in about 50 to 100 words"
#         }
#    2. If the guest is not asking for any informations or documents, you need to respond in JSON schema below:
#         {
#             "answer": "your answer to the user question"
#         }
#    QUESTION: """ + query + """[/INST]
#    ANSWER: """
#     return prompt

# def make_answer_prompt(input, contexts):
#     """A prompt that return the final answer, based on the queried context"""

#     prompt = (
#         """[INST] You are a library assistant that help answering customer QUESTION based on the INFORMATION given.
#         You always answer in a conversational form naturally and politely.
#         You must introduce all the records given, each must contain title, authors and the link to the pdf file.
#   QUESTION: {input}
#   INFORMATION: '{contexts}'
#   [/INST]
#   ANSWER:
#   """
#     ).format(input=input, contexts=contexts)
#     return prompt

# def retrieval_chat_template(question):
#     return {
#         "role":"user",
#         "parts":[f"QUESTION: {question} \n ANSWER:"]
#     }

# def answer_chat_template(question, contexts):
#     return {
#         "role":"user",
#         "parts":[f"QUESTION: {question} \n INFORMATION: {contexts} \n ANSWER:"]
#     }

# def response(args, db_instance):
#     """Create response context, based on input arguments"""
#     keys = list(dict.keys(args))
#     if "answer" in keys:
#         return args['answer'], None  # trả lời trực tiếp
     
#     if "keywords" in keys:
#         # perform query
#         query_texts = args["description"]
#         keywords = args["keywords"]
#         results = utils.db.query_relevant(keywords=keywords, query_texts=query_texts)
#         # print(results)
#         ids = results['metadatas'][0]
#         if len(ids) == 0:
#             # go crawl some
#             new_records = utils.crawl_arxiv(keyword_list=keywords, max_results=10)
#             print("Got new records: ",len(new_records))
#             if type(new_records) == str:
#                 return "Error occured, information not found", new_records
#             utils.db.add(new_records)
#             db_instance.add(new_records)
#             results = utils.db.query_relevant(keywords=keywords, query_texts=query_texts)
#             ids = results['metadatas'][0]
#             print("Re-queried on chromadb, results: ",ids)
#         paper_id = [id['paper_id'] for id in ids]
#         paper_info = db_instance.query_id(paper_id)
#         print(paper_info)
#         records = [] # get title (2), author (3), link (6)
#         result_string = ""
#         if paper_info:
#             for i in range(len(paper_info)):
#                 result_string += "Record no.{} - Title: {}, Author: {}, Link: {}, ".format(i+1,paper_info[i][2],paper_info[i][3],paper_info[i][6])
#                 id = paper_info[i][0]
#                 selected_document = utils.db.query_exact(id)["documents"]
#                 doc_str = "Summary:"
#                 for doc in selected_document:
#                     doc_str+= doc + " "
#                 result_string += doc_str
#                 records.append([paper_info[i][2],paper_info[i][3],paper_info[i][6]])
#             return result_string, records
#         else:
#             return "Information not found", "Information not found"
#         # invoke llm and return result

#     # if "title" in keys:
#     #     title = args['title']
#     #     authors = utils.authors_str_to_list(args['author'])
#     #     paper_info = db_instance.query(title = title,author = authors)
#     #     # if query not found then go crawl brh
#     #     # print(paper_info)

#     #     if len(paper_info) == 0:
#     #         new_records = utils.crawl_exact_paper(title=title,author=authors)
#     #         print("Got new records: ",len(new_records))
#     #         if type(new_records) == str:
#     #             # print(new_records)
#     #             return "Error occured, information not found", "Information not found"
#     #         utils.db.add(new_records)
#     #         db_instance.add(new_records)
#     #         paper_info = db_instance.query(title = title,author = authors)
#     #         print("Re-queried on chromadb, results: ",paper_info)
#     #     # -------------------------------------
#     #     records = [] # get title (2), author (3), link (6)
#     #     result_string = ""
#     #     for i in range(len(paper_info)):
#     #         result_string += "Title: {}, Author: {}, Link: {}".format(paper_info[i][2],paper_info[i][3],paper_info[i][6])
#     #         records.append([paper_info[i][2],paper_info[i][3],paper_info[i][6]])
#     #     # process results:
#     #     if len(result_string) == 0:
#     #         return "Information not found", "Information not found"
#     #     return result_string, records
#         # invoke llm and return result

# def full_chain_single_question(input_prompt, db_instance):
#     try:
#         first_prompt = extract_keyword_prompt(input_prompt)
#         temp_answer = model.generate_content(first_prompt).text

#         args = json.loads(utils.trimming(temp_answer))
#         contexts, results = response(args, db_instance)
#         if not results:
#             # print(contexts)
#             return "Random question, direct return", contexts
#         else:
#             output_prompt = make_answer_prompt(input_prompt,contexts)
#             answer = model.generate_content(output_prompt).text
#             return temp_answer, answer
#     except Exception as e:
#         # print(e)
#         return temp_answer, "Error occured: " + str(e)
    

# def format_chat_history_from_web(chat_history: list):
#     temp_chat = []
#     for message in chat_history:
#         temp_chat.append(
#             {
#                 "role": message["role"],
#                 "parts": [message["content"]]
#             }
#         )
#     return temp_chat

# # def full_chain_history_question(chat_history: list, db_instance):
# #     try:
# #         temp_chat = format_chat_history_from_web(chat_history)
# #         print('Extracted temp chat: ',temp_chat)
# #         first_prompt = extract_keyword_prompt(temp_chat[-1]["parts"][0])
# #         temp_answer = model.generate_content(first_prompt).text

# #         args = json.loads(utils.trimming(temp_answer))
# #         contexts, results = response(args, db_instance)
# #         print('Context extracted: ',contexts)
# #         if not results:
# #             return "Random question, direct return", contexts
# #         else:
# #             QA_Prompt = make_answer_prompt(temp_chat[-1]["parts"][0], contexts)
# #             temp_chat[-1]["parts"] = QA_Prompt
# #             print(temp_chat)
# #             answer = model.generate_content(temp_chat).text
# #             return temp_answer, answer
# #     except Exception as e:
# #         # print(e)
# #         return temp_answer, "Error occured: " + str(e)

# def full_chain_history_question(chat_history: list, db_instance):
#     try:
#         temp_chat = format_chat_history_from_web(chat_history)
#         question = temp_chat[-1]['parts'][0]
#         first_answer = model_retrieval.generate_content(temp_chat).text
        
#         print(first_answer)
#         args = json.loads(utils.trimming(first_answer))

#         contexts, results = response(args, db_instance)
#         if not results:
#             return "Random question, direct return", contexts
#         else:
#             print('Context to answers: ',contexts)
#             answer_chat = answer_chat_template(question, contexts)
#             temp_chat[-1] = answer_chat
#             answer = model_answer.generate_content(temp_chat).text
#             return first_answer, answer
#     except Exception as e:
#         if first_answer:
#             return first_answer, "Error occured: " + str(e)
#         else:
#             return "No answer", "Error occured: " + str(e)