Update app.py

app.py

@@ -1,14 +1,7 @@
 import os
 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'
-    
-if os.name == 'nt':
-    asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
-    
+
+
 import openai
 import pandas as pd
 import streamlit as st
@@ -45,86 +38,71 @@ def oai_key_isvalid(api_key):
 def initialize_session_state():
     """Initialize all session state variables"""
     if 'prompt' not in st.session_state:
-        st.session_state.prompt = None
+        st.session_state['prompt'] = None
     if 'agent' not in st.session_state:
-        st.session_state.agent = None
+        st.session_state['agent'] = None
     if 'api_key' not in st.session_state:
-        st.session_state.api_key = os.getenv('OPENAI_API_KEY', 'none')
+        st.session_state['api_key'] = os.getenv('OPENAI_API_KEY', 'none')
     if 'base_url' not in st.session_state:
-        st.session_state.base_url = ''
+        st.session_state['base_url'] = ''
     if 'model1_select' not in st.session_state:
-        st.session_state.model1_select = ''
+        st.session_state['model1_select'] = ''
     if 'model2_select' not in st.session_state:
-        st.session_state.model2_select = ''
+        st.session_state['model2_select'] = ''
     if 'serp_api' not in st.session_state:
-        st.session_state.serp_api = 'none'
+        st.session_state['serp_api'] = 'none'
     if 'semantic_scholar_url' not in st.session_state:
-        st.session_state.semantic_scholar_url = 'none'
+        st.session_state['semantic_scholar_url'] = 'none'
     if 'file_type' not in st.session_state:
-        st.session_state.file_type = 'None'
+        st.session_state['file_type'] = 'None'
     if 'domain' not in st.session_state:
-        st.session_state.domain = 'Organic solar cell'
-
-# Set width of sidebar BEFORE initializing session state
-st.markdown(
-    """
-    <style>
-    [data-testid="stSidebar"][aria-expanded="true"]{
-        min-width: 500px;
-        max-width: 500px;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True,
-)
-
-# Initialize session state at the very beginning
-initialize_session_state()
-
-# Now safely access session state
-ss = st.session_state
+        st.session_state['domain'] = 'Organic solar cell'
 
 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"
    
     try:
-        ss.agent = TeLLAgent(
+        st.session_state['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'),
+            openai_api_key=st.session_state.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
+        st.session_state['agent'] = None
     
-    return ss.agent
+    return st.session_state.get('agent')
 
 def on_api_key_change():
-    api_key = ss.get('api_key') or os.getenv('OPENAI_API_KEY')
-    # Check if key is valid
+    api_key = st.session_state.get('api_key') or os.getenv('OPENAI_API_KEY')
     if not oai_key_isvalid(api_key):
         st.write("Please input a valid OpenAI API key.")
 
 def run_prompt(prompt, file_path='...', image_path='...'):
     """Execute prompt with agent"""
     try:
-        if ss.get('domain') == 'Drug discovery':
+        if st.session_state.get('domain') == 'Drug discovery':
             agent = instantiate_agent(
-                model1=ss.get('model1_select'), 
-                model2=ss.get('model2_select'), 
+                model1=st.session_state.get('model1_select'), 
+                model2=st.session_state.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'), 
+                model1=st.session_state.get('model1_select'), 
+                model2=st.session_state.get('model2_select'), 
                 file_path=file_path, 
                 image_path=image_path
             )
@@ -138,7 +116,7 @@ def run_prompt(prompt, file_path='...', image_path='...'):
         with st.chat_message("assistant"):
             try:
                 response = agent.run(prompt)
-                if ss.get('file_type') == 'CSV (.csv)':
+                if st.session_state.get('file_type') == 'CSV (.csv)':
                     try: 
                         fx = pd.DataFrame(list(response))
                         st.markdown(":red[Prediction finished! ]")
@@ -161,173 +139,198 @@ def run_prompt(prompt, file_path='...', image_path='...'):
     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',
-    '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.',
-        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',
-        label_visibility="collapsed"
-    ) 
+def main():
+    """Main application function"""
+    initialize_session_state()
+# 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'
     
-    # Input model to use
-    st.text_input(
-        'Input global planning model to use',
-        key='model1_select',
-    )
+    if os.name == 'nt':
+        asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
     
-    st.text_input(
-        'Input local execution model to use',
-        key='model2_select',
+    # Set width of sidebar
+    st.markdown(
+        """
+        <style>
+        [data-testid="stSidebar"][aria-expanded="true"]{
+            min-width: 500px;
+            max-width: 500px;
+        }
+        </style>
+        """,
+        unsafe_allow_html=True,
     )
     
-    st.text_input(
-        'Input SERP API KEY (optional).',
-        placeholder='Input SERP API KEY (optional)',
-        key='serp_api',
-        type='password',
-        label_visibility="collapsed"
-    ) 
+    # Initialize session state
     
-    st.text_input(
-        '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
+    pre_prompts = [
+        'Generate a donor with PCE = 10% ',
+        'The history and development of Y6',
+        'Predict the LogP of PM6',
+        'Predict the PCE of Y6'
+    ]
 
-    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', '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% 🧨',
-            on_click=lambda: run_prompt(pre_prompts[0]),
-        )
-        st.button(
-            r'📚 The history and development of Y6',
-            on_click=lambda: run_prompt(pre_prompts[1]),
+    # 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.',
+            type='password',
+            key='api_key',
+            on_change=on_api_key_change,
+            label_visibility="collapsed"
         )
-    with cols[1]:
-        st.button(
-            r"🎄Predict the LogP of PM6",
-            on_click=lambda: run_prompt(pre_prompts[2]),
+        
+        st.text_input(
+            '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.button(
-            r'💎 Predict the PCE of Y6',
-            on_click=lambda: run_prompt(pre_prompts[3]),
+        
+        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',
+            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',
+            label_visibility="collapsed"
+        ) 
+        
+        user_api_key = st.session_state.get('api_key')
+        user_base_url = st.session_state.get('base_url')
 
-    st.selectbox(
-        '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"])
-    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')
+        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"] = st.session_state.get('serp_api', 'none')
+        os.environ["SEMANTIC_SCHOLAR_API_KEY"] = st.session_state.get('semantic_scholar_url', 'none')
         
-    st.selectbox(
-        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'
+        # 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% 🧨',
+                on_click=lambda: run_prompt(pre_prompts[0]),
             )
-        else:
-            instantiate_agent(
-                model1='gpt-4o-2024-11-20', 
-                model2='gpt-4o-2024-11-20'
+            st.button(
+                r'📚 The history and development of Y6',
+                on_click=lambda: run_prompt(pre_prompts[1]),
             )
-    
-    # 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
+        with cols[1]:
+            st.button(
+                r"🎄Predict the LogP of PM6",
+                on_click=lambda: run_prompt(pre_prompts[2]),
+            )
+            st.button(
+                r'💎 Predict the PCE of Y6',
+                on_click=lambda: run_prompt(pre_prompts[3]),
             )
-        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(f.name), file_path=f.name) 
-                f.close()
-      
-        if ss.get('file_type') == 'Figure (.jpg, .png, .jpeg)':
-            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(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()
-    else:
-        run_prompt(prompt)
+        st.selectbox(
+            'Select the file type',
+            ['None', 'CSV (.csv)', 'Figure (.jpg, .png, .jpeg)', 'PDF (.pdf)'],
+            key='file_type',
+        )
+        
+        uploaded_file = None
+        if st.session_state.get('file_type') == 'Figure (.jpg, .png, .jpeg)':
+            uploaded_file = st.file_uploader("Choose a Figure", type=["jpg", "jpeg", "png"])
+        if st.session_state.get('file_type') == 'PDF (.pdf)':
+            uploaded_file = st.file_uploader("Choose a PDF file")
+        if st.session_state.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 - only initialize agent if not already done
+        if st.session_state.get('agent') is None:
+            if st.session_state.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
+        agent = st.session_state.get('agent')
+        if agent is not None and hasattr(agent, 'agent_executor2'):
+            try:
+                tools = 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"):
+        uploaded_file = None
+        if st.session_state.get('file_type') == 'Figure (.jpg, .png, .jpeg)':
+            # Re-get uploaded file for main area
+            pass
+        elif st.session_state.get('file_type') == 'PDF (.pdf)':
+            pass
+        elif st.session_state.get('file_type') == 'CSV (.csv)':
+            pass
+            
+        # Check if file was uploaded in sidebar
+        # Note: uploaded_file needs to be handled differently in Streamlit
+        # For now, just run the prompt
+        run_prompt(prompt)
+
+# Run the main function
+if __name__ == "__main__":
+    main()