Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,10 +8,9 @@ import base64
 from PIL import Image
 import io
 
-# Import your agent function (assuming it's in the same directory or properly imported)
-# from your_agent_file import call_agent
+from main import call_agent as agent_invoker
+from main import DemoState
 
-# Mock function for demonstration - replace with your actual call_agent function
 def call_agent(query):
     """
     Mock function - replace this with your actual call_agent function
@@ -23,16 +22,20 @@ def call_agent(query):
         {"id": 3, "task": "Analyze hydrological flow patterns", "tool": "HydrologyAnalyzer"},
         {"id": 4, "task": "Generate flood risk assessment maps", "tool": "LLM Reasoning"}
     ]
-    
+    state =  agent_invoker(query)
     # Mock state object
     class MockState:
         def __init__(self):
-            self.query = query
-            self.response = mock_response
-            self.output_files_path = [f"outputs/geospatial_plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}.json"]
-    
+            self.query = state.get('query', query)
+            self.response = state.get('response', mock_response)
+            self.output_files_path = state.get('output_files_path', [f"outputs/geospatial_plan_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}.json"])
     return MockState()
 
+
+
+
+
+
 def get_image_base64(image_path):
     """Convert local image to base64 string for embedding in HTML"""
     try:
@@ -688,8 +691,8 @@ with col2:
     </div>
     """, unsafe_allow_html=True)
     
-    # Chat messages container
-    st.markdown('<div class="chat-messages">', unsafe_allow_html=True)
+
+
     
     # Display chat messages
     for message in st.session_state.messages:
@@ -707,11 +710,7 @@ with col2:
                     <strong>🤖 AI Task Planner:</strong>
                 </div>
                 """, unsafe_allow_html=True)
-                
-                # Use Streamlit's native components for better rendering
-                st.markdown("""
-                <div class="task-plan-container">
-                """, unsafe_allow_html=True)
+
                 
                 for task in message["tasks"]:
                     st.markdown(f"""
@@ -752,10 +751,9 @@ with col2:
         </div>
         """, unsafe_allow_html=True)
     
-    st.markdown('</div>', unsafe_allow_html=True)
+
     
-    # Input section
-    st.markdown('<div class="input-section">', unsafe_allow_html=True)
+
     
     # Text input for query
     user_input = st.text_area(
@@ -836,7 +834,7 @@ with col2:
             use_container_width=True
         )
     
-    st.markdown('</div>', unsafe_allow_html=True)
+
 
 # Clear chat button (at the bottom)
 if st.session_state.messages: