Update app.py

app.py

@@ -3,30 +3,25 @@ import asyncio
 # Init with fake key
 if 'OPENAI_API_KEY' not in os.environ:
     os.environ['OPENAI_API_KEY'] = 'none'
+    os.environ["SERP_API_KEY"] = 'none'
+    os.environ["SEMANTIC_SCHOLAR_API_KEY"] = 'none'
     
-   
-    os.environ["SERP_API_KEY"] =  'none'
-    os.environ["SEMANTIC_SCHOLAR_API_KEY"] =  'none'
 if os.name == 'nt':
     asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
     
 import openai
 import pandas as pd
 import streamlit as st
- 
 from PIL import Image
-from  agent import TeLLAgent, make_tools
-from streamlit_callback_handler import \
-    StreamlitCallbackHandlerChem
+from agent import TeLLAgent, make_tools
+from streamlit_callback_handler import StreamlitCallbackHandlerChem
 import base64 
-import pandas as pd
 from dotenv import load_dotenv
-from langchain_openai import ChatOpenAI , OpenAI
-import base64
+from langchain_openai import ChatOpenAI, OpenAI
 from io import BytesIO
-from PIL import Image
 import tempfile
- 
+
+load_dotenv()
 
 def convert_to_base64(pil_image):
     buffered = BytesIO()
@@ -38,29 +33,39 @@ def oai_key_isvalid(api_key):
     """Check if a given OpenAI key is valid"""
     try:
         if os.getenv("OPENAI_API_BASE"):
-            llm = ChatOpenAI(openai_api_key = api_key, base_url=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)
+            llm = ChatOpenAI(openai_api_key=api_key)
             out = llm.invoke("This is a test")
         return True
     except:
         return False
-    
 
 def initialize_session_state():
-    """Initialize session state variables"""
+    """Initialize all session state variables"""
     if 'prompt' not in st.session_state:
         st.session_state.prompt = None
-    # Add other session state initializations here as needed
-
-# Call initialization function
-initialize_session_state()
+    if 'agent' not in st.session_state:
+        st.session_state.agent = None
+    if 'api_key' not in st.session_state:
+        st.session_state.api_key = os.getenv('OPENAI_API_KEY', 'none')
+    if 'base_url' not in st.session_state:
+        st.session_state.base_url = ''
+    if 'model1_select' not in st.session_state:
+        st.session_state.model1_select = ''
+    if 'model2_select' not in st.session_state:
+        st.session_state.model2_select = ''
+    if 'serp_api' not in st.session_state:
+        st.session_state.serp_api = 'none'
+    if 'semantic_scholar_url' not in st.session_state:
+        st.session_state.semantic_scholar_url = 'none'
+    if 'file_type' not in st.session_state:
+        st.session_state.file_type = 'None'
+    if 'domain' not in st.session_state:
+        st.session_state.domain = 'Organic solar cell'
 
-# Now you can safely use session state
-ss = st.session_state
- 
-# Set width of sidebar
+# Set width of sidebar BEFORE initializing session state
 st.markdown(
     """
     <style>
@@ -68,150 +73,185 @@ st.markdown(
         min-width: 500px;
         max-width: 500px;
     }
+    </style>
     """,
     unsafe_allow_html=True,
 )
 
+# Initialize session state at the very beginning
+initialize_session_state()
 
-def instantiate_agent(model1, model2, file_path = '...',  image_path ='...',  tools=None):
-    
+# Now safely access session state
+ss = st.session_state
+
+def instantiate_agent(model1, model2, file_path='...', image_path='...', tools=None):
+    """Create or update agent instance"""
     if not model1:
         model1 = "deepseek-ai/DeepSeek-R1"
     
-     
     if not model2:
-        model2 =  "deepseek-ai/DeepSeek-v3.1"
+        model2 = "deepseek-ai/DeepSeek-v3.1"
    
-
-    ss.agent = TeLLAgent( tools=tools,
-        model1 = model1,
-        model2 = model2,
-        tools_model='gpt-4o-2024-11-20',
-        temp=0.1,
-        openai_api_key=ss.get('api_key'),
-        file_path = file_path,
-        image_path =image_path)
+    try:
+        ss.agent = TeLLAgent(
+            tools=tools,
+            model1=model1,
+            model2=model2,
+            tools_model='gpt-4o-2024-11-20',
+            temp=0.1,
+            openai_api_key=ss.get('api_key'),
+            file_path=file_path,
+            image_path=image_path
+        )
+    except Exception as e:
+        st.error(f"Error instantiating agent: {str(e)}")
+        ss.agent = None
+    
     return ss.agent
 
