Using Google API

app.py

@@ -2,37 +2,63 @@ import os
 
 import gradio as gr
 from openai import OpenAI
+import google.generativeai as genai
 
 
 LOAD_KEYS_FROM_FILES = False
+USE_LOCAL_ASR_PIPELINE = True
+
 
 if LOAD_KEYS_FROM_FILES:
     # Load OpenAI API key
     with open('KEY_OPENAI', 'r') as file:
-        os.environ['OPENAI_API_KEY'] = file.read().replace('\n', '')
+        OPENAI_API_KEY = file.read().replace('\n', '')
 
     # Hugging Face API key, used for the serverless access to ASR model
     with open('KEY_HF', 'r') as file:  
         os.environ['HUGGINGFACE_API_KEY'] = file.read().replace('\n', '')
 
+    # Google AI API key, to acces Gemini model
+    with open('KEY_GOOGLE_AI', 'r') as file:
+        GOOGLE_API_KEY = file.read().replace('\n', '')
+else:
+    # testar se existe a variável OPENAI_API_KEY
+    OPENAI_API_KEY = "" if 'OPENAI_API_KEY' not in os.environ else os.environ['OPENAI_API_KEY']
+    GOOGLE_API_KEY = "" if 'GOOGLE_API_KEY' not in os.environ else os.environ['GOOGLE_API_KEY']
+    assert 'HUUGINGFACE_API_KEY' in os.environ, "Hugging Face API key not found in environment variables"
+
+
+USE_OPENAI_LLM = (GOOGLE_API_KEY == "")
+
+OPENAI_CLIENT = None
+if USE_OPENAI_LLM:
+    OPENAI_CLIENT = OpenAI(api_key=OPENAI_API_KEY)
+else:
+    genai.configure(api_key=GOOGLE_API_KEY)
+
+GOOGLE_GEN_CONFIG = genai.types.GenerationConfig(
+            candidate_count=1,
+            temperature=0.5)
 
 AUDIO_OUT_FILE_PREFIX = "output"   # prefixo do nome do arquivo de áudio .wav
 
-TEMPLATE_SYSTEM_MESSAGE = """Sua função é entreter uma criança com idade entre 6 e 8 anos que adora futebol. Diretrizes para a conversa: 
-- Seu nome é {NAME}
+TEMPLATE_SYSTEM_MESSAGE = """Você é assistente virtual com a função é entreter uma criança de idade entre 6 e 8 anos que adora futebol. Diretrizes para a conversa: 
+- Você é {GENRE}, seu nome é {NAME}.
 - Pergunte o nome da criança. 
 - Fale sobre futebol, times, jogadores, seleções e grandes jogos. 
 - Em especial, foque no futebol do Brasil, Inglaterra e Espanha. 
 - {PERSONALITY}
