import streamlit as st
import requests
import os

st.title('Medi Chat')

auth = os.environ['auth']
API_URL = os.environ['API_URL']

# for dedicated
# headers = {
# 	"Accept" : "application/json",
# 	"Authorization": f"{auth}",
# 	"Content-Type": "application/json" 
# }

# for serverless
headers = {"Authorization": f"Bearer {auth}"}

print('auth token pulled')

if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.json()

if prompt := st.chat_input():
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)

    # Make the query with the latest messages
    response = query({
    	"inputs": f"<<SYS>> You are an anesthesiologist instructor. Keep your answers short and provide them in steps where applicable. <</SYS>> [INST] {st.session_state.messages} [/INST]",
    	"parameters": {
    		"max_new_tokens": 600
    	}
    })
    print(f'response: {response}')

    msg = response[0]['generated_text'].strip()
    
    # Append the generated text to the messages and display it
    st.session_state.messages.append({"role": "assistant", "content": msg})
    st.chat_message("assistant").write(msg)