Spaces:

Kirtan001
/

test

Paused

test / src /app.py

UI: Rebranding and cleanup

808e67d 9 months ago

4.72 kB

import os
import sys
import streamlit as st
from pathlib import Path

# Add project root to path for local execution
sys.path.append(str(Path(__file__).parent.parent))

import psutil
from textwrap import dedent

from src.rag_engine import SatelliteRAG
from src.config import settings

def check_system_health():
    """Check memory availability before processing."""
    process = psutil.Process(os.getpid())
    mem_info = process.memory_info()
    mem_mb = mem_info.rss / 1024 / 1024
    
    # Hugging Face Free Tier ~16GB, Streamlit Cloud ~1GB
    # If using > 80% (just a heuristics check, not hard limit)
    # real check is total memory available
    total_mem = psutil.virtual_memory().total / 1024 / 1024
    avail_mem = psutil.virtual_memory().available / 1024 / 1024
    
    return {
        "used_mb": mem_mb,
        "available_mb": avail_mem,
        "warning": avail_mem < 250 # Warn if < 250MB available
    }

# Setup Page
st.set_page_config(page_title="Space Satellite Assistant", page_icon="🛰️")
st.title("🛰️ Space Satellite Assistant")

# Sidebar
with st.sidebar:
    st.markdown("### Configuration")
    if not settings.GROQ_API_KEY:
        st.error("⚠️ GROQ_API_KEY not found in .env!")
    else:
        st.success("System Online")
    
    st.markdown("---")
    st.markdown("**Engine:** SatelliteRAG v1.0")
    st.markdown("**Knowledge Base:** 637 Satellites")

# Initialize Resources (Cached at startup)
@st.cache_resource
def get_engine():
    try:
        with st.spinner("Initializing RAG Engine (Loading Models & DB)..."):
            return SatelliteRAG()
    except Exception as e:
        # We catch strictly to prevent early app crash
        return f"ERROR: {str(e)}"

# Early Engine Check
engine_status = get_engine()

with st.sidebar:
    st.markdown("---")
    st.markdown("### System Status")
    if isinstance(engine_status, str) and "ERROR" in engine_status:
        st.error(f"🔴 Engine: {engine_status}")
    elif engine_status is None:
        st.warning("🟠 Engine: Initializing...")
    else:
        st.success("🟢 Engine: Ready")

