app.py

import os
import streamlit as st
from tasks.data_ingestion_task.data_ingestion_task import DataIngestionTask
from tasks.query_handling_task.query_handling_task import QueryHandlingTask
from utils.llama_index_utils import setup_directories

# Title Section
st.title("(PDF) Chat con documentos de Procesos 🗒️")

# Subtitle Section
st.markdown("Retrieval-Augmented Generation")
st.markdown("iniciar chat ...🚀")

# Session State Initialization
if 'messages' not in st.session_state:
 st.session_state.messages = [{'role': 'assistant', "content": 'Hola! Selecciona un pdf para cargar, y hazme una pregunta.'}]

# Sidebar Section
with st.sidebar:
 st.image('image_logo.jpeg', use_column_width=True)
 st.title("Menu:")
 
 uploaded_file = st.file_uploader("Sube un archivo PDF y dale click al botón enviar y procesar.")

 if st.button("Enviar y Procesar"):
 if uploaded_file is not None:
 with st.spinner("Procesando..."):
 # Ensure the data directory exists
 data_dir, persist_dir = setup_directories()
 
 # Save the uploaded file in the data directory
 filepath = os.path.join(data_dir, "saved_pdf.pdf")
 with open(filepath, "wb") as f:
 f.write(uploaded_file.getbuffer())
 
 # Use DataIngestionTask to process the PDF
 ingestion_task = DataIngestionTask(
 config_path='tasks/data_ingestion_task/config.txt',
 input_structure_path='tasks/data_ingestion_task/input_structure.json',
 output_structure_path='tasks/data_ingestion_task/output_structure.json'
 )
 ingestion_task.execute({})
 
 st.success("PDF processed successfully")
 else:
 st.error("Please upload a PDF file before processing.")

 st.subheader("Select Model:")
 model_selected = st.session_state.get('model_selected', '')
 
 col1, col2 = st.columns(2)
 
 if col1.button("GEMMA", key="gemmabtn"):
 st.session_state.model_selected = "GEMMA"
 
 if col2.button("GEMINI", key="geminibtn"):
 st.session_state.model_selected = "GEMINI"
 
 if model_selected:
 st.write(f"Selected Model: **{model_selected}**")

# Chat Input Section
user_prompt = st.chat_input("Pregunta acerca del contenido en el archivo PDF:")
if user_prompt:
 st.session_state.messages.append({'role': 'user', "content": user_prompt})
 
 # Use QueryHandlingTask to handle the query
 query_task = QueryHandlingTask(
 config_path='tasks/query_handling_task/config.txt',
 input_structure_path='tasks/query_handling_task/input_structure.json',
 output_structure_path='tasks/query_handling_task/output_structure.json'
 )
 response = query_task.execute({'query': user_prompt})
 
 st.session_state.messages.append({'role': 'assistant', "content": response})

# Chat Message Display
for message in st.session_state.messages:
 with st.chat_message(message['role']):
 st.write(message['content'])