Update app.py

app.py

@@ -1,6 +1,8 @@
 import os
 import time
+import json
 import pandas as pd
+import numpy as np
 from sqlalchemy import create_engine
 from langchain_openai import ChatOpenAI
 from langchain_community.agent_toolkits import create_sql_agent
@@ -8,26 +10,26 @@ from langchain_community.utilities import SQLDatabase
 from huggingface_hub import InferenceClient
 import gradio as gr
 from dotenv import load_dotenv
-import logging
 
 load_dotenv()
 
 CSV_FILE_PATH = "tabela_tiago_formated.csv"
 SQL_DB_PATH = "data.db"
-HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
+HF_API_KEY = os.getenv("HF_API_KEY")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 LLAMA_MODEL = "meta-llama/Llama-3.3-70B-Instruct"
+LLAMA_MODEL_FINAL = "Qwen/QwQ-32B"
+
+hf_client = InferenceClient(
+    provider="sambanova",
+    api_key=HF_TOKEN,
+)
 
-hf_client = InferenceClient(api_key=HUGGINGFACE_API_KEY)
 os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
 
 query_cache = {}
-history_log = []  
-recent_history = []  
-show_history_flag = False
-
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
+# Criar banco de dados SQL a partir do JSON
 def create_or_load_sql_database(csv_path, sql_db_path):
     if os.path.exists(sql_db_path):
         print("Banco de dados SQL já existe. Carregando...")
@@ -41,11 +43,11 @@ def create_or_load_sql_database(csv_path, sql_db_path):
         print("Banco de dados SQL criado com sucesso!")
         return engine
 
-engine = create_or_load_sql_database(CSV_FILE_PATH, SQL_DB_PATH)
+engine = create_or_load_sql_database(JSON_FILE_PATH, SQL_DB_PATH)
 db = SQLDatabase(engine=engine)
 
 llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
-sql_agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True, max_iterations=40, return_intermediate_steps=True)
+sql_agent = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True, max_iterations=20, return_intermediate_steps=True)
 
 def generate_initial_context(db_sample):
     return (
@@ -63,97 +65,117 @@ 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, chat_history):
     initial_context = generate_initial_context(db_sample)
-    formatted_history = "\n".join(
-        [f"{msg['role'].capitalize()}: {msg['content']}" for msg in recent_history[-2:]]
-    )
-    
+
+    formatted_history = ""  # Não inclui histórico temporariamente
+
+    #formatted_history = "\n".join(
+        #[f"{msg['role'].capitalize()}: {msg['content']}" for msg in chat_history[-0:]]
+    #)
+
     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")
-    
+
+    print("------- Modelo Llama 3.3: Gerando query SQL -------")
+    print(f"[DEBUG] Contexto enviado ao Llama:\n{full_prompt}\n")
     start_time = time.time()
