Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,10 +1,11 @@
 import streamlit as st
 import pandas as pd
 import os
+import asyncio
 from llama_index.llms.openai import OpenAI
 from llama_index.core.agent import ReActAgent
+from llama_index.core.workflow import Context 
 from llama_index.core.tools import FunctionTool
-from llama_index.core.llms import ChatMessage
 
 # --- Import tools class ---
 from tools import PokemonAdvisorTools
@@ -25,58 +26,57 @@ def initialize_agent():
         advisor.get_data_shape
     ]
     tools = [FunctionTool.from_defaults(fn=func) for func in tool_methods]
-    llm = OpenAI(model="gpt-4o-mini", temperature=0.1)
+    llm = OpenAI(model="gpt-4o-mini")
     
-    system_prompt = """### ROLE
-You are the **cAsh, Pokemon Investment Advisor**, an expert algorithmic trading assistant for the Pokémon TCG market.
-You rely **strictly** on data. You do not guess. You do not hallucinate prices.
-
-### TOOL USAGE PROTOCOL
-1. **Verify First:** If a user asks about a card, ALWAYS use `get_card_info` first to ensure it exists.
-2. **Safety Check:** BEFORE recommending ANY card, you MUST call `assess_risk_volatility` (default to 6 months).
-3. **Refining Search:** If the user corrects you (e.g., "Not that card"), use `get_card_info` again with the corrected name or use `analyze_set_performance` to broaden the scope.
-4. **Profit Hunting:** If a user asks "What should I buy?", use `find_grading_opportunities` or `get_market_movers`.
-
-### TONE
-Professional, objective, concise, and user-friendly.
-"""
+    system_prompt = """
+        ### ROLE
+        You are the **cAsh, Pokemon Investment Advisor**, an expert algorithmic trading assistant for the Pokémon TCG market.
+        You rely **strictly** on data. You do not guess. You do not hallucinate prices.
+        
+        ### TOOL USAGE PROTOCOL
+        1. **Verify First:** If a user asks about a card, ALWAYS use `get_card_info` first to ensure it exists.
+        2. **Safety Check:** BEFORE recommending ANY card, you MUST call `assess_risk_volatility` (default to 6 months).
+        3. **Refining Search:** If the user corrects you (e.g., "Not that card"), use `get_card_info` again with the corrected name or use `analyze_set_performance` to broaden the scope.
+        4. **Profit Hunting:** If a user asks "What should I buy?", use `find_grading_opportunities` or `get_market_movers`.
+        
+        ### TONE
+        Professional, objective, concise, and user-friendly.
+        """
 
     agent = ReActAgent(
         tools=tools, 
         llm=llm, 
-        verbose=True,
-        max_iterations=10
+        verbose=True, 
+        system_prompt=system_prompt,
+        timeout=120
     )
-    
-    # Set system prompt via chat history
-    agent.chat_history = [ChatMessage(role="system", content=system_prompt)]
-    
-    return agent
+    # Context should be initialized once
+    ctx = Context()
+    return agent, ctx
 
-agent = initialize_agent()
+agent, global_ctx = initialize_agent()
+print("agent has initialized")
 
 # --- 2. Page Configuration ---
 st.set_page_config(page_title="cAsh Robo-Advisor", page_icon="🤖", layout="wide")
 
 # --- 3. Chat Logic Function ---
-def process_chat(user_message):
-    """Handles the synchronous call to the agent and updates session state."""
+async def process_chat(user_message):
+    """Handles the async call to the agent and updates session state."""
     st.session_state.messages.append({"role": "user", "content": user_message})
     
     # Placeholder for the assistant's response while it thinks
     with st.chat_message("assistant"):
         response_placeholder = st.empty()
-        with st.spinner("Thinking..."):
-            try:
-                # Execute the agent chat
-                response = agent.chat(user_message)
-                final_text = str(response)
-                response_placeholder.markdown(final_text)
-                st.session_state.messages.append({"role": "assistant", "content": final_text})
-            except Exception as e:
-                error_msg = f"⚠️ **Agent Error:** {str(e)}"
-                response_placeholder.markdown(error_msg)
-                st.error(f"Full error: {repr(e)}")
+        try:
+            # Execute the agent workflow
+            response = await agent.run(user_msg=user_message, ctx=global_ctx)
+            final_text = str(response)
+            response_placeholder.markdown(final_text)
+            st.session_state.messages.append({"role": "assistant", "content": final_text})
+        except Exception as e:
+            error_msg = f"⚠️ **Agent Error:** {str(e)}"
+            response_placeholder.markdown(error_msg)
 
 # --- 4. Sidebar & UI Styling ---
 st.title("🤖 cAsh Robo-Advisor")
@@ -90,7 +90,6 @@ if "messages" not in st.session_state:
 st.sidebar.header("Example Queries")
 examples = [
     "What are the top 3 grading opportunities right now?",
-    "Is investing in Charizard VMAX risky?",
     "What cards are trending down?",
     "Show me profitable cards by Tomokazu Komiya."
 ]
@@ -98,8 +97,7 @@ examples = [
 # Sidebar helper: clicking an example sends it immediately
 for ex in examples:
     if st.sidebar.button(ex):
-        process_chat(ex)
-        st.rerun()
+        asyncio.run(process_chat(ex))
 
 # --- 5. Main Chat Interface ---
 # Display existing messages
@@ -109,6 +107,8 @@ for message in st.session_state.messages:
 
 # User Input
 if prompt := st.chat_input("Ask your Pokemon Quants advisor..."):
+    # Clear the "user" UI immediately then run logic
     with st.chat_message("user"):
         st.markdown(prompt)
-    process_chat(prompt)
+    
+    asyncio.run(process_chat(prompt))