Update src/streamlit_app.py

src/streamlit_app.py

@@ -6,35 +6,42 @@ 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 tools import PokemonAdvisorTools
-
-# --- 1. Setup Tools & LLM ---
-# Instantiate the class to load data
-advisor = PokemonAdvisorTools()
-
-# Create the list of tool methods to wrap
-tool_methods = [
-    advisor.get_card_info,
-    advisor.find_grading_opportunities,
-    advisor.assess_risk_volatility,
-    advisor.get_roi_metrics,
-    advisor.get_recent_price_spikes,
-    advisor.analyze_set_performance,
-    advisor.find_cards_by_artist,
-    advisor.get_market_movers
-]
 
-# Wrap tools into LlamaIndex FunctionTools
-tools = [FunctionTool.from_defaults(fn=func) for func in tool_methods]
+# --- Import tools class ---
+from tools import PokemonAdvisorTools
 
-# Initialize the LLM (Ensure OPENAI_API_KEY is set in your env)
-llm = OpenAI(model="gpt-4o-mini", temperature=0.6)
+# --- 1. Page Configuration (MUST be first Streamlit command) ---
+st.set_page_config(page_title="cAsh Robo-Advisor", page_icon="🤖")
 
-# --- 2. System Prompt ---
-system_prompt = """
-### ROLE
-You are the **cAsh, Pokemon Investment Advisor**, an expert algorithmic trading assistant for the Pokemon TCG market.
-You rely **strictly** on data. You do not guess. You do not hallucinate prices. Always answer in English.
+# --- 2. Initialize Agent & Tools (with caching to avoid recreating on every rerun) ---
+@st.cache_resource
+def initialize_agent():
+    """Initialize the agent and tools once and cache them."""
+    # Instantiate the class to load data
+    advisor = PokemonAdvisorTools()
+    
+    # Create the list of tool methods to wrap
+    tool_methods = [
+        advisor.get_card_info,
+        advisor.find_grading_opportunities,
+        advisor.assess_risk_volatility,
+        advisor.get_roi_metrics,
+        advisor.get_recent_price_spikes,
+        advisor.analyze_set_performance,
+        advisor.find_cards_by_artist,
+        advisor.get_market_movers
+    ]
+    
+    # Wrap tools into LlamaIndex FunctionTools
+    tools = [FunctionTool.from_defaults(fn=func) for func in tool_methods]
+    
+    # Initialize the LLM (Ensure OPENAI_API_KEY is set in your env)
+    llm = OpenAI(model="gpt-4o-mini")
+    
+    # System Prompt
+    system_prompt = """### ROLE
+You are the **Poke-Alpha 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.
@@ -43,74 +50,60 @@ You rely **strictly** on data. You do not guess. You do not hallucinate prices.
 4. **Profit Hunting:** If a user asks "What should I buy?", use `find_grading_opportunities` or `get_market_movers`.
 
 ### TONE
-Professional, objective, and user-friendly.
-"""
-# --- 3. Initialize Agent & Global Memory ---
-# We use timeout=120 because deep reasoning sequences can take time
-agent = ReActAgent(
-    tools=tools, 
-    llm=llm, 
-    verbose=True, 
-    system_prompt=system_prompt,
-    timeout=120
-)
-
-# Global Context acts as the "Server Memory" for this session
-global_ctx = Context(agent)
-
-# --- 4. Define the Chat Function (Async) ---
-async def ask_advisor(user_message, history):
-    st.session_state.messages.append({"role": "user", "content": user_message})
-    """
-    Async function to handle the chat. 
-    It uses 'global_ctx' to maintain memory of previous turns/corrections.
-    """
+Professional, objective, concise, and user-friendly."""
+    
+    # Initialize Agent
+    agent = ReActAgent(
+        tools=tools,
+        llm=llm,
+        verbose=True,
+        system_prompt=system_prompt,
+        timeout=120
+    )
+    
+    return agent
+
+# Get cached agent
+agent = initialize_agent()
+
+# --- 3. Initialize Session State ---
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+if "context" not in st.session_state:
+    # Create a new Context for this session
+    st.session_state.context = Context(agent)
+
+# --- 4. Define the Chat Function (Sync wrapper for async) ---
+async def ask_advisor_async(user_message, ctx):
+    """Async function to handle the chat."""
     if not user_message:
         return "Please enter a message."
-
+    
     try:
         # Execute the agent workflow
-        # The 'ctx' argument passes the memory of previous interactions
-        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})        
-        # Return the final text response
+        response = await agent.run(user_msg=user_message, ctx=ctx)
         return str(response)
-        
     except Exception as e:
         return f"⚠️ **Agent Error:** {str(e)}\n\n*Check the console logs for detailed tool output.*"
 
-# --- 2. Page Configuration ---
-st.set_page_config(page_title="cAsh Robo-Advisor", page_icon="🤖", layout="wide")
+def ask_advisor(user_message):
+    """Synchronous wrapper for the async function."""
+    return asyncio.run(ask_advisor_async(user_message, st.session_state.context))
 
-# --- 3. Chat Logic Function ---
-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()
-        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 ---
+# --- 5. Styling & Header ---
 st.title("🤖 cAsh Robo-Advisor")
-st.markdown("---")
+st.markdown("""
+**Your AI Quantitative Analyst for Pokemon Cards.**
 
-# Initialize Chat History
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+Ask about:
+* **Arbitrage:** "What are the best grading opportunities?"
+* **Risk:** "Is Pikachu EX a safe investment?"
+* **Trends:** "What is crashing right now?"
+* **Sets:** "How is Evolving Skies performing?"
+""")
 
-# Sidebar Examples
+# --- 6. Sidebar Examples ---
 st.sidebar.header("Example Queries")
 examples = [
     "What are the top 3 grading opportunities right now?",
@@ -119,22 +112,34 @@ examples = [
     "Show me profitable cards by Tomokazu Komiya."
 ]
 
-# Sidebar helper: clicking an example sends it immediately
 for ex in examples:
-    if st.sidebar.button(ex):
-        asyncio.run(process_chat(ex))
-
-# --- 5. Main Chat Interface ---
-# Display existing messages
+    if st.sidebar.button(ex, key=f"example_{ex}"):
+        # Add to messages and set a flag to process
+        st.session_state.messages.append({"role": "user", "content": ex})
+        # Process the example query
+        with st.spinner("Thinking..."):
+            response = ask_advisor(ex)
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        st.rerun()
+
+# --- 7. Display Chat History ---
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# User Input
+# --- 8. Chat Input Logic ---
 if prompt := st.chat_input("Ask your Pokemon Quants advisor..."):
-    # Clear the "user" UI immediately then run logic
+    # Display user message
     with st.chat_message("user"):
         st.markdown(prompt)
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    # Generate Response
+    with st.chat_message("assistant"):
+        with st.spinner("Analyzing..."):
+            response = ask_advisor(prompt)
+        st.markdown(response)
+    st.session_state.messages.append({"role": "assistant", "content": response})
     
-    asyncio.run(ask_advisor(prompt, global_ctx))
-    #ask_advisor(prompt)
+    # Force a rerun to update the chat display
+    st.rerun()