app.py

import streamlit as st
import requests
import os
from dotenv import load_dotenv

load_dotenv()

def reset_conversation():
    '''
    Resets Conversation
    '''
    st.session_state.messages = []
    return None

# Define model links for the Addiction Recovery and Mental Health models
model_links = {
    "Addiction recovery AI": "NousResearch/Nous-Hermes-2-Yi-34B",
    "Mental health AI": "NousResearch/Nous-Hermes-2-Yi-34B"
}

# Define pre-instructions for each model
model_pre_instructions = {
    "Addiction recovery AI": "This AI is trained to provide support and guidance related to addiction recovery.",
    "Mental health AI": "This AI is trained to provide support and guidance related to mental health."
}

# Function to interact with the selected model via the Together API
def interact_with_together_api(messages, model_link):
    all_messages = []

    if not messages:  # If history is empty
        all_messages.append({"role": "user", "content": ""})
        history = [("", "")]  # Add dummy values to prevent unpacking error

    for human, assistant in messages:
        all_messages.append({"role": "user", "content": human})
        all_messages.append({"role": "assistant", "content": assistant})

    all_messages.append({"role": "user", "content": messages[-1][1]})

    url = "https://api.together.xyz/v1/chat/completions"
    payload = {
        "model": model_link,
        "temperature": 1.05,
        "top_p": 0.9,
        "top_k": 50,
        "repetition_penalty": 1,
        "n": 1,
        "messages": all_messages,
    }

    TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "Authorization": f"Bearer {TOGETHER_API_KEY}",
    }

    response = requests.post(url, json=payload, headers=headers)
    response.raise_for_status()  # Ensure HTTP request was successful

    # Extract response from JSON
    response_data = response.json()
    assistant_response = response_data["choices"][0]["message"]["content"]

    return assistant_response

# Create sidebar with model selection dropdown and reset button
selected_model = st.sidebar.selectbox("Select Model", list(model_links.keys()))
st.sidebar.write(f"Pre-Instructions: {model_pre_instructions[selected_model]}")
st.sidebar.button('Reset Chat', on_click=reset_conversation)

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message[0]):
        st.markdown(message[1])

# Accept user input
if prompt := st.chat_input(f"Hi, I'm {selected_model}, ask me a question"):
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(prompt)
    # Add user message to chat history
    st.session_state.messages.append(("user", prompt))

    # Interact with the selected model
    assistant_response = interact_with_together_api(st.session_state.messages, model_links[selected_model])

    # Display assistant response in chat message container
    with st.empty():
        st.markdown("AI is typing...")
        st.empty()
        st.markdown(assistant_response)

    # Add assistant response to chat history
    st.session_state.messages.append(("assistant", assistant_response))