Update app.py

app.py

@@ -16,9 +16,24 @@ CSV_FILE_PATH = "anomalia_vendas.csv"
 SQL_DB_PATH = "data.db"
 HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
-LLAMA_MODEL = "meta-llama/Llama-3.3-70B-Instruct"
 
-hf_client = InferenceClient(api_key=HUGGINGFACE_API_KEY)
+LLAMA_MODELS = {
+    "LLaMA 70B": "meta-llama/Llama-3.3-70B-Instruct",
+    "LlaMA 8B": "meta-llama/Llama-3.1-8B-Instruct",
+    "Qwen 32B": "Qwen/QwQ-32B"
+}
+
+MAX_TOKENS_MAP = {
+    "meta-llama/Llama-3.3-70B-Instruct": 900,
+    "meta-llama/Llama-3.1-8B-Instruct": 600,
+    "Qwen/QwQ-32B": 8192
+}
+
+hf_client = InferenceClient(
+    provider="sambanova",
+    api_key=HF_TOKEN,
+)
+
 os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
 
 query_cache = {}
@@ -50,7 +65,6 @@ sql_agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True
 def generate_initial_context(db_sample):
     return (
         f"Você é um assistente que gera queries SQL objetivas e eficientes. Sempre inclua LIMIT 15 nas queries. Aqui está o banco de dados:\n\n"
-        ###f"Colunas: {', '.join(db_sample.columns)}\n"
         f"Exemplos do banco de dados:\n{db_sample.head().to_string(index=False)}\n\n"
         "\n***IMPORTANTE***: Detecte automaticamente o idioma da pergunta do usuário e responda sempre no mesmo idioma.\n"
         "Essa base de dados representa o sellout de 2025, janeiro, fevereiro e março até dia 11, de uma farmácia.\n"
@@ -65,7 +79,10 @@ def is_greeting(user_query):
     greetings = ["olá", "oi", "bom dia", "boa tarde", "boa noite", "oi, tudo bem?"]
     return user_query.lower().strip() in greetings
 
-def query_with_llama(user_query, db_sample):
+def query_with_llama(user_query, db_sample, selected_model_name):
+    model_id = LLAMA_MODELS[selected_model_name]
+    max_tokens = MAX_TOKENS_MAP.get(model_id, 512)
+    
     initial_context = generate_initial_context(db_sample)
     formatted_history = "\n".join(
         [f"{msg['role'].capitalize()}: {msg['content']}" for msg in recent_history[-2:]]
@@ -73,25 +90,28 @@ def query_with_llama(user_query, db_sample):
     
     full_prompt = f"{initial_context}\n\nHistórico recente:\n{formatted_history}\n\nPergunta do usuário:\n{user_query}"
     
-    logging.info(f"[DEBUG] Contexto enviado ao Llama:\n{full_prompt}\n")
+    logging.info(f"[DEBUG] Contexto enviado ao ({selected_model_name}):\n{full_prompt}\n")
     
     start_time = time.time()
+    
     try:
         response = hf_client.chat.completions.create(
-            model=LLAMA_MODEL,
+            model=model_id,
             messages=[{"role": "system", "content": full_prompt}],
-            max_tokens=500,
+            max_tokens=max_tokens,
             stream=False
         )
+        
         llama_response = response["choices"][0]["message"]["content"]
         end_time = time.time()
-        logging.info(f"[DEBUG] Resposta do Llama para o Agent SQL:\n{llama_response.strip()}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")
-        return llama_response.strip()
+        logging.info(f"[DEBUG] Resposta do {selected_model_name} para o Agent SQL:\n{llama_response.strip()}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")
+        return llama_response.strip(), model_id
+        
     except Exception as e:
-        logging.error(f"[ERRO] Falha ao interagir com o Llama: {e}")
-        return None
+        logging.error(f"[ERRO] Falha ao interagir com o modelo {selected_model_name}: {e}")
+        return None, model_id
 
-def query_sql_agent(user_query):
+def query_sql_agent(user_query, selected_model_name):
     try:
         if user_query in query_cache:
             print(f"[CACHE] Retornando resposta do cache para a consulta: {user_query}")
@@ -103,7 +123,8 @@ def query_sql_agent(user_query):
             return greeting_response
 
         column_data = pd.read_sql_query("SELECT * FROM anomalia_vendas LIMIT 10", engine)
-        llama_instruction = query_with_llama(user_query, column_data)
+        llama_instruction = query_with_llama(user_query, column_data, selected_model_name)
+        
         if not llama_instruction:
             return "Erro: O modelo Llama não conseguiu gerar uma instrução válida."
 
@@ -117,19 +138,27 @@ def query_sql_agent(user_query):
     except Exception as e:
         return f"Erro ao consultar o agente SQL: {e}"
 
-def chatbot_response(user_input):
+def chatbot_response(user_input, selected_model_name):
     start_time = time.time()
-    response = query_sql_agent(user_input)
+    response = query_sql_agent(user_input, selected_model_name)
     end_time = time.time()
-    
-    history_log.append({"Pergunta": user_input, "Resposta": response, "Tempo de Resposta (s)": round(end_time - start_time, 2)})
+
+    model_id = LLAMA_MODELS[selected_model_name]
+
+    history_log.append({
+        "Modelo LLM": model_id,
+        "Pergunta": user_input,
+        "Resposta": response,
+        "Tempo de Resposta (s)": round(end_time - start_time, 2)
+    })
+
     recent_history.append({"role": "user", "content": user_input})
     recent_history.append({"role": "assistant", "content": response})
-    
+
     if len(recent_history) > 4:
         recent_history.pop(0)
         recent_history.pop(0)
-    
+
     return response
 
 def toggle_history():
@@ -138,24 +167,33 @@ def toggle_history():
     return history_log if show_history_flag else {}
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Anomalia Agent")
-    chatbot = gr.Chatbot(height=600)
-    msg = gr.Textbox(placeholder="Digite sua pergunta aqui...", label=" ", lines=1)
-    
-    def respond(message, chat_history):
-        response = chatbot_response(message)
-        chat_history.append((message, response))
-        return "", chat_history
-    
+    gr.Markdown("# 🧠 Anomalia Agent")
+
     with gr.Row():
-        btn = gr.Button("Enviar", variant="primary")
-        history_btn = gr.Button("Histórico", variant="secondary")
-        
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    btn.click(respond, [msg, chatbot], [msg, chatbot])
-    
-    history_output = gr.JSON()
-    history_btn.click(toggle_history, inputs=[], outputs=history_output)
+        with gr.Column(scale=1):
+            gr.Markdown("## ⚙️ Configurações")
+            model_selector = gr.Dropdown(
+                choices=list(LLAMA_MODELS.keys()),
+                label="Escolha o Modelo LLM para gerar a query SQL",
+                value="LLaMA 70B"
+            )
+
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot(height=600)
+            msg = gr.Textbox(placeholder="Digite sua pergunta aqui...", label=" ", lines=1)
+            btn = gr.Button("Enviar", variant="primary")
+            history_btn = gr.Button("Histórico", variant="secondary")
+
+            def respond(message, chat_history, selected_model_name):
+                response = chatbot_response(message, selected_model_name)
+                chat_history.append((message, response))
+                return "", chat_history
+
+            msg.submit(respond, [msg, chatbot, model_selector], [msg, chatbot])
+            btn.click(respond, [msg, chatbot, model_selector], [msg, chatbot])
+
+            history_output = gr.JSON()
+            history_btn.click(toggle_history, inputs=[], outputs=history_output)
 
 if __name__ == "__main__":
     demo.launch(share=False)