Update app.py

app.py

@@ -1,4 +1,4 @@
-import cohere
+from premai import Prem
 import streamlit as st
 from streamlit.components.v1 import html
 import re
@@ -6,43 +6,45 @@ import urllib.parse
 
 st.title("Prem Chat UI")
 
-if "api_key" not in st.session_state:
+if "api_key" not in st.session_state and "project_id" not in st.session_state:
     api_key = st.text_input("Enter your API Key", type="password")
-    if api_key:
+    project_id = st.text_input("Enter your project ID")
+    if api_key and project_id:
         st.session_state.api_key = api_key
-        client = cohere.Client(api_key=api_key)
+        st.session_state.project_id = project_id
+        client = Prem(api_key=api_key)
         st.rerun()
     else:
-        st.warning("Please enter your API key to use the app. You can obtain your API key from here: https://dashboard.cohere.com/api-keys")
+        st.warning("Please enter your API key and Project ID to use the app.")
         st.stop()
 else:
-    client = cohere.Client(api_key=st.session_state.api_key)
+    client = Prem(api_key=st.session_state.api_key)
 
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-def get_ai_response(prompt, chat_history):
-    with st.chat_message("ai"):
+def get_ai_response(prompt, messages):
+    with st.chat_message("assistant"):
         penalty_kwargs = {
             "frequency_penalty" if penalty_type == "Frequency Penalty" else "presence_penalty": penalty_value
         }
         
-        stream = client.chat_stream(
-            message=prompt,
+        response = client.chat.completions.create(
+            project_id=st.session_state.project_id,
+            messages=messages,
+            stream=True,
             model=model,
-            preamble=preamble,
-            chat_history=chat_history,
+            system_prompt=system_prompt,
             temperature=temperature,
-            k=k,
-            p=p,
-            **penalty_kwargs
+            top_p=top_p,
+            **penalty_kwargs,
         )
         
         response = ""
         placeholder = st.empty()
-        for event in stream:
-            if event.event_type == "text-generation":
-                content = event.text
+        for chunk in response:
+            if chunk.choices[0].delta["content"]:
+                content = chunk.choices[0].delta["content"]
                 response += content
                 placeholder.markdown(response)
     
@@ -50,9 +52,8 @@ def get_ai_response(prompt, chat_history):
 
 def display_messages():
     for i, message in enumerate(st.session_state.messages):
-        name = "user" if message["role"] == "USER" else "ai"
-        with st.chat_message(name):
-            st.markdown(message["text"])
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
             col1, col2, col3, col4 = st.columns([1, 1, 1, 1])
             with col1:
                 if st.button("Edit", key=f"edit_{i}_{len(st.session_state.messages)}"):
@@ -63,7 +64,7 @@ def display_messages():
                     del st.session_state.messages[i]
                     st.rerun()
             with col3:
-                text_to_copy = message["text"]
+                text_to_copy = message["content"]
                 # PythonでURLエンコード
                 text_to_copy_escaped = urllib.parse.quote(text_to_copy)
 
@@ -81,7 +82,7 @@ def display_messages():
                 """
                 html(copy_button_html, height=50)
 
-            if i == len(st.session_state.messages) - 1 and message["role"] == "CHATBOT":
+            if i == len(st.session_state.messages) - 1 and message["role"] == "assistant":
                 with col4:
                     if st.button("Retry", key=f"retry_{i}_{len(st.session_state.messages)}"):
                         if len(st.session_state.messages) >= 2:
@@ -95,7 +96,7 @@ def display_messages():
                 col1, col2 = st.columns([1, 1])
                 with col1:
                     if st.form_submit_button("Save"):
-                        st.session_state.messages[i]["text"] = new_content
+                        st.session_state.messages[i]["content"] = new_content
                         del st.session_state.edit_index
                         st.rerun()
                 with col2:
@@ -105,9 +106,9 @@ def display_messages():
     
     if "retry_flag" in st.session_state and st.session_state.retry_flag == True:
         if len(st.session_state.messages) > 0:  # メッセージリストが空でないことを確認
-            prompt = st.session_state.messages[-1]["text"]
+            prompt = st.session_state.messages[-1]["content"]
             response = get_ai_response(prompt, st.session_state.messages[:-1])
-            st.session_state.messages.append({"role": "CHATBOT", "text": response})
+            st.session_state.messages.append({"role": "assistant", "content": response})
             st.session_state.retry_flag = False
             st.rerun()
         else:
@@ -118,12 +119,12 @@ with st.sidebar:
     # Copy Conversation History button
     log_text = ""
     for message in st.session_state.messages:
-        if message["role"] == "USER":
+        if message["role"] == "user":
             log_text += "<USER>\n"
-            log_text += message["text"] + "\n\n"
+            log_text += message["content"] + "\n\n"
         else:
             log_text += "<ASSISTANT>\n"
-            log_text += message["text"] + "\n\n"
+            log_text += message["content"] + "\n\n"
     log_text = log_text.rstrip("\n")
         
     # PythonでURLエンコード
@@ -144,11 +145,10 @@ with st.sidebar:
     html(copy_log_button_html, height=50)
     
     st.header("Advanced Settings")
-    model = st.selectbox("Model", options=["command-r-plus", "command-r"], index=0)
-    preamble = st.text_area("Preamble", height=100)
+    model = st.selectbox("Model", options=["claude-3-haiku", "command-r-plus"], index=0)
+    system_prompt = st.text_area("System prompt", height=100)
     temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.3, step=0.1)
-    k = st.slider("Top-K", min_value=0, max_value=500, value=0, step=1)
-    p = st.slider("Top-P", min_value=0.01, max_value=0.99, value=0.75, step=0.01)
+    top_p = st.slider("Top-P", min_value=0.01, max_value=0.99, value=0.75, step=0.01)
     penalty_type = st.selectbox("Penalty Type", options=["Frequency Penalty", "Presence Penalty"])
     penalty_value = st.slider("Penalty Value", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
 
@@ -162,9 +162,9 @@ with st.sidebar:
         for message in messages:
             if message.strip() in ["<USER>", "<ASSISTANT>"]:
                 if role and text:
-                    st.session_state.messages.append({"role": role, "text": text.strip()})
+                    st.session_state.messages.append({"role": role, "content": text.strip()})
                     text = ""
-                role = "USER" if message.strip() == "<USER>" else "CHATBOT"
+                role = "user" if message.strip() == "<USER>" else "assistant"
             else:
                 text += message
         if role and text:
@@ -181,7 +181,7 @@ with st.sidebar:
     if st.button("Update API Key"):
         if new_api_key:
             st.session_state.api_key = new_api_key
-            client = cohere.Client(api_key=new_api_key)
+            client = Prem(api_key=new_api_key)
             st.success("API Key updated successfully!")
         else:
             st.warning("Please enter a valid API Key.")
@@ -196,6 +196,6 @@ if prompt := st.chat_input("What is up?"):
 
     response = get_ai_response(prompt, chat_history)
         
-    st.session_state.messages.append({"role": "USER", "text": prompt})
-    st.session_state.messages.append({"role": "CHATBOT", "text": response})
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.session_state.messages.append({"role": "assistant", "content": response})
     st.rerun()