-- Fale, no máximo, três frases por mensagem. 
+- Fale, no máximo, três frases por mensagem.
 """
 
 # Mapeia a personalidade no template e na temperatura
 PERSONALITIES = {
-    "nova": ("Seja amigável e alegre, como uma criança. Tente iniciar novos assuntos, quando a conversa estiver repetitiva. Conte piadas de futebol, de vez em quando.", 0.8),
-    "echo": ("Seja amigável e objetivo. Tente manter-se no mesmo assunto. Conte alguma curiosidade sobre um grande craque, de vez em quando.", 0.2)
+    "nova": ("Sua personalidade é bastante amigável e alegre, como uma criança. Tente iniciar novos assuntos, quando a conversa estiver repetitiva. Conte piadas de futebol, de vez em quando.", 0.5, "F"),
+    "echo": ("Sua personalidade é amigável, mas objetivo. Tente manter-se no mesmo assunto. Conte alguma curiosidade sobre um grande craque, de vez em quando.", 0.2, "M")
 }
 
+INITIAL_PERSON = "nova"
+
 
 # Função para converter o histórico de chat para o formato esperado pela API do OpenAI
 def to_openai_chat_history(system_prompt, curr_message, chat_history):
@@ -41,21 +67,46 @@ def to_openai_chat_history(system_prompt, curr_message, chat_history):
         chat_history = chat_history[0:3] + chat_history[-5:]
     for turn in chat_history:
         user_message, bot_message = turn
-        prompt.append( {'role': 'user', 'content' : user_message} )
-        prompt.append( {'role': 'assistant', 'content' : bot_message} )
-    prompt.append( {'role': 'user', 'content' : curr_message } )
+        prompt.append( {'role': 'user', 'content': user_message} )
+        prompt.append( {'role': 'assistant', 'content': bot_message} )
+    prompt.append( {'role': 'user', 'content': curr_message } )
+    return prompt
+
+
+# Função para converter o histórico de chat para o formato esperado pela API do Google AI
+def to_google_history(curr_message, chat_history):
+    prompt = []
+    for turn in chat_history:
+        user_message, bot_message = turn
+        prompt.append( {'role':'user', 'parts': [user_message]} )
+        prompt.append( {'role': 'model', 'parts': [bot_message]} )
+    prompt.append( {'role': 'user', 'parts': [curr_message]} )
+    
     return prompt
 
 
 def respond(system_prompt, user_message, chat_history, temperature, voice="echo"):
-        openai_history = to_openai_chat_history(system_prompt, user_message, chat_history)
+        if USE_OPENAI_LLM:
+            openai_history = to_openai_chat_history(system_prompt, user_message, chat_history)
 
-        bot_response = OPENAI_CLIENT.chat.completions.create(
-            messages=openai_history,
-            temperature=temperature,
-            model="gpt-3.5-turbo")
+            bot_response = OPENAI_CLIENT.chat.completions.create(
+                messages=openai_history,
+                temperature=temperature,
+                model="gpt-3.5-turbo")
 
-        assistant_msg = bot_response.choices[0].message.content
+            assistant_msg = bot_response.choices[0].message.content
+        
+        else:
+            GOOGLE_GEN_CONFIG.temperature = temperature
+            model = genai.GenerativeModel('gemini-1.5-pro-latest', 
+                               system_instruction=system_prompt,
+                               generation_config=GOOGLE_GEN_CONFIG)
+            
+            google_history = to_google_history(user_message, chat_history)
+            
+            bot_response = model.generate_content(google_history, generation_config=GOOGLE_GEN_CONFIG)
+
+            assistant_msg = bot_response.text
 
         # salva o audio 
         response = OPENAI_CLIENT.audio.speech.create(
@@ -74,21 +125,25 @@ def respond(system_prompt, user_message, chat_history, temperature, voice="echo"
         return "", chat_history, output_audio_file
 
 
-def reset_and_apply(openai_key, voice):
-    global OPENAI_CLIENT
-    OPENAI_CLIENT = OpenAI(api_key=openai_key)
+def reset_and_apply(voice):
     return [("", "Olá, vamos falar de futebol?")], AUDIO_OUT_FILE_PREFIX + f"-001-{voice}.wav"
 
 def reset_openai_client(openai_key):
-    global OPENAI_CLIENT
-    OPENAI_CLIENT = OpenAI(api_key=openai_key)
+    global USE_OPENAI_LLM, OPENAI_CLIENT, OPENAI_API_KEY, GOOGLE_API_KEY
+    USE_OPENAI_LLM = (GOOGLE_API_KEY != "")
+    OPENAI_API_KEY = openai_key
+    OPENAI_CLIENT = OpenAI(api_key=OPENAI_API_KEY)
 
-def on_voice_change(voice):
-    persona_description, persona_temperature = PERSONALITIES[voice]
-    return TEMPLATE_SYSTEM_MESSAGE.format(NAME=voice.upper(), PERSONALITY=persona_description), persona_temperature
+def reset_google_client(google_key):
+    global GOOGLE_API_KEY, USE_OPENAI_LLM
+    GOOGLE_API_KEY = google_key
+    USE_OPENAI_LLM = (google_key != "")
 
+def on_voice_change(voice):
+    persona_description, persona_temperature, sex = PERSONALITIES[voice]
+    genre = "menina" if sex=="F" else "menino"
+    return TEMPLATE_SYSTEM_MESSAGE.format(NAME=voice.upper(), PERSONALITY=persona_description, GENRE=genre), persona_temperature
 
-USE_LOCAL_ASR_PIPELINE = True
 
 # With pipeline (downloaded model)
 if USE_LOCAL_ASR_PIPELINE:
@@ -111,19 +166,17 @@ else:
 
 
 def transcribe(audio_in):
-    if USE_LOCAL_ASR_PIPELINE:
-        # sampling rate and audio data
-        #sr, y = audio_in
-        #y2 = y.astype(np.float32)
-        #y2 /= np.max(np.abs(y))
-        #response = ASR_PIPELINE({"sampling_rate": sr, "raw": y})
-        response = ASR_PIPELINE(audio_in)
-    else:
-        # Option2: serverless API
-        with open(audio_in, "rb") as f:
-            data = f.read()
-        response = requests.post(ASR_API_URL, headers=ASR_API_HEADERS, data=data)
-        response = response.json()
+    # from numpy data:
+    #sr, y = audio_in  # sampling rate and audio data
+    #y2 = y.astype(np.float32)
+    #y2 /= np.max(np.abs(y))
+    #response = ASR_PIPELINE({"sampling_rate": sr, "raw": y})
+
+    # using serverless API
+    with open(audio_in, "rb") as f:
+        data = f.read()
+    response = requests.post(ASR_API_URL, headers=ASR_API_HEADERS, data=data)
+    response = response.json()
     
     #print(response)
     return response['text']
@@ -135,21 +188,17 @@ def transcribe_and_respond(audio_in, system_txtbox, user_msg_txb, *args):
     return outputs
 
 
-OPENAI_CLIENT = None
-INITIAL_VOICE = "nova"
-
 
 with gr.Blocks() as demo:
     # aqui, é resetado e instanciado o cliente
-    initial_chat_history, initial_audio = reset_and_apply(os.environ['OPENAI_API_KEY'], INITIAL_VOICE)
+    initial_chat_history, initial_audio = reset_and_apply(INITIAL_PERSON)
 
     chatbot_area = gr.Chatbot(value=initial_chat_history)
     audio_out = gr.Audio(label="Escute a última mensagem", value=initial_audio, autoplay=True, interactive=False)
     
     user_msg_txb = gr.Textbox(label="Mensagem")
 
-    #audio_in = gr.Audio(label="Mensagem de Áudio", sources=['microphone'], interactive=True, type='numpy')
-    audio_in = gr.Audio(label="Mensagem de Áudio", sources=['microphone'], interactive=True, type='filepath')
+    #audio_in = gr.Audio(label="Mensagem de Áudio", sources=['microphone'], interactive=True, type='filepath')  # TODO: tentar type='numpy'
     
     submit_btn = gr.Button("Enviar")
     
@@ -157,16 +206,16 @@ with gr.Blocks() as demo:
     reset_btn = gr.Button("Reiniciar")
     
     with gr.Accordion(label="Configurações",open=False):
-        if LOAD_KEYS_FROM_FILES:
-            openai_key = gr.Textbox(label="OPENAI API Key", value=os.environ['OPENAI_API_KEY'])
-        else:
-            openai_key = gr.Textbox(label="OPENAI API Key", placeholder="Insert your API key here")
+        openai_key = gr.Textbox(label="OpenAI API Key (GPT 3.5)", value="", placeholder="Insira a chave aqui")
         openai_key.change(reset_openai_client, inputs=[openai_key])
 
+        openai_key = gr.Textbox(label="Google API Key (Gemini 1.5)", value="", placeholder="Insira a chave aqui")
+        openai_key.change(reset_google_client, inputs=[openai_key])
+
         # opções de vozes e personalidades
-        voice_ddown = gr.Dropdown(label="Personalidade (muda os dois abaixo)", choices=["nova", "echo"], value=INITIAL_VOICE)
+        voice_ddown = gr.Dropdown(label="Personalidade (muda os dois abaixo)", choices=["nova", "echo"], value=INITIAL_PERSON)
 
-        initial_system_message, initial_temperature = on_voice_change(INITIAL_VOICE)        
+        initial_system_message, initial_temperature = on_voice_change(INITIAL_PERSON)        
         temperature_sldr = gr.Slider(label="Diversidade de respostas", minimum=0.0, maximum=1.0, value=initial_temperature, step=0.1)
         
         with gr.Accordion(label="Avançado",open=False):
@@ -178,8 +227,7 @@ with gr.Blocks() as demo:
 
     reset_btn.click(reset_and_apply, inputs=[openai_key, voice_ddown], outputs=[chatbot_area, audio_out])
 
-    #audio_in.stop_recording( transcribe, inputs=[audio_in], outputs=[user_msg_txb] )
-    audio_in.stop_recording( transcribe_and_respond, inputs=[audio_in, system_txtbox, user_msg_txb, chatbot_area, temperature_sldr, voice_ddown], outputs=[user_msg_txb, chatbot_area, audio_out] )
+    #audio_in.stop_recording( transcribe_and_respond, inputs=[audio_in, system_txtbox, user_msg_txb, chatbot_area, temperature_sldr, voice_ddown], outputs=[user_msg_txb, chatbot_area, audio_out] )
     submit_btn.click(respond, inputs=[system_txtbox, user_msg_txb, chatbot_area, temperature_sldr, voice_ddown], outputs=[user_msg_txb, chatbot_area, audio_out])    # Click on the button
     user_msg_txb.submit(respond, inputs=[system_txtbox, user_msg_txb, chatbot_area, temperature_sldr, voice_ddown], outputs=[user_msg_txb, chatbot_area, audio_out]) # Press enter to submit - same effect