Fix frontend (#15)

* Fix frontend

* Fix feature creation from overture tool

src/geo_assistant/frontend/app.py

@@ -4,6 +4,8 @@ import uuid
 
 import httpx
 import streamlit as st
+import streamlit.components.v1 as components
+import folium
 from dotenv import load_dotenv
 
 # Load environment variables from .env file
@@ -14,6 +16,8 @@ API_BASE_URL = os.environ.get("API_BASE_URL", "http://localhost:8000")
 
 st.set_page_config(page_title="Geo Assistant", page_icon="💬")
 
+st.title("Geo Assistant")
+
 # Initialize session state
 if "thread_id" not in st.session_state:
     st.session_state.thread_id = str(uuid.uuid4())
@@ -44,58 +48,106 @@ def stream_chat(user_message: str):
         response.raise_for_status()
 
         for line in response.iter_lines():
-            print("=" * 100)
-            print(line)
-            print("=" * 100)
-
             if not line:
                 continue
 
             data = json.loads(line)
-            print("=" * 100)
-            print(data)
-            print("=" * 100)
             state = data.get("state", {})
-            messages = state.get("messages", [])
-
-            if not messages:
-                continue
+            messages = state.pop("messages", [])
 
             for msg in messages:
                 msg_type = msg.get("type", "")
                 content = msg.get("content", "")
+                if not content:
+                    continue
+                with st.chat_message(msg_type):
+                    st.markdown(content)
 
-                yield msg_type, content
+            # Check for GeoJSON features and render map if present
+            geojson_features = {}
+            for key, value in state.items():
+                if value and isinstance(value, dict) and value.get("type") == "Feature":
+                    geojson_features[key] = value
+                    # with st.chat_message("tool"):
+                    #     st.code(json.dumps(value, indent=2), language="json")
+                elif value:
+                    with st.chat_message("tool"):
+                        st.code(json.dumps(value, indent=2), language="json")
+
+            # Render map if GeoJSON features are present
+            if geojson_features:
+                # Helper function to extract coordinates from geometry
+                def get_coords_from_geometry(geom):
+                    """Extract all coordinates from a GeoJSON geometry."""
+                    geom_type = geom.get("type", "")
+                    coords = geom.get("coordinates", [])
+
+                    if geom_type == "Point":
+                        return [coords]
+                    elif geom_type == "LineString":
+                        return coords
+                    elif geom_type == "Polygon":
+                        return coords[0] if coords else []
+                    elif geom_type == "MultiPoint":
+                        return coords
+                    elif geom_type == "MultiLineString":
+                        return [c for line in coords for c in line]
+                    elif geom_type == "MultiPolygon":
+                        return [c for poly in coords for c in poly[0]] if coords else []
+                    return []
+
+                # Calculate center from all features
+                all_lons, all_lats = [], []
+                for feature in geojson_features.values():
+                    geom = feature.get("geometry", {})
+                    coords = get_coords_from_geometry(geom)
+                    for coord in coords:
+                        if len(coord) >= 2:
+                            all_lons.append(coord[0])
+                            all_lats.append(coord[1])
+
+                if all_lons and all_lats:
+                    center_lat = sum(all_lats) / len(all_lats)
+                    center_lon = sum(all_lons) / len(all_lons)
+                else:
+                    center_lat, center_lon = 0.0, 0.0
+
+                m = folium.Map(location=[center_lat, center_lon], zoom_start=10)
+
+                # Add features to map with different colors
+                colors = {"place": "blue", "search_area": "red"}
+
+                def make_style_function(color):
+                    """Create a style function with the given color."""
+                    return lambda x: {
+                        "fillColor": color,
+                        "color": color,
+                        "weight": 2,
+                        "fillOpacity": 0.3,
+                    }
+
+                for key, feature in geojson_features.items():
+                    color = colors.get(key, "green")
+                    folium.GeoJson(
+                        feature,
+                        style_function=make_style_function(color),
+                        tooltip=key,
+                    ).add_to(m)
+
+                # Fit map to bounds if we have coordinates
+                if all_lons and all_lats:
+                    m.fit_bounds(
+                        [[min(all_lats), min(all_lons)], [max(all_lats), max(all_lons)]]
+                    )
+
+                # Display the map
+                with st.chat_message("tool"):
+                    st.markdown("**Map View**")
+                    map_html = m._repr_html_()
+                    components.html(map_html, height=400)
 
 
-# Main UI
-st.title("Geo Assistant")
-
-# Display chat history
-for item in st.session_state.chat_history:
-    role = item["role"]
-    content = item["content"]
-
-    with st.chat_message(role):
-        if role == "assistant":
-            # For assistant messages, check if it's a tool message
-            if item.get("is_tool"):
-                st.code(content, language="json")
-            else:
-                st.markdown(content)
-        else:
-            st.markdown(content)
-
-# Chat input
 if prompt := st.chat_input("Type your message..."):
-    st.session_state.chat_history.append({"role": "user", "content": prompt})
-
-    for msg_type, content in stream_chat(prompt):
-        if msg_type == "tool":
-            st.session_state.chat_history.append({"role": "tool", "content": content})
-        elif msg_type in ["ai", "assistant"]:
-            st.session_state.chat_history.append(
-                {"role": "assistant", "content": content}
-            )
-
-    st.rerun()
+    with st.chat_message("user"):
+        st.markdown(prompt)
+    stream_chat(prompt)