Update app.py

app.py

@@ -1,7 +1,5 @@
 import os
 import asyncio
-
-
 import openai
 import pandas as pd
 import streamlit as st
@@ -16,6 +14,8 @@ import tempfile
 
 load_dotenv()
 
+# --- Helper Functions ---
+
 def convert_to_base64(pil_image):
     buffered = BytesIO()
     pil_image.save(buffered, format="PNG")
@@ -25,119 +25,128 @@ def convert_to_base64(pil_image):
 def oai_key_isvalid(api_key):
     """Check if a given OpenAI key is valid"""
     try:
+        # Simple check without invoking expensive calls if possible, 
+        # or keep the invoke if you want strict validation.
+        if not api_key or api_key == 'none': return False
+        
         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")
+        llm.invoke("Hi")
         return True
     except:
         return False
 
 def initialize_session_state():
     """Initialize all session state variables"""
-    if 'prompt' not in st.session_state:
-        st.session_state['prompt'] = None
-    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'
-
-def instantiate_agent(model1, model2, file_path='...', image_path='...', tools=None):
-    """Create or update agent instance"""
+    # UI States
+    if 'prompt' not in st.session_state: st.session_state['prompt'] = 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'
     
-    try:
-        st.session_state['agent'] = TeLLAgent(
-            tools=tools,
-            model1=model1,
-            model2=model2,
-            tools_model='gpt-4o-2024-11-20',
-            temp=0.1,
-            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)}")
-        st.session_state['agent'] = None
+    # Config States (Defaults)
+    if 'model1_select' not in st.session_state: st.session_state['model1_select'] = 'deepseek-ai/DeepSeek-R1'
+    if 'model2_select' not in st.session_state: st.session_state['model2_select'] = 'deepseek-ai/DeepSeek-v3.1'
+    if 'api_key' not in st.session_state: st.session_state['api_key'] = os.getenv('OPENAI_API_KEY', '')
     
-    return st.session_state.get('agent')
+    # Agent & Logic States
+    if 'agent' not in st.session_state: st.session_state['agent'] = None
+    if 'last_agent_config' not in st.session_state: st.session_state['last_agent_config'] = {}
+    
+    # File Paths
+    if 'current_file_path' not in st.session_state: st.session_state['current_file_path'] = '...'
+    if 'current_image_path' not in st.session_state: st.session_state['current_image_path'] = '...'
 
-def on_api_key_change():
-    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 update_agent_if_needed(current_config):
+    """
+    Compare current_config with the last used config.
+    If changed, instantiate the agent automatically.
+    """
+    last_config = st.session_state.get('last_agent_config', {})
+    
+    # Check if config changed or agent is missing
+    if st.session_state['agent'] is None or current_config != last_config:
+        
+        # Validations
+        if not current_config['api_key'] or current_config['api_key'] == 'none':
+            # Don't try to init if no key, just return
+            return 
 
-def run_prompt(prompt, file_path='...', image_path='...'):
-    """Execute prompt with agent"""
-    try:
-        if st.session_state.get('domain') == 'Drug discovery':
-            agent = instantiate_agent(
-                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=st.session_state.get('model1_select'), 
-                model2=st.session_state.get('model2_select'), 
-                file_path=file_path, 
+        model1 = current_config['model1'] or "deepseek-ai/DeepSeek-R1"
+        model2 = current_config['model2'] or "deepseek-ai/DeepSeek-v3.1"
+        domain = current_config['domain']
+        file_path = current_config['file_path']
+        image_path = current_config['image_path']
+        
+        tools_setting = 'drug' if domain == 'Drug discovery' else None
+
+        try:
+            # Create the agent
+            new_agent = TeLLAgent(
+                tools=tools_setting,
+                model1=model1,
+                model2=model2,
+                tools_model='gpt-4o-2024-11-20',
+                temp=0.1,
+                openai_api_key=current_config['api_key'],
+                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 st.session_state.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.session_state['agent'] = new_agent
+            st.session_state['last_agent_config'] = current_config # Update history
+            
+            # Optional: Toast notification to show user it updated
+            # st.toast(f"Agent updated! (Model: {model1})", icon="🤖")
+            
+        except Exception as e:
+            st.error(f"Auto-initialization failed: {str(e)}")
+            st.session_state['agent'] = None
+
+def run_prompt(prompt):
+    """Execute prompt using the already initialized agent"""
+    agent = st.session_state.get('agent')
+    
+    if agent is None:
+        st.error("Agent is not ready. Please check your API Key and Model settings.")
+        return
+            
+    st.chat_message("user").write(prompt)
+    
+    with st.chat_message("assistant"):
+        try:
+            response = agent.run(prompt)
+            if st.session_state.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)
-            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)}")
+            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)}")
+
+def on_api_key_change():
+    """Callback when API key changes"""
+    # Logic is now handled by auto-update, but we can do validation here
+    pass
 
 def main():
-    """Main application function"""
     initialize_session_state()
-# Init with fake key
+    
+    # Environment Defaults
     if 'OPENAI_API_KEY' not in os.environ:
         os.environ['OPENAI_API_KEY'] = 'none'
         os.environ["SERP_API_KEY"] = 'none'
@@ -146,186 +155,129 @@ def main():
     if os.name == 'nt':
         asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
     
-    # Set width of sidebar
+    # Sidebar Setup
     st.markdown(
         """
         <style>
         [data-testid="stSidebar"][aria-expanded="true"]{
-            min-width: 500px;
-            max-width: 500px;
+            min-width: 500px; max-width: 500px;
         }
         </style>
         """,
         unsafe_allow_html=True,
     )
     
-    # Initialize session state
-    
-    
-    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
+    # --- Sidebar Inputs ---
     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"
