Update app.py

app.py

@@ -1,16 +1,16 @@
 # ---------------------------------------------------------------------------------
 # Aplicación principal para cargar el modelo, generar prompts y explicar los datos
 # ---------------------------------------------------------------------------------
-
 import streamlit as st  # type: ignore
 import os
 import re
 import pandas as pd  # type: ignore
-from dotenv import load_dotenv  # type: ignore # Para cambios locales
+from dotenv import load_dotenv  # type: ignore
+# Para cambios locales
 from supabase import create_client, Client  # type: ignore
-
 # from pandasai import SmartDataframe  # type: ignore
-from pandasai import SmartDatalake  # type: ignore # Porque ya usamos más de un df (más de una tabla de nuestra db)
+from pandasai import SmartDatalake  # type: ignore
+# Porque ya usamos más de un df (más de una tabla de nuestra db)
 from pandasai.llm.local_llm import LocalLLM # type: ignore
 from pandasai import Agent
 import plotly.graph_objects as go
@@ -21,13 +21,10 @@ import time
 # Funciones auxiliares
 # ---------------------------------------------------------------------------------
 
-
 def generate_graph_prompt(user_query):
     prompt = f"""
             You are a senior data scientist analyzing European labor force data.
-
             Given the user's request: "{user_query}"
-
             1. Plot the relevant data using graph_objects plotly:
             - Use `df.query("geo == 'X'")` to filter the country, instead of chained comparisons.
             - Avoid using filters like `df[df['geo'] == 'Germany']`.
@@ -57,11 +54,9 @@ def generate_graph_prompt(user_query):
     return prompt
 
 #TODO: Continuar mejorando el prompt
-
 # ---------------------------------------------------------------------------------
 # Configuración de conexión a Supabase
 # ---------------------------------------------------------------------------------
-
 # Cargar variables de entorno desde archivo .env
 load_dotenv()
 
@@ -94,7 +89,6 @@ def load_data(table):
             else:
                 st.info("Response object does not have 'data' or known error attributes. Check the logs.")
                 return pd.DataFrame()
-
         else:
             st.error("Supabase client not initialized. Check environment variables.")
             return pd.DataFrame()
@@ -112,7 +106,7 @@ def load_data(table):
 labor_data = load_data("labor")
 fertility_data = load_data("fertility")
 # population_data = load_data("population")
-# predictions_data = load_data("predictions") 
+# predictions_data = load_data("predictions")
 
 # TODO: Buscar la forma de disminuir la latencia (muchos datos = mucha latencia)
 
@@ -120,25 +114,24 @@ fertility_data = load_data("fertility")
 # Inicializar LLM desde Ollama con PandasAI
 # ---------------------------------------------------------------------------------
 
-# ollama_llm = LocalLLM(api_base="http://localhost:11434/v1", 
+# ollama_llm = LocalLLM(api_base="http://localhost:11434/v1",
 #                       model="gemma3:12b",
-#                       temperature=0.1,  
+#                       temperature=0.1,
 #                       max_tokens=8000)
 
 lm_studio_llm = LocalLLM(api_base="http://localhost:1234/v1") # el modelo es gemma-3-12b-it-qat
-
 # sdl = SmartDatalake([labor_data, fertility_data, population_data, predictions_data], config={"llm": ollama_llm}) # DataFrame PandasAI-ready.
 # sdl = SmartDatalake([labor_data, fertility_data], config={"llm": ollama_llm})
 
 # agent = Agent([labor_data], config={"llm": lm_studio_llm}) # TODO: Probar Agent con multiples dfs
 agent = Agent(
     [
-        labor_data, 
+        labor_data,
         fertility_data
-        ],
+    ],
     config={
         "llm": lm_studio_llm,
-        "enable_cache": False, 
+        "enable_cache": False,
         "enable_filter_extraction": False  # evita errores de parseo
     }
 )
@@ -147,48 +140,60 @@ agent = Agent(
 # Configuración de la app en Streamlit
 # ---------------------------------------------------------------------------------
 
-# Título de la app
-st.title("Europe GraphGen  :blue[Graph generator] :flag-eu:")
-
-# TODO: Poner instrucciones al usuario sobre cómo hacer un muy buen prompt (sin tecnisismos, pensando en el usuario final)
-
-# Entrada de usuario para describir el gráfico
-user_input = st.text_input("What graphics do you have in mind")
-generate_button = st.button("Generate")
-
-if generate_button and user_input:
-    with st.spinner('Generating answer...'):
-        try:
-            print(f"\nGenerating prompt...\n")
-            prompt = generate_graph_prompt(user_input)
-            print(f"\nPrompt generated\n")
-
-            start_time = time.time()
-
-            answer = agent.chat(prompt)
-            print(f"\nAnswer type: {type(answer)}\n")  # Verificar tipo de objeto
-            print(f"\nAnswer content: {answer}\n")  # Inspeccionar contenido de la respuesta
-            print(f"\nFull result: {agent.last_result}\n")
-
-            full_result = agent.last_result
-            explanation = full_result.get("explanation", "")
-            
-            elapsed_time = time.time() - start_time
-            print(f"\nExecution time: {elapsed_time:.2f} seconds\n")
-
-            if isinstance(answer, str) and os.path.isfile(answer):
-                # Si el output es una ruta válida a imagen
-                im = plt.imread(answer)
-                st.image(im)
-                os.remove(answer)  # Limpiar archivo temporal
-
-                if explanation:
-                    st.markdown(f"**Explanation:** {explanation}")
-            else:
-                # Si no es una ruta válida, mostrar como texto
-                st.markdown(str(answer))            
-
-        except Exception as e:
-            st.error(f"Error generating answer: {e}")
-
-# TODO: Output estructurado si vemos que es necesario.
+st.set_page_config(page_title="GraphGen", page_icon="🇪🇺")
+st.title("_Europe GraphGen_  :blue[Graph generator] :flag-eu:")
+st.caption("Mapping Europe's data: Your tool for custom demographic charts")
+
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+    st.session_state.messages.append({"role": "assistant", "content": "What graphic do you have in mind?"})
+
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+prompt = st.chat_input("Type your message here...", key="chat_input_bottom")
+
+if prompt:
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    with st.chat_message("assistant"):
+        with st.spinner('Generating answer...'):
+            try:
+                print(f"\nGenerating prompt...\n")
+                graph_prompt = generate_graph_prompt(prompt)
+                print(f"\nPrompt generated: {graph_prompt}\n")
+
+                start_time = time.time()
+                answer = agent.chat(graph_prompt)
+                print(f"\nAnswer type: {type(answer)}\n")  # Verificar tipo de objeto
+                print(f"\nAnswer content: {answer}\n")  # Inspeccionar contenido de la respuesta
+                print(f"\nFull result: {agent.last_result}\n")
+
+                full_result = agent.last_result
+                explanation = full_result.get("explanation", "")
+
+                elapsed_time = time.time() - start_time
+                print(f"\nExecution time: {elapsed_time:.2f} seconds\n")
+
+                if isinstance(answer, str) and os.path.isfile(answer):
+                    # Si el output es una ruta válida a imagen
+                    im = plt.imread(answer)
+                    st.image(im)
+                    os.remove(answer)  # Limpiar archivo temporal
+
+                    if explanation:
+                        st.markdown(f"**Explanation:** {explanation}")
+                else:
+                    # Si no es una ruta válida, mostrar como texto
+                    st.markdown(str(answer))
+
+            except Exception as e:
+                st.error(f"Error generating answer: {e}")
+
+if st.button("Clear chat"):
+    st.session_state.messages = []
+    st.session_state.messages.append({"role": "assistant", "content": "Chat has been cleared. What graphic do you have in mind now?"})
+    st.rerun()