Update backend/query_llm.py

backend/query_llm.py

@@ -8,6 +8,12 @@ from typing import Any, Dict, Generator, List
 
 from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer
+import google.generativeai as genai
+import os
+import PIL.Image
+import gradio as gr
+#from gradio_multimodalchatbot import MultimodalChatbot
+from gradio.data_classes import FileData
 
 #tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
 tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
@@ -18,6 +24,10 @@ temperature = 0.5
 top_p = 0.7
 repetition_penalty = 1.2
 
+
+# Fetch an environment variable.
+GOOGLE_API_KEY = os.environ.get('GOOGLE_API_KEY')
+genai.configure(api_key=GOOGLE_API_KEY)
 OPENAI_KEY = getenv("OPENAI_API_KEY")
 HF_TOKEN = getenv("HUGGING_FACE_HUB_TOKEN")
 
@@ -160,3 +170,115 @@ def generate_openai(prompt: str, history: str, temperature: float = 0.9, max_new
             print("Unhandled Exception:", str(e))
             gr.Warning("Unfortunately OpenAI is unable to process")
             return "I do not know what happened, but I couldn't understand you."
+
+
+def generate_gemini(prompt: str, history: str, temperature: float = 0.9, max_new_tokens: int = 4000,
+             top_p: float = 0.95, repetition_penalty: float = 1.0):
+    
+    # For better security practices, retrieve sensitive information like API keys from environment variables.
+ 
+    
+    # Initialize genai models
+    model = genai.GenerativeModel('gemini-pro')
+    api_key = os.environ.get("GOOGEL_API_KEY")
+    genai.configure(api_key=api_key)
+    #model = genai.GenerativeModel('gemini-pro')
+    #chat = model.start_chat(history=[])
+
+
+   
+    candidate_count=1
+    max_output_tokens=max_new_tokens
+    temperature=temperature
+    top_p=top_p
+
+    
+    formatted_prompt = format_prompt(prompt, "gemini")
+
+    try:
+        stream = model.generate_content(formatted_prompt,generation_config=genai.GenerationConfig(temperature=temperature,candidate_count=1 ,max_output_tokens=max_new_tokens,top_p=top_p),
+                                            stream=True)
+        output = ""
+        for response in stream:
+            output += response.text
+            yield output
+
+    except Exception as e:
+        if "Too Many Requests" in str(e):
+            print("ERROR: Too many requests on Mistral client")
+            gr.Warning("Unfortunately Mistral is unable to process")
+            return "Unfortunately, I am not able to process your request now."
+        elif "Authorization header is invalid" in str(e):
+            print("Authetification error:", str(e))
+            gr.Warning("Authentication error: HF token was either not provided or incorrect")
+            return "Authentication error"
+        else:
+            print("Unhandled Exception:", str(e))
+            gr.Warning("Unfortunately Mistral is unable to process")
+            return "I do not know what happened, but I couldn't understand you."
+
+    
+
+    
+    # def gemini(input, file, chatbot=[]):
+    #     """
+    #     Function to handle gemini model and gemini vision model interactions.
+    #     Parameters:
+    #     input (str): The input text.
+    #     file (File): An optional file object for image processing.
+    #     chatbot (list): A list to keep track of chatbot interactions.
+    #     Returns:
+    #     tuple: Updated chatbot interaction list, an empty string, and None.
+    #     """
+    
+    #     messages = []
+    #     print(chatbot)
+    
+    #     # Process previous chatbot messages if present
+    #     if len(chatbot) != 0:
+    #         for messages_dict in chatbot:
+    #             user_text = messages_dict[0]['text']
+    #             bot_text = messages_dict[1]['text']
+    #             messages.extend([
+    #                 {'role': 'user', 'parts': [user_text]},
+    #                 {'role': 'model', 'parts': [bot_text]}
+    #             ])
+    #         messages.append({'role': 'user', 'parts': [input]})
+    #     else:
+    #         messages.append({'role': 'user', 'parts': [input]})
+    
+    #     try:
+    #         response = model.generate_content(messages)
+    #         gemini_resp = response.text
+    #         # Construct list of messages in the required format
+    #         user_msg = {"text": input, "files": []}
+    #         bot_msg = {"text": gemini_resp, "files": []}
+    #         chatbot.append([user_msg, bot_msg])
+       
+    #     except Exception as e:
+    #         # Handling exceptions and raising error to the modal
+    #         print(f"An error occurred: {e}")
+    #         raise gr.Error(e)
+    
+    #     return chatbot, "", None
+    
+    # # Define the Gradio Blocks interface
+    # with gr.Blocks() as demo:
+    #     # Add a centered header using HTML
+    #     gr.HTML("<center><h1>Gemini Chat PRO API</h1></center>")
+    
+    #     # Initialize the MultimodalChatbot component
+    #     multi = MultimodalChatbot(value=[], height=800)
+    
+    #     with gr.Row():
+    #         # Textbox for user input with increased scale for better visibility
+    #         tb = gr.Textbox(scale=4, placeholder='Input text and press Enter')
+    
+           
+    
+    #     # Define the behavior on text submission
+    #     tb.submit(gemini, [tb, multi], [multi, tb])
+    
+    
+    # # Launch the demo with a queue to handle multiple users
+    # demo.queue().launch()