src/streamlit_app.py

import streamlit as st
import requests
import os

# Load your OpenRouter API key
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY") or st.secrets.get("OPENROUTER_API_KEY", "")
MODEL_NAME = "deepseek/deepseek-r1-0528:free"

# App title
st.title("WhatsApp-like Chat with DeepSeek-R1")

# Custom CSS
st.markdown("""
<style>
    .stChatMessage { background-color: transparent !important; padding: 0 !important; width: 100%; }
    .user-container, .assistant-container { display: flex; flex-direction: column; margin-bottom: 12px; }
    .user-container { align-items: flex-end; }
    .assistant-container { align-items: flex-start; }
    .user-message, .assistant-message {
        padding: 8px 12px; border-radius: 15px; max-width: 80%; margin-top: 4px;
    }
    .user-message { background-color: #DCF8C6; border-radius: 15px 15px 0 15px; }
    .assistant-message { background-color: #FFF; border: 1px solid #ECECEC; border-radius: 15px 15px 15px 0; }
    .sender-label { font-weight: bold; font-size: 0.8em; color: #555; }
    .user-label { margin-right: 8px; }
    .assistant-label { margin-left: 8px; }
    [data-testid="stChatMessage"] { background-color: transparent !important; }
    [data-testid="stChatMessage"] > div:first-child { display: none !important; }
</style>
""", unsafe_allow_html=True)

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": "You are a helpful assistant."}]

# Display chat messages
for message in st.session_state.messages:
    if message["role"] == "system":
        continue
    with st.chat_message(message["role"]):
        if message["role"] == "user":
            st.markdown(f"""
            <div class="user-container">
                <div class="user-label sender-label">You:</div>
                <div class="user-message">{message['content']}</div>
            </div>
            """, unsafe_allow_html=True)
        else:
            st.markdown(f"""
            <div class="assistant-container">
                <div class="assistant-label sender-label">Bot:</div>
                <div class="assistant-message">{message['content']}</div>
            </div>
            """, unsafe_allow_html=True)

# Input field
if prompt := st.chat_input("Type your message..."):
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Show user message immediately
    with st.chat_message("user"):
        st.markdown(f"""
        <div class="user-container">
            <div class="user-label sender-label">Me:</div>
            <div class="user-message">{prompt}</div>
        </div>
        """, unsafe_allow_html=True)

    # Show assistant reply placeholder
    with st.chat_message("assistant"):
        placeholder = st.empty()
        placeholder.markdown("""
        <div class="assistant-container">
            <div class="assistant-label sender-label">Bot:</div>
            <div class="assistant-message">▌</div>
        </div>
        """, unsafe_allow_html=True)

        # Call OpenRouter
        headers = {
            "Authorization": f"Bearer {OPENROUTER_API_KEY}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://your-app-name.streamlit.app/",
            "X-Title": "DeepSeek WhatsApp Chat"
        }

        body = {
            "model": MODEL_NAME,
            "messages": st.session_state.messages,
            "stream": False  # OpenRouter's free models might not support streaming
        }

        try:
            res = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=body)
            res.raise_for_status()
            full_response = res.json()["choices"][0]["message"]["content"]
        except Exception as e:
            full_response = f"Error: {str(e)}"

        # Render final response
        placeholder.markdown(f"""
        <div class="assistant-container">
            <div class="assistant-label sender-label">Bot:</div>
            <div class="assistant-message">{full_response}</div>
        </div>
        """, unsafe_allow_html=True)

        # Save assistant reply
        st.session_state.messages.append({"role": "assistant", "content": full_response})