app.py

import streamlit as st
from streamlit_chat import message
import requests
# changed
url = 'https://api.webraft.in/v2/chat/completions'
st.set_page_config(page_title="WebraftAI Chat", page_icon=":robot_face:")
st.markdown("<h1 style='text-align: center;'>WebraftAI Chat</h1>", unsafe_allow_html=True)

if 'generated' not in st.session_state:
    st.session_state['generated'] = []
if 'past' not in st.session_state:
    st.session_state['past'] = []

if 'model_name' not in st.session_state:
    st.session_state['model_name'] = []
if 'cost' not in st.session_state:
    st.session_state['cost'] = []
if 'total_tokens' not in st.session_state:
    st.session_state['total_tokens'] = []
if 'total_cost' not in st.session_state:
    st.session_state['total_cost'] = 1
if 'api_key' not in st.session_state:
    st.session_state['api_key']=""

st.sidebar.title("Settings")
model_name = st.sidebar.selectbox("Chat Model:", ("gpt-3.5-turbo","gpt-3.5-turbo-1106","gpt-4","gpt-4-32k","gpt-4-1106-preview","gpt-4-0125-preview","gpt-4o","claude-3-haiku","claude-3-sonnet","claude-3-opus",""))
image_model = st.sidebar.selectbox("Image Model:", ("dalle-3","sdxl", "midjourney","stable-cascade"))
counter_placeholder = st.sidebar.empty()

api_key = st.sidebar.text_input("API_Key", value=st.session_state['api_key'], max_chars=None,placeholder=st.session_state['api_key'], key=None, type="password", label_visibility="visible")
st.session_state['api_key'] = api_key
max_tokens = st.sidebar.text_input("Max_tokens", value=4096, max_chars=None, key=None,placeholder=4096, type="default")
top_p = st.sidebar.text_input("Top_p", value=0.7, max_chars=None, key=None,placeholder=0.7, type="default")
temperature = st.sidebar.text_input("Temperature", value=0.1, max_chars=None, key=None,placeholder=0.1, type="default")
system_message = st.sidebar.text_area("System_message", value="You are a helpful AI assistant.", max_chars=None, key=None,placeholder="You are a helpful AI assistant.")
clear_button = st.sidebar.button("Clear Conversation", key="clear")
headers = {
    'Content-Type': 'application/json',
    'Authorization': f'Bearer {api_key}',
}
if 'messages' not in st.session_state:
    st.session_state['messages'] = [
        {"role": "system", "content": system_message}
    ]
model = model_name

if clear_button:
    st.session_state['generated'] = []
    st.session_state['past'] = []
    st.session_state['messages'] = [
        {"role": "system", "content": system_message}
    ]
    st.session_state['number_tokens'] = []
    st.session_state['model_name'] = []
    st.session_state['cost'] = []
    st.session_state['total_cost'] = 0.0
    st.session_state['total_tokens'] = []


def evaluate(model, question, max_tokens=4096, top_p=0.7,temperature=0.1):

    data = {
        "model": model,
        "max_tokens": max_tokens,
        "messages": question,
        "temperature":temperature,
        "top_p":top_p

    }
    response = requests.post(url, headers=headers, json=data)
    try:

        response = response.json()

        choices = response.get('choices', [])
        if choices:
            sentie = response
            sentence = choices[0].get('message', {}).get('content'," API Returned: \n ```{sentie}``` \n \n Is the API not working? check https://stats.uptimerobot.com/KkZE2T84mj \n If you are still encountering errors then please open a ticket from our discord server at https://discord.gg/webraftai . ")
        else:
            sentie = response
            sentence = f"API Returned: \n ```{sentie}``` \n \n Is the API not working? check https://stats.uptimerobot.com/KkZE2T84mj \n If you are still encountering errors then please open a ticket from our discord server at https://discord.gg/webraftai . "


    except requests.exceptions.JSONDecodeError as e:
        sentie= response.content
        sentence = f"API Returned:\n ```{sentie}``` \n \n Is the API not working? check https://stats.uptimerobot.com/KkZE2T84mj \n If you are still encountering errors then please open a ticket from our discord server at https://discord.gg/webraftai . "
    return sentence


def generate_response(prompt):
    st.session_state['messages'].append({"role": "user", "content": prompt})
    prompt = prompt.lower()
    if prompt.startswith("draw"):
        data = {"prompt": prompt.replace("draw",""),
                "model": image_model
                }

        response = requests.post("https://api.webraft.in/api/images/generations", json=data, headers=headers)
        total_tokens = "153"
        prompt_tokens = "153"
        completion_tokens = "153"
        try:

            response = response.json()
            imageurl = response['data'][0]['url']
            ie = f"Here is the generated image:\n {imageurl} \n \n Model used: `{image_model}`"
            st.session_state['messages'].append({"role": "assistant", "content": ie})

        except requests.exceptions.JSONDecodeError as e:
            sentie = response.content
            ie= f"Image generation API returned:\n ```{sentie}``` \n \n Is the API not working? check https://stats.uptimerobot.com/KkZE2T84mj \n If you are still encountering errors then please open a ticket from our discord server at https://discord.gg/webraftai . "
            st.session_state['messages'].append({"role": "assistant", "content": ie})
        return ie,total_tokens,prompt_tokens,completion_tokens


    question = st.session_state['messages']

    sentence = evaluate(model, question,top_p=float(top_p),max_tokens=int(max_tokens),temperature=float(temperature))

    response = sentence
    st.session_state['messages'].append({"role": "assistant", "content": response})

    # print(st.session_state['messages'])
    total_tokens = "153"
    prompt_tokens = "153"
    completion_tokens = "153"
    return response, total_tokens, prompt_tokens, completion_tokens


response_container = st.container()
container = st.container()

with container:
    with st.form(key='my_form', clear_on_submit=True):
        user_input = st.text_area("You:", key='input', height=2)
        submit_button = st.form_submit_button(label='✉')

    if submit_button and user_input:
        output, total_tokens, prompt_tokens, completion_tokens = generate_response(user_input)
        st.session_state['past'].append(user_input)
        st.session_state['generated'].append(output)
        st.session_state['model_name'].append(model_name)
        st.session_state['total_tokens'].append(total_tokens)


if st.session_state['generated']:
    with response_container:
        for i in range(len(st.session_state['generated'])):
            message(st.session_state["past"][i], is_user=True, key=str(i) + '_user')
            message(st.session_state["generated"][i], key=str(i))