-
 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, file_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:
-           
-           response = agent.run(prompt)
-           if ss.get('file_type') == 'CSV (.csv)':
-               try: 
-                   fx = pd.DataFrame(list(response))
-                   st.markdown(":red[Prediction finished! ]")	
-                   st.download_button( "⬇️Download  the predicted files as .csv", fx.to_csv(), "predict results.csv", use_container_width=True)
-               except:
-                   st.write(response)
-           else:
-                st.write(response)
-        except openai.AuthenticationError:
-            st.write("Please input a valid OpenAI API key")
-        except openai.APIError:
-            # Handle specific API errors here
-            print("OpenAI API error, please try again!")
-    
+def run_prompt(prompt, file_path='...', image_path='...'):
+    """Execute prompt with agent"""
+    try:
+        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
+            )
+        
+        if agent is None:
+            st.error("Failed to create agent. Please check your API key and settings.")
+            return
+            
+        st.chat_message("user").write(prompt)
+        
+        with st.chat_message("assistant"):
+            try:
+                response = agent.run(prompt)
+                if ss.get('file_type') == 'CSV (.csv)':
+                    try: 
+                        fx = pd.DataFrame(list(response))
+                        st.markdown(":red[Prediction finished! ]")
+                        st.download_button(
+                            "⬇️Download the predicted files as .csv", 
+                            fx.to_csv(), 
+                            "predict results.csv", 
+                            use_container_width=True
+                        )
+                    except:
+                        st.write(response)
+                else:
+                    st.write(response)
+            except openai.AuthenticationError:
+                st.write("Please input a valid OpenAI API key")
+            except openai.APIError as e:
+                st.error(f"OpenAI API error: {str(e)}")
+            except Exception as e:
+                st.error(f"Error running prompt: {str(e)}")
+    except Exception as e:
+        st.error(f"Error in run_prompt: {str(e)}")
 
 pre_prompts = [
     'Generate a donor with PCE = 10% ',
-    ('The history and development of Y6'
-       
-    ),
-    (
-        'Predict the LogP of PM6'
-    ),
+    'The history and development of Y6',
+    'Predict the LogP of PM6',
     'Predict the PCE of Y6'
 ]
 
 # sidebar
 with st.sidebar:
-
     st.header("🤖 :blue[TeLLAgent] ")
+    
     # Input OpenAI api key
     st.text_input(
         'Input your OpenAI API key.',
-        placeholder = '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"
     )
+    
     st.text_input(
-        'Input  base url (optional).',
-        placeholder = 'Input  base url (optional)',
-        key='base_url',type='password',
+        'Input base url (optional).',
+        placeholder='Input base url (optional)',
+        key='base_url',
+        type='password',
         label_visibility="collapsed"
     ) 
+    
     # Input model to use
     st.text_input(
         'Input global planning model to use',
-        
         key='model1_select',
     )
+    
     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',
+        'Input SERP API KEY (optional).',
+        placeholder='Input SERP API KEY (optional)',
+        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',
+        'Input SEMANTIC SCHOLAR API KEY (optional).',
+        placeholder='Input SEMANTIC SCHOLAR API KEY (optional)',
+        key='semantic_scholar_url',
+        type='password',
         label_visibility="collapsed"
     ) 
+    
     user_api_key = ss.get('api_key')
     user_base_url = ss.get('base_url')
 
-
     if user_api_key:
         os.environ['OPENAI_API_KEY'] = user_api_key
 
     if user_base_url and user_base_url.strip() != "" and user_base_url != "none":
         os.environ["OPENAI_API_BASE"] = user_base_url
     else:
-      
         if "OPENAI_API_BASE" in os.environ:
             del os.environ["OPENAI_API_BASE"]
    
