app.py

# Force redeploy trigger - version 1.5
import streamlit as st
from utils.config import config
import requests
import json
import os
from core.memory import load_user_state
from core.llm import LLMClient

# Set page config
st.set_page_config(page_title="AI Life Coach", page_icon="🧘", layout="centered")

# Initialize session state for ngrok URL
if "ngrok_url" not in st.session_state:
    st.session_state.ngrok_url = config.ollama_host

# Sidebar for user selection
st.sidebar.title("🧘 AI Life Coach")
user = st.sidebar.selectbox("Select User", ["Rob", "Sarah"])

# Ngrok URL input in sidebar
st.sidebar.markdown("---")
st.sidebar.subheader("Ollama Connection")
ngrok_input = st.sidebar.text_input("Ngrok URL", value=st.session_state.ngrok_url)
if st.sidebar.button("Update Ngrok URL"):
    st.session_state.ngrok_url = ngrok_input
    st.sidebar.success("Ngrok URL updated!")
    st.experimental_rerun()

st.sidebar.markdown("---")

# Get environment info
BASE_URL = os.environ.get("SPACE_ID", "")  # Will be set in HF Spaces
IS_HF_SPACE = bool(BASE_URL)

# Headers to skip ngrok browser warning
NGROK_HEADERS = {
    "ngrok-skip-browser-warning": "true",
    "User-Agent": "AI-Life-Coach-App"
}

# Session state for model status
if "model_status" not in st.session_state:
    st.session_state.model_status = "checking"

if "available_models" not in st.session_state:
    st.session_state.available_models = []

# Fetch Ollama status
def get_ollama_status(ngrok_url):
    try:
        # Try to connect to the remote Ollama service directly
        response = requests.get(
            f"{ngrok_url}/api/tags", 
            headers=NGROK_HEADERS,
            timeout=10
        )
        if response.status_code == 200:
            models = response.json().get("models", [])
            model_names = [m.get("name") for m in models]
            st.session_state.available_models = model_names
            
            if models:
                return {
                    "running": True, 
                    "model_loaded": models[0].get("name"),
                    "remote_host": ngrok_url,
                    "available_models": model_names
                }
            else:
                st.session_state.model_status = "no_models"
                return {
                    "running": False, 
                    "model_loaded": None,
                    "remote_host": ngrok_url,
                    "message": "Connected to Ollama but no models found"
                }
    except Exception as e:
        st.session_state.model_status = "unreachable"
        # If direct connection fails, return error info
        return {
            "running": False, 
            "model_loaded": None,
            "error": str(e),
            "remote_host": ngrok_url
        }

# Poll for model availability
def poll_model_status(ngrok_url):
    if st.session_state.model_status in ["checking", "no_models"]:
        try:
            response = requests.get(
                f"{ngrok_url}/api/tags", 
                headers=NGROK_HEADERS,
                timeout=5
            )
            if response.status_code == 200:
                models = response.json().get("models", [])
                model_names = [m.get("name") for m in models]
                st.session_state.available_models = model_names
                
                if config.local_model_name in model_names:
                    st.session_state.model_status = "ready"
                elif models:
                    st.session_state.model_status = "different_models"
                else:
                    st.session_state.model_status = "no_models"
        except:
            st.session_state.model_status = "unreachable"

# After user selects name, load conversation history
def get_conversation_history(user_id):
    user_state = load_user_state(user_id)
    if user_state and "conversation" in user_state:
        return json.loads(user_state["conversation"])
    return []

# Check Ollama status with the current ngrok URL
ollama_status = get_ollama_status(st.session_state.ngrok_url)

# Poll for model status (run once per session)
poll_model_status(st.session_state.ngrok_url)

# Display Ollama status
use_fallback = not ollama_status.get("running", False) or config.use_fallback

if use_fallback:
    st.sidebar.warning("🌐 Using Hugging Face fallback (Ollama not available)")
    if "error" in ollama_status:
        st.sidebar.caption(f"Error: {ollama_status['error'][:50]}...")
