src/streamlit_app.py

import streamlit as st
import requests
import os

# Configure the page
st.set_page_config(
    page_title="Document Assistant",
    page_icon="📄",
    layout="wide"
)

# Add a title and description
st.title("Document Assistant")

# Get API URL from environment variable
API_URL = os.getenv("API_URL")


def query_rag_system(query: str):
    """Send query to the RAG endpoint and return results"""
    try:
        # FastAPI expects the query as a raw string in the body
        # Based on FastAPI code: query: str = Body(...)
        response = requests.post(
            f"{API_URL}/rag",
            json=query,  # This should work with your current FastAPI setup
            headers={'Content-Type': 'application/json'},
        )
        response.raise_for_status()
        return response.json()
    except requests.exceptions.Timeout:
        st.error("⏱️ Request timed out. The API might be starting up, please try again.")
        return None
    except requests.exceptions.RequestException as e:
        st.error(f"❌ Error connecting to the API: {str(e)}")
        return None
    except Exception as e:
        st.error(f"❌ Unexpected error: {str(e)}")
        return None

# Create the main query interface
st.markdown("Enter your query to get AI-generated responses based on the documents.")

# Use columns for better layout
col1, col2 = st.columns([4, 1])

with col1:
    query = st.text_area(
        "Enter your query:",
        height=100,
        placeholder="Ask me anything about the document..."
    )

with col2:
    st.write("")  # Add some spacing
    st.write("")  # Add some spacing
    submit_button = st.button("Submit Query", type="primary")

# Add some example queries
with st.expander("💡 Example Queries"):
    st.markdown("""
    - What should be considered when isolating pressure systems to ensure safety?
    - What is the reference density for wet solids in the Alfa Laval separator bowl?
    - What should be stored or discharged in accordance with current rules and directives during valve maintenance?
    """)

# Add a spinner for loading state
if submit_button and query.strip():
    with st.spinner("🤔 Processing your query..."):
        result = query_rag_system(query)
        if result:
            st.subheader("Response")
            response_text = result.get('response', 'No response received')
            st.markdown(response_text)
            
                
elif submit_button and not query.strip():
    st.warning("⚠️ Please enter a query before submitting.")

# Add footer with instructions
st.markdown("---")
st.markdown("""
### 📖 How to use:
1. Enter your question in the text area above
2. Click 'Submit Query' to get an AI-generated response
3. The response will be based on the documents in the system
""")

# Add some spacing at the bottom
st.write("")
st.write("")