Update src/chatbox_v1.py

src/chatbox_v1.py

@@ -1,60 +1,60 @@
-import  streamlit as st 
-from Rag_milvus import query_qdrant, obtener_colecciones, query_qdrant_sinumbral
-from Llm_local import get_response_from_mistral, generarPages
-from sentence_transformers import SentenceTransformer
-
-col1, col2 = st.columns([1, 4])
-with col1:
-    st.image("Procuradurialogo.jpg", width=600)
-
-with col2:
-    st.markdown("""
-        <div style='display: flex; align-items: center; height: 100%;'>
-            <h1 style='margin: 0; text-align: center;'>ProcurIA</h1>
-        </div>
-    """, unsafe_allow_html=True)
-
-st.sidebar.title("Menú de Funciones")
-generarPages()
-#Inicializamos nuestro historial de chat
-if "messages" not in st.session_state:
-     st.session_state.messages = [{"role": "assistant", "content": "Hola!, en que puedo ayudarte?"}]
-
-#Definimos modelo
-model = SentenceTransformer("all-MiniLM-L6-v2")
-
-#Elegimos una coleccion
-colecciones = obtener_colecciones()
-coleccion_seleccionada = st.sidebar.selectbox("Selecciona una colección", colecciones)
-
-# Mostrar el historial de mensajes
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Entrada del usuario
-if prompt := st.chat_input("Escribe tus dudas"):
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    
-    with st.chat_message("user"):
-        st.markdown(prompt)
-
-    with st.chat_message("assistant"):
-        if coleccion_seleccionada == "Todas las colecciones":
-            colecciones_disponibles = obtener_colecciones()
-            results = []
-            umbral=1
-            for coleccion in colecciones_disponibles[1:]:  
-                coleccion_results = query_qdrant_sinumbral(prompt,model,coleccion)
-                results.extend(coleccion_results) 
-        else:
-            umbral=0.56
-            results = query_qdrant(prompt, model, coleccion_seleccionada,5,umbral)
-
-        if not results:
-            response = "Disculpa, no tengo información para responder esa pregunta."
-        else:
-            response = st.write_stream(get_response_from_mistral(prompt, results))
-        
-    st.session_state.messages.append({"role": "assistant", "content": response})
+import  streamlit as st 
+from Rag_milvus import query_qdrant, obtener_colecciones, query_qdrant_sinumbral
+from Llm_local import get_response_from_mistral, generarPages
+from sentence_transformers import SentenceTransformer
+
+col1, col2 = st.columns([1, 4])
+with col1:
+    #st.image("Procuradurialogo.jpg", width=600)
+
+with col2:
+    st.markdown("""
+        <div style='display: flex; align-items: center; height: 100%;'>
+            <h1 style='margin: 0; text-align: center;'>ProcurIA</h1>
+        </div>
+    """, unsafe_allow_html=True)
+
+st.sidebar.title("Menú de Funciones")
+generarPages()
+#Inicializamos nuestro historial de chat
+if "messages" not in st.session_state:
+     st.session_state.messages = [{"role": "assistant", "content": "Hola!, en que puedo ayudarte?"}]
+
+#Definimos modelo
+model = SentenceTransformer("all-MiniLM-L6-v2")
+
+#Elegimos una coleccion
+colecciones = obtener_colecciones()
+coleccion_seleccionada = st.sidebar.selectbox("Selecciona una colección", colecciones)
+
+# Mostrar el historial de mensajes
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Entrada del usuario
+if prompt := st.chat_input("Escribe tus dudas"):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    with st.chat_message("assistant"):
+        if coleccion_seleccionada == "Todas las colecciones":
+            colecciones_disponibles = obtener_colecciones()
+            results = []
+            umbral=1
+            for coleccion in colecciones_disponibles[1:]:  
+                coleccion_results = query_qdrant_sinumbral(prompt,model,coleccion)
+                results.extend(coleccion_results) 
+        else:
+            umbral=0.56
+            results = query_qdrant(prompt, model, coleccion_seleccionada,5,umbral)
+
+        if not results:
+            response = "Disculpa, no tengo información para responder esa pregunta."
+        else:
+            response = st.write_stream(get_response_from_mistral(prompt, results))
+        
+    st.session_state.messages.append({"role": "assistant", "content": response})
     st.write(results)