else:
    st.sidebar.success(f"🧠 Ollama Model: {ollama_status['model_loaded']}")
    st.sidebar.info(f"Connected to: {ollama_status['remote_host']}")

# Model status indicator
model_status_container = st.sidebar.empty()
if st.session_state.model_status == "ready":
    model_status_container.success("✅ Model Ready")
elif st.session_state.model_status == "checking":
    model_status_container.info("🔍 Checking model...")
elif st.session_state.model_status == "no_models":
    model_status_container.warning("⚠️ No models found")
elif st.session_state.model_status == "different_models":
    model_status_container.warning("⚠️ Different models available")
else:  # unreachable
    model_status_container.error("❌ Ollama unreachable")

# Main chat interface
st.title("🧘 AI Life Coach")
st.markdown("Talk to your personal development assistant.")

# Show detailed status
with st.expander("🔍 Connection Status"):
    st.write("Ollama Status:", ollama_status)
    st.write("Model Status:", st.session_state.model_status)
    st.write("Available Models:", st.session_state.available_models)
    st.write("Environment Info:")
    st.write("- Is HF Space:", IS_HF_SPACE)
    st.write("- Base URL:", BASE_URL or "Not in HF Space")
    st.write("- Configured Ollama Host:", config.ollama_host)
    st.write("- Current Ngrok URL:", st.session_state.ngrok_url)
    st.write("- Using Fallback:", use_fallback)

# Function to send message to Ollama
def send_to_ollama(user_input, conversation_history, ngrok_url):
    try:
        payload = {
            "model": config.local_model_name,
            "messages": conversation_history,
            "stream": False
        }
        
        response = requests.post(
            f"{ngrok_url}/api/chat",
            json=payload,
            headers=NGROK_HEADERS,
            timeout=60
        )
        
        if response.status_code == 200:
            response_data = response.json()
            return response_data.get("message", {}).get("content", "")
        else:
            st.error(f"Ollama API error: {response.status_code}")
            st.error(response.text[:200])
            return None
    except Exception as e:
        st.error(f"Connection error: {e}")
        return None

# Function to send message to Hugging Face (fallback)
def send_to_hf(user_input, conversation_history):
    try:
        # Initialize LLM client for Hugging Face
        llm_client = LLMClient(provider="huggingface")
        
        # Format prompt for HF
        prompt = ""
        for msg in conversation_history:
            role = msg["role"]
            content = msg["content"]
            if role == "system":
                prompt += f"System: {content}\n"
            elif role == "user":
                prompt += f"Human: {content}\n"
            elif role == "assistant":
                prompt += f"Assistant: {content}\n"
        prompt += "Assistant:"
        
        response = llm_client.generate(prompt, max_tokens=500, stream=False)
        return response
    except Exception as e:
        st.error(f"Hugging Face API error: {e}")
        return None

# Display conversation history
conversation = get_conversation_history(user)
for msg in conversation:
    role = msg["role"].capitalize()
    content = msg["content"]
    st.markdown(f"**{role}:** {content}")

# Chat input
user_input = st.text_input("Your message...", key="input")
if st.button("Send"):
    if user_input.strip() == "":
        st.warning("Please enter a message.")
    else:
        # Display user message
        st.markdown(f"**You:** {user_input}")

        # Prepare conversation history
        conversation_history = [{"role": msg["role"], "content": msg["content"]} 
                              for msg in conversation[-5:]]  # Last 5 messages
        conversation_history.append({"role": "user", "content": user_input})
        
        # Send to appropriate backend
        with st.spinner("AI Coach is thinking..."):
            if use_fallback:
                ai_response = send_to_hf(user_input, conversation_history)
                backend_used = "Hugging Face"
            else:
                ai_response = send_to_ollama(user_input, conversation_history, st.session_state.ngrok_url)
                backend_used = "Ollama"
            
            if ai_response:
                st.markdown(f"**AI Coach ({backend_used}):** {ai_response}")
                # Note: In a production app, we'd save the conversation to Redis here
            else:
                st.error(f"Failed to get response from {backend_used}.")