-        )
+        # 1. Credentials
+        api_key_input = st.text_input('Input your OpenAI API key.', type='password', key='api_key', label_visibility="collapsed")
+        base_url_input = st.text_input('Input base url (optional).', key='base_url', type='password', label_visibility="collapsed") 
         
-        st.text_input(
-            'Input base url (optional).',
-            placeholder='Input base url (optional)',
-            key='base_url',
-            type='password',
-            label_visibility="collapsed"
-        ) 
+        # 2. Model Configuration
+        # Defaults are set in initialize_session_state, user can change them here
+        model1_input = st.text_input('Input global planning model to use', key='model1_select')
+        model2_input = st.text_input('Input local execution model to use', key='model2_select')
         
-        # Input model to use
-        st.text_input(
-            'Input global planning model to use',
-            key='model1_select',
-        )
+        # 3. External Tools Keys
+        serp_input = st.text_input('Input SERP API KEY (optional).', key='serp_api', type='password', label_visibility="collapsed") 
+        scholar_input = st.text_input('Input SEMANTIC SCHOLAR API KEY (optional).', key='semantic_scholar_url', type='password', label_visibility="collapsed") 
         
-        st.text_input(
-            'Input local execution model to use',
-            key='model2_select',
-        )
+        # Update OS Environ
+        if api_key_input: os.environ['OPENAI_API_KEY'] = api_key_input
+        if base_url_input and base_url_input != "none": os.environ["OPENAI_API_BASE"] = base_url_input
+        elif "OPENAI_API_BASE" in os.environ: del os.environ["OPENAI_API_BASE"]
+        os.environ["SERP_API_KEY"] = serp_input if serp_input else 'none'
+        os.environ["SEMANTIC_SCHOLAR_API_KEY"] = scholar_input if scholar_input else 'none'
         
-        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')
-
-        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')
-        
-        # 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]),
-            )
-        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]),
-            )
-
-        st.selectbox(
-            'Select the file type',
-            ['None', 'CSV (.csv)', 'Figure (.jpg, .png, .jpeg)', 'PDF (.pdf)'],
-            key='file_type',
-        )
+        # 4. Domain & Files
+        st.selectbox(r'📚 Choose the domain', ['Organic solar cell', 'Drug discovery'], key='domain')
+        st.selectbox('Select the file type', ['None', 'CSV (.csv)', 'Figure (.jpg, .png, .jpeg)', 'PDF (.pdf)'], key='file_type')
         
+        # --- File Processing (Immediate) ---
+        # We process files HERE so the paths are ready for the Agent Auto-Init immediately after.
         uploaded_file = None
-        if st.session_state.get('file_type') == 'Figure (.jpg, .png, .jpeg)':
+        current_file_path = '...'
+        current_image_path = '...'
+        
+        if st.session_state['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)':
+            if uploaded_file:
+                st.image(uploaded_file, width=200)
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp:
+                    img_str = base64.b64encode(uploaded_file.getvalue()).decode("utf-8")
+                    temp.write(base64.b64decode(img_str))
+                    current_image_path = temp.name
+                    st.session_state['current_image_path'] = temp.name # Persist
+
+        elif st.session_state['file_type'] == 'PDF (.pdf)':
             uploaded_file = st.file_uploader("Choose a PDF file")
-        if st.session_state.get('file_type') == 'CSV (.csv)':
+            if uploaded_file:
+                with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as f:
+                    f.write(uploaded_file.read())
+                    current_file_path = f.name
+                    st.session_state['current_file_path'] = f.name
+
+        elif st.session_state['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',
-        )
+            if uploaded_file:
+                with tempfile.NamedTemporaryFile(suffix='.csv', delete=False) as f:
+                    f.write(uploaded_file.read())
+                    current_file_path = f.name
+                    st.session_state['current_file_path'] = f.name
         
-        # 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'
-                )
+        # --- AUTOMATIC AGENT INITIALIZATION ---
+        # Collect current config into a dict
+        current_config = {
+            'api_key': api_key_input,
+            'base_url': base_url_input,
+            'model1': model1_input,
+            'model2': model2_input,
+            'domain': st.session_state['domain'],
+            'file_path': current_file_path,   # Uses the temp path created just now
+            'image_path': current_image_path  # Uses the temp path created just now
+        }
         
-        # Safely display tools
+        # Trigger the check & init logic
+        update_agent_if_needed(current_config)
+
+        # Show Tools status
         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 = 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)}")
+                st.dataframe(tool_list, use_container_width=True, hide_index=True, height=150)
+            except: pass
         else:
-            st.info("Agent not initialized. Please check your API key.")
+            if not api_key_input:
+                st.warning("🔑 Waiting for API Key...")
+        
+        # Pre-defined Prompts
+ 
+        st.markdown('# What can I ask?')
+        
+        pre_prompts = [
+            'Generate a donor with PCE = 10% ',
+            'The history and development of Y6',
+            'Predict the LogP of PM6',
+            'Predict the PCE of Y6'
+        ]
+        
+        cols = st.columns(2)
+        with cols[0]:
+            st.button(r'👑 Generate a donor...', on_click=lambda: run_prompt(pre_prompts[0]))
+            st.button(r'📚 History of Y6', on_click=lambda: run_prompt(pre_prompts[1]))
+        with cols[1]:
+            st.button(r"🎄Predict LogP of PM6", on_click=lambda: run_prompt(pre_prompts[2]))
+            st.button(r'💎 Predict PCE of Y6', on_click=lambda: run_prompt(pre_prompts[3]))
 
-    # 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
+    # --- Main Chat Area ---
+    if prompt := st.chat_input("Say something..."):
+        # Append context about file if needed, though Agent handles files via init args mostly.
         run_prompt(prompt)
 
-# Run the main function
 if __name__ == "__main__":
     main()