# Chat Logic
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display History
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Handle Input
if prompt := st.chat_input("Ask about any satellite (e.g., 'What is Gaofen 1?')..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        # Use the already initialized engine
        engine = engine_status

        if isinstance(engine, SatelliteRAG):
            with st.spinner("Analyzing satellite data..."):
                try:
                    # Health Check
                    health = check_system_health()
                    if health["warning"]:
                        st.warning(f"⚠️ Low Memory Warning: Only {health['available_mb']:.0f}MB available. Query might be slow.")
                    
                    # Construct Chat History
                    # We need pairs of (User, AI) from session_state.messages
                    # Excluding the current new prompt which is already appended but not part of 'history' yet for this context
                    chat_history = []
                    msgs = st.session_state.messages[:-1] 
                    for i in range(0, len(msgs) - 1, 2):
                        if msgs[i]["role"] == "user" and msgs[i+1]["role"] == "assistant":
                            chat_history.append((msgs[i]["content"], msgs[i+1]["content"]))
                    
                    response, docs = engine.query(prompt, chat_history=chat_history)

                    st.markdown(response)

                    # Show sources in expander
                    with st.expander(f"View Source Context ({len(docs)} chunks)"):
                        for i, doc in enumerate(docs, 1):
                            st.markdown(f"**Source {i}:** {doc.metadata.get('name')}")
                            st.text(doc.page_content[:300] + "...")

                    st.session_state.messages.append({"role": "assistant", "content": response})
                except Exception as e:
                    st.error("⚠️ An error occurred during analysis.")
                    st.markdown(f"**Error Details:** `{str(e)}`")
                    st.info("💡 **Tip:** If this happens frequently, the system might be out of memory (OOM). Try waiting a few seconds and refreshing.")
        else:
            st.error("⚠️ RAG Engine is not ready. Please check the sidebar for status.")

	import os
	import sys
	import streamlit as st
	from pathlib import Path

	# Add project root to path for local execution
	sys.path.append(str(Path(__file__).parent.parent))

	import psutil
	from textwrap import dedent

	from src.rag_engine import SatelliteRAG
	from src.config import settings

	def check_system_health():
	"""Check memory availability before processing."""
	process = psutil.Process(os.getpid())
	mem_info = process.memory_info()
	mem_mb = mem_info.rss / 1024 / 1024

	# Hugging Face Free Tier ~16GB, Streamlit Cloud ~1GB
	# If using > 80% (just a heuristics check, not hard limit)
	# real check is total memory available
	total_mem = psutil.virtual_memory().total / 1024 / 1024
	avail_mem = psutil.virtual_memory().available / 1024 / 1024

	return {
	"used_mb": mem_mb,
	"available_mb": avail_mem,
	"warning": avail_mem < 250 # Warn if < 250MB available
	}

	# Setup Page
	st.set_page_config(page_title="Space Satellite Assistant", page_icon="🛰️")
	st.title("🛰️ Space Satellite Assistant")

	# Sidebar
	with st.sidebar:
	st.markdown("### Configuration")
	if not settings.GROQ_API_KEY:
	st.error("⚠️ GROQ_API_KEY not found in .env!")
	else:
	st.success("System Online")

	st.markdown("---")
	st.markdown("Engine: SatelliteRAG v1.0")
	st.markdown("Knowledge Base: 637 Satellites")

	# Initialize Resources (Cached at startup)
	@st.cache_resource
	def get_engine():
	try:
	with st.spinner("Initializing RAG Engine (Loading Models & DB)..."):
	return SatelliteRAG()
	except Exception as e:
	# We catch strictly to prevent early app crash
	return f"ERROR: {str(e)}"

	# Early Engine Check
	engine_status = get_engine()

	with st.sidebar:
	st.markdown("---")
	st.markdown("### System Status")
	if isinstance(engine_status, str) and "ERROR" in engine_status:
	st.error(f"🔴 Engine: {engine_status}")
	elif engine_status is None:
	st.warning("🟠 Engine: Initializing...")
	else:
	st.success("🟢 Engine: Ready")

	# Chat Logic
	if "messages" not in st.session_state:
	st.session_state.messages = []

	# Display History
	for message in st.session_state.messages:
	with st.chat_message(message["role"]):
	st.markdown(message["content"])

	# Handle Input
	if prompt := st.chat_input("Ask about any satellite (e.g., 'What is Gaofen 1?')..."):
	st.session_state.messages.append({"role": "user", "content": prompt})
	with st.chat_message("user"):
	st.markdown(prompt)

	with st.chat_message("assistant"):
	# Use the already initialized engine
	engine = engine_status

	if isinstance(engine, SatelliteRAG):
	with st.spinner("Analyzing satellite data..."):
	try:
	# Health Check
	health = check_system_health()
	if health["warning"]:
	st.warning(f"⚠️ Low Memory Warning: Only {health['available_mb']:.0f}MB available. Query might be slow.")

	# Construct Chat History
	# We need pairs of (User, AI) from session_state.messages
	# Excluding the current new prompt which is already appended but not part of 'history' yet for this context
	chat_history = []
	msgs = st.session_state.messages[:-1]
	for i in range(0, len(msgs) - 1, 2):
	if msgs[i]["role"] == "user" and msgs[i+1]["role"] == "assistant":
	chat_history.append((msgs[i]["content"], msgs[i+1]["content"]))

	response, docs = engine.query(prompt, chat_history=chat_history)

	st.markdown(response)

	# Show sources in expander
	with st.expander(f"View Source Context ({len(docs)} chunks)"):
	for i, doc in enumerate(docs, 1):
	st.markdown(f"Source {i}: {doc.metadata.get('name')}")
	st.text(doc.page_content[:300] + "...")

	st.session_state.messages.append({"role": "assistant", "content": response})
	except Exception as e:
	st.error("⚠️ An error occurred during analysis.")
	st.markdown(f"Error Details: `{str(e)}`")
	st.info("💡 Tip: If this happens frequently, the system might be out of memory (OOM). Try waiting a few seconds and refreshing.")
	else:
	st.error("⚠️ RAG Engine is not ready. Please check the sidebar for status.")