import streamlit as st
import cohere

st.set_page_config(page_title="Cohere Chat Interface", layout="wide")

MODEL_PFPS = {
    "command-a-03-2025": "/media/pfp/cohere-pfp.png",
    "command-r7b-12-2024": "/media/pfp/cohere-pfp.png",
    "command-r-plus-04-2024": "/media/pfp/cohere-pfp.png",
    "command-r-plus": "/media/pfp/cohere-pfp.png",
    "command-r-08-2024": "/media/pfp/cohere-pfp.png",
    "command-r-03-2024": "/media/pfp/cohere-pfp.png",
    "command-r": "/media/pfp/cohere-pfp.png",
    "command": "/media/pfp/cohere-pfp.png",
    "command-nightly": "/media/pfp/cohere-pfp.png",
    "command-light": "/media/pfp/cohere-pfp.png",
    "command-light-nightly": "/media/pfp/cohere-pfp.png"
}

USER_PFP = "https://example.com/user-default.png"

MODEL_INFO = {
    "command-a-03-2025": {
        "description": "Command A is our most performant model to date, excelling at tool use, agents, retrieval augmented generation (RAG), and multilingual use cases. Command A has a context length of 256K, only requires two GPUs to run, and has 150% higher throughput compared to Command R+ 08-2024.",
        "context_window": "256K tokens",
        "output_tokens": "8K tokens"
    },
    "command-r7b-12-2024": {
        "description": "command-r7b-12-2024 is a small, fast update delivered in December 2024. It excels at RAG, tool use, agents, and similar tasks requiring complex reasoning and multiple steps.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    },
    "command-r-plus-04-2024": {
        "description": "Command R+ is an instruction-following conversational model that performs language tasks at a higher quality, more reliably, and with a longer context than previous models. It is best suited for complex RAG workflows and multi-step tool use.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    },
    "command-r-plus": {
        "description": "command-r-plus is an alias for command-r-plus-04-2024, so if you use command-r-plus in the API, that's the model you're pointing to.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    },
    "command-r-08-2024": {
        "description": "command-r-08-2024 is an update of the Command R model, delivered in August 2024.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    },
    "command-r-03-2024": {
        "description": "Command R is an instruction-following conversational model that performs language tasks at a higher quality, more reliably, and with a longer context than previous models. It can be used for complex workflows like code generation, retrieval augmented generation (RAG), tool use, and agents.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    },
    "command-r": {
        "description": "command-r is an alias for command-r-03-2024, so if you use command-r in the API, that's the model you're pointing to.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    },
    "command": {
        "description": "An instruction-following conversational model that performs language tasks with high quality, more reliably and with a longer context than our base generative models.",
        "context_window": "4K tokens",
        "output_tokens": "4K tokens"
    },
    "command-nightly": {
        "description": "Nightly version of command - experimental and unstable. Not recommended for production use.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    },
    "command-light": {
        "description": "Smaller, faster version of command with similar capabilities.",
        "context_window": "4K tokens",
        "output_tokens": "4K tokens"
    },
    "command-light-nightly": {
        "description": "Nightly version of command-light - experimental and unstable. Not for production use.",
        "context_window": "128K tokens",
        "output_tokens": "4K tokens"
    }
}

with st.sidebar:
    st.title("Configuration")
    api_key = st.text_input("Cohere API Key", type="password")
    
    selected_model = st.selectbox(
        "Select Model",
        options=list(MODEL_INFO.keys()),
        format_func=lambda x: x.upper()
    )
    
    st.divider()
    st.subheader("Model Details")
    st.image(MODEL_PFPS[selected_model], width=80)
    st.markdown(f"**{selected_model}**")
    st.markdown(MODEL_INFO[selected_model]["description"])
    st.markdown(f"**Context Window:** {MODEL_INFO[selected_model]['context_window']}")
    st.markdown(f"**Max Output:** {MODEL_INFO[selected_model]['output_tokens']}")

st.title(f"Chat with {selected_model.upper()}")
st.image(MODEL_PFPS[selected_model], width=50)

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

for message in st.session_state.messages:
    avatar = USER_PFP if message["role"] == "user" else MODEL_PFPS[selected_model]
    with st.chat_message(message["role"], avatar=avatar):
        st.markdown(message["content"])

if prompt := st.chat_input("Type your message..."):
    if not api_key:
        st.error("API key required - enter in sidebar")
        st.stop()
    
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user", avatar=USER_PFP):
        st.markdown(prompt)
    
    try:
        co = cohere.ClientV2(api_key)
        
        with st.chat_message("assistant", avatar=MODEL_PFPS[selected_model]):
            response = co.chat(
                model=selected_model,
                messages=st.session_state.messages
            )
            
            if hasattr(response, 'text'):
                full_response = response.text
            else:
                full_response = "Error: Unexpected API response format"
            
            st.markdown(full_response)
        
        st.session_state.messages.append({"role": "assistant", "content": full_response})
    
    except cohere.CohereError as e:
        st.error(f"Cohere API Error: {str(e)}")
    except Exception as e:
        st.error(f"General Error: {str(e)}")