app.py

import streamlit as st
import requests
import json
import os
import datetime

# Constants
SPACE_URL = "https://z7svds7k42bwhhgm.us-east-1.aws.endpoints.huggingface.cloud"
HF_API_KEY = os.getenv("HF_API_KEY")  # Retrieve the Hugging Face API key from system variables
EOS_TOKEN = "<|end|>"
CHAT_HISTORY_DIR = "chat_histories"
IMAGE_PATH = "DubsChat.png"
IMAGE_PATH_2 = "Reboot AI.png"
Dubs_PATH = "Dubs.png"

# Ensure the directory exists
try:
    os.makedirs(CHAT_HISTORY_DIR, exist_ok=True)
except OSError as e:
    st.error(f"Failed to create chat history directory: {e}")

# Streamlit Configurations
st.set_page_config(page_title="DUBSChat", page_icon=IMAGE_PATH, layout="wide")

# If you are using a custom "logo" method:
st.logo(IMAGE_PATH_2)

# -------------------------
#    Utility Functions
# -------------------------
def save_chat_history(session_name, messages):
    """
    Save the chat history to a JSON file.
    """
    file_path = os.path.join(CHAT_HISTORY_DIR, f"{session_name}.json")
    try:
        with open(file_path, "w") as f:
            json.dump(messages, f)
    except IOError as e:
        st.error(f"Failed to save chat history: {e}")


def load_chat_history(file_name):
    """
    Load the chat history from a JSON file.
    """
    file_path = os.path.join(CHAT_HISTORY_DIR, file_name)
    try:
        with open(file_path, "r") as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        st.error("Failed to load chat history. Starting with a new session.")
        return []


def get_saved_sessions():
    """
    Get the list of saved chat sessions.
    """
    return [f.replace(".json", "") for f in os.listdir(CHAT_HISTORY_DIR) if f.endswith(".json")]

# -------------------------
#    Sidebar Configuration
# -------------------------
with st.sidebar:
    # Reset chat
    if st.button("New Chat"):
        st.session_state["messages"] = [
            {"role": "system", "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."},
            {"role": "assistant", "content": "Hello! How can I assist you today?"}
        ]
        st.session_state["session_name"] = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        save_chat_history(st.session_state["session_name"], st.session_state["messages"])
        st.success("Chat reset and new session started.")

    # Hugging Face or “Dubs” API key input
    dubs_key = st.text_input("Enter Dubs Key", key="chatbot_api_key", type="password")
    
    # Load past sessions
    saved_sessions = get_saved_sessions()
    if saved_sessions:
        selected_session = st.radio("Past Sessions:", saved_sessions)
        if st.button("Load Session"):
            st.session_state["messages"] = load_chat_history(f"{selected_session}.json")
            st.session_state["session_name"] = selected_session
            st.success(f"Loaded session: {selected_session}")
    else:
        st.write("No past sessions available.")

# -------------------------
# Chat History Initialization
# -------------------------
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {"role": "system", "content": "You are DUBS, a helpful assistant capable of conversing in a friendly and knowledgeable way."},
        {"role": "assistant", "content": "Hello! How can I assist you today?"}
    ]
if "session_name" not in st.session_state:
    st.session_state["session_name"] = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

# -------------------------
#        Main Chat UI
# -------------------------
st.image(IMAGE_PATH, width=250)
st.markdown("Empowering you with a Sustainable AI")

# Display existing chat history
for message in st.session_state["messages"]:
    if message["role"] == "user":
        st.chat_message("user").write(message["content"])
    elif message["role"] == "assistant":
        st.chat_message("assistant", avatar=Dubs_PATH).write(message["content"])

# -------------------------
#    Streaming Logic
# -------------------------
def stream_response(prompt_text, api_key):
    """
    Stream text from the HF Inference Endpoint (or any streaming API).
    Yields each chunk of text as it arrives.
    """
    try:
    # Match the structure of your working payload:
        payload = {
            "inputs": prompt_text,
            "parameters": {
                "max_new_tokens": 250,
                "return_full_text": False,
                "stream": True
            }
        }
        headers = {
                	"Accept" : "application/json",
                    "Authorization": f"Bearer {api_key}",
                	"Content-Type": "application/json" 
                  }

        # POST request with stream=True to get partial chunks
        response = requests.post(
            SPACE_URL,
            json=payload,
            headers=headers,
            stream=True
        )
        response.raise_for_status()

        # The endpoint presumably returns lines of JSON. Adjust parsing if needed:
        for line in response.iter_lines():
            if line:
                data = json.loads(line.decode("utf-8"))
                # Example: data might be [{"generated_text": "..."}]
                # Adjust if your endpoint returns different JSON keys
                chunk = data[0].get("generated_text", "")
                yield chunk

    except requests.exceptions.Timeout:
        yield "The request timed out. Please try again later."
    except requests.exceptions.RequestException as e:
        yield f"Error: {e}"
    except json.JSONDecodeError:
        yield "Error decoding server response."

# -------------------------
#       User Input
# -------------------------
if prompt := st.chat_input():
    if not dubs_key:
        st.warning("Please provide a valid Dubs Key.")
    else:
        # 1) Add the user's message to session state
        st.session_state["messages"].append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)

        # 2) Build combined chat history for the model prompt
        chat_history = "".join(
            [f"<|{msg['role']}|>{msg['content']}<|end|>" for msg in st.session_state["messages"]]
        )

        # 3) Create a placeholder for the assistant’s streamed response
        with st.spinner("Dubs is thinking... Woof Woof! 🐾"):
            assistant_message_placeholder = st.chat_message("assistant", avatar=Dubs_PATH).empty()

            full_response = ""
            # 4) Stream chunks from the API
            for chunk in stream_response(chat_history, HF_API_KEY):
                full_response += chunk
                # Continuously update the placeholder with the partial response
                assistant_message_placeholder.write(full_response)

        # 5) Save the final assistant message in session state
        st.session_state["messages"].append({"role": "assistant", "content": full_response})
        # 6) Persist updated chat history
        save_chat_history(st.session_state["session_name"], st.session_state["messages"])