File size: 4,717 Bytes
ad06665 808e67d ad06665 37be6ad ad06665 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
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.")
|