-    os.environ["SERP_API_KEY"] =  ss.get('serp_api')
-    os.environ["SEMANTIC_SCHOLAR_API_KEY"] =  ss.get('semantic_scholar_url')
+    os.environ["SERP_API_KEY"] = ss.get('serp_api', 'none')
+    os.environ["SEMANTIC_SCHOLAR_API_KEY"] = ss.get('semantic_scholar_url', 'none')
     
     # Display prompt examples
     st.markdown('# What can I ask?')
     cols = st.columns(2)
     with cols[0]:
         st.button(
-           r'👑  Generate a donor with PCE = 10%     🧨     ',
+            r'👑 Generate a donor with PCE = 10% 🧨',
             on_click=lambda: run_prompt(pre_prompts[0]),
         )
         st.button(
-            r'📚 The history and development of Y6     ',
+            r'📚 The history and development of Y6',
             on_click=lambda: run_prompt(pre_prompts[1]),
         )
     with cols[1]:
         st.button(
-            r"🎄Predict the LogP of PM6  ",
+            r"🎄Predict the LogP of PM6",
             on_click=lambda: run_prompt(pre_prompts[2]),
         )
         st.button(
@@ -220,70 +260,79 @@ with st.sidebar:
         )
 
     st.selectbox(
-            'Select the file type ',
-            ['None', 'CSV (.csv)', 'Figure (.jpg, .png, .jpeg)', 'PDF (.pdf)'],
-            key='file_type',
-        )
+        'Select the file type',
+        ['None', 'CSV (.csv)', 'Figure (.jpg, .png, .jpeg)', 'PDF (.pdf)'],
+        key='file_type',
+    )
+    
     uploaded_file = None
     if ss.get('file_type') == 'Figure (.jpg, .png, .jpeg)':
-        uploaded_file = st.file_uploader("Choose a Figure", type = ["jpg", "jpeg", "png"])
+        uploaded_file = st.file_uploader("Choose a Figure", type=["jpg", "jpeg", "png"])
     if ss.get('file_type') == 'PDF (.pdf)':
         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')   
+        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,
-        use_container_width=True,
-        hide_index=True,
-        height=200
+        r'📚 Choose the domain',
+        ['Organic solar cell', 'Drug discovery'], 
+        key='domain',
     )
+    
+    # Display available tools - only initialize agent if not already done
+    if ss.agent is None:
+        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'
+            )
+    
+    # Safely display tools
+    if ss.agent is not None and hasattr(ss.agent, 'agent_executor2'):
+        try:
+            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,
+                use_container_width=True,
+                hide_index=True,
+                height=200
+            )
+        except Exception as e:
+            st.warning(f"Could not load tools: {str(e)}")
+    else:
+        st.info("Agent not initialized. Please check your API key.")
 
 # Execute agent on user input
 if prompt := st.chat_input("Say something and/or attach files"):
-
     if uploaded_file is not None:
         if ss.get('file_type') == 'CSV (.csv)':
-            with tempfile.NamedTemporaryFile(  suffix ='.csv' ,delete=False) as f:
-                 f.write(uploaded_file.read())
-                 run_prompt(prompt + str(' ') +  str(f.name), file_path =  f.name) 
-                 f.close()
+            with tempfile.NamedTemporaryFile(suffix='.csv', delete=False) as f:
+                f.write(uploaded_file.read())
+                run_prompt(prompt + ' ' + str(f.name), file_path=f.name) 
+                f.close()
       
         if ss.get('file_type') == 'Figure (.jpg, .png, .jpeg)':
-    
-            st.image(uploaded_file, width = 500)
+            st.image(uploaded_file, width=500)
             with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp:
-    
-                   mg_str = base64.b64encode(uploaded_file.getvalue()).decode("utf-8")
-                   temp.write(base64.b64decode(mg_str))
-             
-            run_prompt(prompt+ str(' ') + str(temp.name), image_path = temp.name )
+                mg_str = base64.b64encode(uploaded_file.getvalue()).decode("utf-8")
+                temp.write(base64.b64decode(mg_str))
+            run_prompt(prompt + ' ' + str(temp.name), image_path=temp.name)
    
         if ss.get('file_type') == 'PDF (.pdf)':
-                with tempfile.NamedTemporaryFile(  suffix ='.pdf' ,delete=False) as f:
-                     f.write(uploaded_file.read())
-                     run_prompt(prompt, file_path =  f.name) 
-                     f.close()
-      
-                    # with open("input.png","wb") as af:
-                    #              mg_str = base64.b64encode(files.getvalue()).decode("utf-8")
-                    #              af.write(base64.b64decode(mg_str))        
-                     
-                    # run_prompt(prompt.text+str(f.name), image_path =f.name )
-                # except:
-                #     st.markdown("Please input correct files or query ")
+            with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as f:
+                f.write(uploaded_file.read())
+                run_prompt(prompt, file_path=f.name) 
+                f.close()
     else:
-          run_prompt(prompt)
+        run_prompt(prompt)