Update app.py

app.py

@@ -37,8 +37,12 @@ def convert_to_base64(pil_image):
 def oai_key_isvalid(api_key):
     """Check if a given OpenAI key is valid"""
     try:
-        llm = ChatOpenAI(openai_api_key = api_key, base_url="https://www.dmxapi.com/v1/")
-        out = llm.invoke("This is a test")
+        if os.getenv("OPENAI_API_BASE"):
+            llm = ChatOpenAI(openai_api_key = api_key, base_url=os.getenv("OPENAI_API_BASE"))
+            out = llm.invoke("This is a test")
+        else:
+            llm = ChatOpenAI(openai_api_key = api_key)
+            out = llm.invoke("This is a test")
         return True
     except:
         return False
@@ -60,8 +64,8 @@ st.markdown(
 )
 
 
-def instantiate_agent(model1, model2, file_path = '...',  image_path ='...'):
-    ss.agent = TeLLAgent(
+def instantiate_agent(model1, model2, file_path = '...',  image_path ='...',  tools=None):
+    ss.agent = TeLLAgent( tools=tools,
         model1 = model1,
         model2 = model2,
         tools_model='gpt-4o-2024-11-20',
@@ -70,11 +74,7 @@ def instantiate_agent(model1, model2, file_path = '...',  image_path ='...'):
                 image_path =image_path)
     return ss.agent
 
-instantiate_agent(model1 = 'gpt-4o-2024-11-20', model2 = 'gpt-4o-2024-11-20')
-tools = ss.agent.agent_executor2.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')
@@ -84,7 +84,10 @@ def on_api_key_change():
         st.write("Please input a valid OpenAI API key.")
 
 def run_prompt(prompt, file_path = '...', image_path = '...'):
-    agent = instantiate_agent(model1 = ss.get('model1_select'), model2 = ss.get('model2_select'), file_path = file_path, image_path =image_path)
+    if ss.get('domain') =='Drug discovery':
+        agent = instantiate_agent(model1 = ss.get('model1_select'), model2 = ss.get('model2_select'), file_path = file_path, image_path =image_path, tools = 'drug')
+    else:
+        agent = instantiate_agent(model1 = ss.get('model1_select'), model2 = ss.get('model2_select'), file_path = file_path, image_path =image_path)
     st.chat_message("user").write(prompt)
     with st.chat_message("assistant") :
         try:
@@ -107,7 +110,7 @@ def run_prompt(prompt, file_path = '...', image_path = '...'):
     
 
 pre_prompts = [
-    'Who are you?',
+    'Generate a donor with PCE = 10% ',
     ('The history and development of Y6'
        
     ),
@@ -133,30 +136,30 @@ with st.sidebar:
     st.text_input(
         'Input  base url (optional).',
         placeholder = 'Input  base url (optional)',
-        key='base_url', type='password',
+        key='base_url',type='password',
         label_visibility="collapsed"
     ) 
     # Input model to use
-    st.selectbox(
-        'Select global planning model to use',
-        ['deepseek-r1-250528'],
+    st.text_input(
+        'Input global planning model to use',
+        
         key='model1_select',
     )
-    st.selectbox(
-        'Select local execution model to use',
-        [ 'deepseek-v3-241226', 'gpt-4o-2024-11-20'],
+    st.text_input(
+        'Input local execution model to use',
+     
         key='model2_select',
     )
     st.text_input(
         'Input  SERP API KEY (optional).',
         placeholder = 'Input SERP API KEY (optional)',
-        key='serp_api', type='password',
+        key='serp_api',type='password',
         label_visibility="collapsed"
     ) 
     st.text_input(
         'Input  SEMANTIC SCHOLAR API KEY (optional).',
         placeholder = 'Input SEMANTIC SCHOLAR API KEY (optional)',
-        key='semantic_scholar_url', type='password',
+        key='semantic_scholar_url',type='password',
         label_visibility="collapsed"
     ) 
     os.environ['OPENAI_API_KEY'] = ss.get('api_key')
@@ -170,7 +173,7 @@ with st.sidebar:
     cols = st.columns(2)
     with cols[0]:
         st.button(
-           r'👑  Who     are    you  ?  🧨               ',
+           r'👑  Generate a donor with PCE = 10%     🧨     ',
             on_click=lambda: run_prompt(pre_prompts[0]),
         )
         st.button(
@@ -199,8 +202,19 @@ with st.sidebar:
         uploaded_file = st.file_uploader("Choose a PDF file")
     if ss.get('file_type') == 'CSV (.csv)':
         uploaded_file = st.file_uploader("Choose a csv file", type = 'csv')   
-        
+    st.selectbox(
+            r'📚 Choose the domain   ',
+            ['Organic solar cell', 'Drug discovery'], key='domain',
+        )  
     # Display available tools
+    if ss.get('domain') == 'Drug discovery':
+        instantiate_agent(model1 = 'gpt-4o-2024-11-20', model2 = 'gpt-4o-2024-11-20' ,tools = 'drug')
+    else:
+        instantiate_agent(model1 = 'gpt-4o-2024-11-20', model2 = 'gpt-4o-2024-11-20' )
+    tools = ss.agent.agent_executor2.tools
+
+    tool_list = pd.Series(   {f"✅ {t.name}": t.description for t in tools}).reset_index()
+    tool_list.columns = ['Tool', 'Description']
     st.markdown(f"# {len(tool_list)} available tools")
     st.dataframe(
         tool_list,