+
     try:
         response = hf_client.chat.completions.create(
             model=LLAMA_MODEL,
             messages=[{"role": "system", "content": full_prompt}],
-            max_tokens=500,
+            max_tokens=600,
             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()
+        print(f"[DEBUG] Resposta do Llama: {llama_response.strip()}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")
+        chat_history.append({"role": "assistant", "content": llama_response})
+        return llama_response.strip(), chat_history
     except Exception as e:
-        logging.error(f"[ERRO] Falha ao interagir com o Llama: {e}")
-        return None
+        print(f"[ERRO] Falha ao interagir com o Llama: {e}")
+        return None, chat_history
+
+def refine_response_with_llama(user_query, sql_response):
+    prompt = f"""
+    Você é um assistente especializado em consolidar informações. Abaixo estão as entradas:\n
+    Pergunta: 
+    {user_query}
+    Resposta do Agente:
+    {sql_response}\n
+    """
+
+    print("------- Modelo Final: Refinando resposta -------")
+    print(f"[DEBUG] Contexto enviado ao Llama Final:\n{prompt}\n")
+    start_time = time.time()
 
-def query_sql_agent(user_query):
+    try:
+        response = hf_client.chat.completions.create(
+            model=LLAMA_MODEL_FINAL,
+            messages=[{"role": "system", "content": prompt}],
+            max_tokens=8000
+        )
+        
+        final_response = response["choices"][0]["message"]["content"]
+        end_time = time.time()
+        print(f"[DEBUG] Resposta do Llama Final: {final_response.strip()}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")
+        return final_response.strip()
+        
+    except Exception as e:
+        print(f"[ERRO] Falha ao interagir com o Llama Final: {e}")
+        return "Erro ao consolidar a resposta. Por favor, tente novamente."
+
+def query_sql_agent(user_query, chat_history):
     try:
         if user_query in query_cache:
-            print(f"[CACHE] Retornando resposta do cache para a consulta: {user_query}")
-            return query_cache[user_query]
+            print("[CACHE] Retornando resposta do cache para a consulta: {user_query}")
+            return query_cache[user_query], chat_history
 
         if is_greeting(user_query):
             greeting_response = "Olá! Estou aqui para ajudar com suas consultas. Pergunte algo relacionado aos dados carregados no agente!"
-            query_cache[user_query] = greeting_response  # Armazena saudação no cache
-            return greeting_response
+            return greeting_response, chat_history
 
         column_data = pd.read_sql_query("SELECT * FROM tabela_tiago_formated LIMIT 10", engine)
-        llama_instruction = query_with_llama(user_query, column_data)
+
+        llama_instruction, chat_history = query_with_llama(user_query, column_data, chat_history)
         if not llama_instruction:
-            return "Erro: O modelo Llama não conseguiu gerar uma instrução válida."
+            return "Erro: O modelo Llama não conseguiu gerar uma instrução válida.", chat_history
 
         print("------- Agent SQL: Executando query -------")
+        print(f"[DEBUG] Instrução gerada pelo Llama:\n{llama_instruction}\n")
+        start_time = time.time()
         response = sql_agent.invoke({"input": llama_instruction})
-        sql_response = response.get("output", "Erro ao obter a resposta do agente...")
+        sql_response = response.get("output", "Erro ao obter a resposta do agente.")
+        end_time = time.time()
+        print(f"[DEBUG] Resposta do Agent SQL: {sql_response}\n[Tempo de execução: {end_time - start_time:.2f}s]\n")
+
+        final_response = refine_response_with_llama(user_query, sql_response)
+        query_cache[user_query] = final_response  # Armazenar no cache
+        return final_response, chat_history
 
-        query_cache[user_query] = sql_response
-        return sql_response
-        
     except Exception as e:
-        return f"Erro ao consultar o agente SQL: {e}"
+        return f"Erro ao consultar o agente SQL: {e}", chat_history
 
-def chatbot_response(user_input):
-    start_time = time.time()
-    response = query_sql_agent(user_input)
-    end_time = time.time()
-    
-    history_log.append({"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():
-    global show_history_flag
-    show_history_flag = not show_history_flag
-    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
-    
-    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)
+def gradio_interface():
+    chat_history = []
+
+    def respond(user_input):
+        nonlocal chat_history
+        response, chat_history = query_sql_agent(user_input, chat_history)
+        chat_history.append({"role": "user", "content": user_input})
+        chat_history.append({"role": "assistant", "content": response})
+        return [{"role": "user", "content": user_input}, {"role": "assistant", "content": response}]
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# Tributario Agent")
+        chatbot = gr.Chatbot(label="Tributario Result Agent", type="messages")
+        user_input = gr.Textbox(label="Digite sua pergunta")
+        submit_button = gr.Button("Enviar")
+
+        submit_button.click(respond, inputs=user_input, outputs=chatbot)
+
+    return demo
 
 if __name__ == "__main__":
+    demo = gradio_interface()
     demo.launch(share=False)