Update app.py

app.py

@@ -1,64 +1,161 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+import os
+
+# Init with fake key
+if 'OPENAI_API_KEY' not in os.environ:
+    os.environ['OPENAI_API_KEY'] = 'none'
+
+import openai
+import pandas as pd
+import streamlit as st
+from IPython.core.display import HTML
+from PIL import Image
+from langchain.callbacks import wandb_tracing_enabled
+from chemcrow.agents import ChemCrow, make_tools
+from chemcrow.frontend.streamlit_callback_handler import \
+    StreamlitCallbackHandlerChem
+from utils import oai_key_isvalid
+
+from dotenv import load_dotenv
+
+load_dotenv()
+ss = st.session_state
+ss.prompt = None
+
+icon = Image.open('assets/logo0.png')
+st.set_page_config(
+    page_title="ChemCrow",
+    page_icon = icon
 )
 
+# Set width of sidebar
+st.markdown(
+    """
+    <style>
+    [data-testid="stSidebar"][aria-expanded="true"]{
+        min-width: 450px;
+        max-width: 450px;
+    }
+    """,
+    unsafe_allow_html=True,
+)
+
+
+def instantiate_agent(model):
+    ss.agent = ChemCrow(
+        model=model,
+        tools_model=model,
+        temp=0.1,
+        openai_api_key=ss.get('api_key'),
+        api_keys={
+            'RXN4CHEM_API_KEY': st.secrets['RXN4CHEM_API_KEY'],
+            'CHEMSPACE_API_KEY': st.secrets['CHEMSPACE_API_KEY']
+        }
+    ).agent_executor
+    return ss.agent
+
+instantiate_agent('gpt-4-0613')
+tools = ss.agent.tools
+
+tool_list = pd.Series(
+    {f"✅ {t.name}":t.description for t in tools}
+).reset_index()
+tool_list.columns = ['Tool', 'Description']
+
+def on_api_key_change():
+    api_key = ss.get('api_key') or os.getenv('OPENAI_API_KEY')
+    # Check if key is valid
+    if not oai_key_isvalid(api_key):
+        st.write("Please input a valid OpenAI API key.")
+
+def run_prompt(prompt):
+    agent = instantiate_agent(ss.get('model_select'))
+    st.chat_message("user").write(prompt)
+    with st.chat_message("assistant"):
+        st_callback = StreamlitCallbackHandlerChem(
+            st.container(),
+            max_thought_containers = 3,
+            collapse_completed_thoughts = False,
+            output_placeholder=ss
+        )
+        try:
+            with wandb_tracing_enabled():
+                response = agent.run(prompt, callbacks=[st_callback])
+                st.write(response)
+        except openai.error.AuthenticationError:
+            st.write("Please input a valid OpenAI API key")
+        except openai.error.APIError:
+            # Handle specific API errors here
+            print("OpenAI API error, please try again!")
+
+
+pre_prompts = [
+    'How can I synthesize safinamide?',
+    (
+        'Predict the product of a mixture of Ethylidenecyclohexane and HBr. '
+        'Then predict the same reaction, adding methyl peroxide into the '
+        'mixture. Compare the two products and explain the reaction mechanism.'
+    ),
+    (
+        'What is the boiling point of the reaction product between '
+        'isoamyl alcohol and acetic acid?'
+    ),
+    'Tell me how to synthesize vanilline, and the price of the precursors.'
+]
+
+# sidebar
+with st.sidebar:
+    chemcrow_logo = Image.open('assets/chemcrow-logo-bold-new.png')
+    st.image(chemcrow_logo)
+
+    # Input OpenAI api key
+    st.text_input(
+        'Input your OpenAI API key.',
+        placeholder = 'Input your OpenAI API key.',
+        type='password',
+        key='api_key',
+        on_change=on_api_key_change,
+        label_visibility="collapsed"
+    )
+
+    # Input model to use
+    st.selectbox(
+        'Select model to use',
+        ['gpt-4-0613', 'gpt-3.5-turbo', 'gpt-4o-mini'],
+        key='model_select',
+    )
+
+    # Display prompt examples
+    st.markdown('# What can I ask?')
+    cols = st.columns(2)
+    with cols[0]:
+        st.button(
+            "How can I synthesize safinamide?",
+            on_click=lambda: run_prompt(pre_prompts[0]),
+        )
+        st.button(
+            "Explain mechanism of bromoaddition reaction",
+            on_click=lambda: run_prompt(pre_prompts[1]),
+        )
+    with cols[1]:
+        st.button(
+            'Predict properties of a reaction product',
+            on_click=lambda: run_prompt(pre_prompts[2]),
+        )
+        st.button(
+            'Synthesize molecule with price of precursors',
+            on_click=lambda: run_prompt(pre_prompts[3]),
+        )
+
+    st.markdown('---')
+    # Display available tools
+    st.markdown(f"# {len(tool_list)} available tools")
+    st.dataframe(
+        tool_list,
+        use_container_width=True,
+        hide_index=True,
+        height=200
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+# Execute agent on user input
+if user_input := st.chat_input():
+    run_prompt(user_input)