Update app.py

app.py

@@ -1,20 +1,33 @@
+ 
 import os
 import asyncio
+# Init with fake key
+if 'OPENAI_API_KEY' not in os.environ:
+    os.environ['OPENAI_API_KEY'] = 'none'
+    os.environ["OPENAI_API_BASE"] = '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
+from langchain_openai import ChatOpenAI , OpenAI
+import base64
 from io import BytesIO
+from PIL import Image
 import tempfile
-
-load_dotenv()
-
-# --- Helper Functions ---
+ 
 
 def convert_to_base64(pil_image):
     buffered = BytesIO()
@@ -25,259 +38,225 @@ 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"))
+            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.invoke("Hi")
+            llm = ChatOpenAI(openai_api_key = api_key)
+            out = llm.invoke("This is a test")
         return True
     except:
         return False
-
-def initialize_session_state():
-    """Initialize all session state variables"""
-    # 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'
     
-    # Config States (Defaults)
-    if 'model1_select' not in st.session_state: st.session_state['model1_select'] = 'deepseek-reasoner'
-    if 'model2_select' not in st.session_state: st.session_state['model2_select'] = 'deepseek-reasoner'
-    if 'api_key' not in st.session_state: st.session_state['api_key'] = os.getenv('OPENAI_API_KEY', '')
-    
-    # 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 update_agent_if_needed(current_config):
-    """
-    Compare current_config with the last used config.
-    If changed, instantiate the agent automatically.
+load_dotenv()
+ss = st.session_state
+ss.prompt = None
+ 
+# Set width of sidebar
+st.markdown(
     """
-    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 
+    <style>
+    [data-testid="stSidebar"][aria-expanded="true"]{
+        min-width: 500px;
+        max-width: 500px;
+    }
+    """,
+    unsafe_allow_html=True,
+)
 
-        model1 = current_config['model1'] or "deepseek-reasoner"
-        model2 = current_config['model2'] or "deepseek-reasoner"
-        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
-            )
-            
-            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 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',
+        temp=0.1,
+        openai_api_key=ss.get('api_key')    ,        file_path = file_path,
+                image_path =image_path)
+    return ss.agent
 
-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
-            
+
+
+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"):
+    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:
+           
+           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 openai.APIError:
+            # Handle specific API errors here
+            print("OpenAI API error, please try again!")
+    
 
-def on_api_key_change():
-    """Callback when API key changes"""
-    # Logic is now handled by auto-update, but we can do validation here
-    pass
+pre_prompts = [
+    'Generate a donor with PCE = 10% ',
+    ('The history and development of Y6'
+       
+    ),
+    (
+        'Predict the LogP of PM6'
+    ),
+    'Predict the PCE of Y6'
+]
 
-def main():
-    initialize_session_state()
-    
-    # Environment Defaults
-    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())
-    
-    # Sidebar Setup
-    st.markdown(
-        """
-        <style>
-        [data-testid="stSidebar"][aria-expanded="true"]{
-            min-width: 500px; max-width: 500px;
-        }
-        </style>
-        """,
-        unsafe_allow_html=True,
+# 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"
     )
-    
-    # --- Sidebar Inputs ---
-    with st.sidebar:
-        st.header("🤖 :blue[TeLLAgent] ")
-        
-        # 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") 
-        
-        # 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')
-        
-        # 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") 
-        
-        # 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'
-        
-        # 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
-        current_file_path = '...'
-        current_image_path = '...'
+    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',
         
-        if st.session_state['file_type'] == 'Figure (.jpg, .png, .jpeg)':
-            uploaded_file = st.file_uploader("Choose a Figure", type=["jpg", "jpeg", "png"])
-            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
+        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',
+        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"
+    ) 
+    os.environ['OPENAI_API_KEY'] = ss.get('api_key')
+    os.environ["OPENAI_API_BASE"] = ss.get('base_url')
+   
+    os.environ["SERP_API_KEY"] =  ss.get('serp_api')
+    os.environ["SEMANTIC_SCHOLAR_API_KEY"] =  ss.get('semantic_scholar_url')
+    
+    # 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]),
+        )
 
-        elif st.session_state['file_type'] == 'PDF (.pdf)':
-            uploaded_file = st.file_uploader("Choose a PDF file")
-            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
+    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')   
+    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
 
-        elif st.session_state['file_type'] == 'CSV (.csv)':
-            uploaded_file = st.file_uploader("Choose a csv file", type='csv')
-            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
-        
-        # --- 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
-        }
-        
-        # Trigger the check & init logic
-        update_agent_if_needed(current_config)
+    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
+    )
 
-        # 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.columns = ['Tool', 'Description']
-                st.dataframe(tool_list, use_container_width=True, hide_index=True, height=150)
-            except: pass
-        else:
-            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"):
 
-    # --- 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)
+    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()
+      
+        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(' ') + 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 ")
+    else:
+          run_prompt(prompt)
 
-if __name__ == "__main__":
-    main()