feat: add buffer tool (#8)

* wip buffer tool

* add search area to attribute

* add geopandas dependency

* add buffer tool and fix graph to add in the overture tool

* make async

* simpler buffering

* cleaner comments

Co-authored-by: Wei Ji <23487320+weiji14@users.noreply.github.com>

* Change from FeatureCollection to Feature

And ensure that output search area is a Polygon

* update test comments

* remove unhelpful test checks

---------

Co-authored-by: Wei Ji <23487320+weiji14@users.noreply.github.com>
Co-authored-by: Daniel Wiesmann <yellowcap@users.noreply.github.com>

src/geo_assistant/agent/graph.py

@@ -4,7 +4,8 @@ from langgraph.checkpoint.memory import InMemorySaver
 from langchain.agents import create_agent
 from geo_assistant.agent.state import GeoAssistantState
 from geo_assistant.agent.llms import llm
-from geo_assistant.tools.overture import get_overture_locations
+from geo_assistant.tools.overture import get_place
+from src.geo_assistant.tools.buffer import get_search_area
 
 SYSTEM_PROMPT = """
 You are a helpful assistant that can answer questions and help with tasks.
@@ -19,7 +20,7 @@ async def create_graph():
     checkpointer = InMemorySaver()
     graph = create_agent(
         model=llm,
-        tools=[get_overture_locations],  # [get_overture_locations, geocode_division],
+        tools=[get_place, get_search_area],
         system_prompt=SYSTEM_PROMPT.format(
             now=datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         ),

src/geo_assistant/agent/state.py

@@ -1,7 +1,8 @@
-from langchain.agents import AgentState as BaseAgentState
-from geojson_pydantic import FeatureCollection
+from langchain.agents import AgentState
+from geojson_pydantic import Feature
 from typing import Optional
 
 
-class GeoAssistantState(BaseAgentState):
-    place: Optional[FeatureCollection] = None
+class GeoAssistantState(AgentState):
+    place: Optional[Feature]
+    search_area: Optional[Feature]

src/geo_assistant/tools/buffer.py

@@ -0,0 +1,62 @@
+import geopandas as gpd
+from langgraph.types import Command
+from langgraph.prebuilt import InjectedState
+from langchain_core.tools.base import InjectedToolCallId
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import tool
+from typing import Annotated
+from geo_assistant.agent.state import GeoAssistantState
+
+
+@tool
+async def get_search_area(
+    buffer_size_km: float,
+    state: Annotated[GeoAssistantState, InjectedState],
+    tool_call_id: Annotated[str, InjectedToolCallId] = "",
+) -> Command:
+    """Get a search area buffer in km around the place defined in the agent state."""
+
+    place_feature = state["place"]
+
+    if not place_feature:
+        return Command(
+            update={
+                "messages": [
+                    ToolMessage(
+                        content="No place defined in the agent state to create a search area around.",
+                        tool_call_id=tool_call_id,
+                    )
+                ],
+            }
+        )
+
+    # Convert GeoJSON feature to GeoDataFrame
+    gdf = gpd.GeoDataFrame.from_features(features=[place_feature])
+    gdf.crs = "EPSG:4326"
+
+    gdf_m = gdf.to_crs(epsg=3857)  # latlon to Web Mercator for meter-based buffering
+
+    gdf_m["geometry"] = gdf_m["geometry"].buffer(
+        buffer_size_km * 1000
+    )  # Buffer in meters
+    gdf = gdf_m.to_crs(epsg=4326)  # Back to WGS84
+
+    # Convert back to GeoJSON feature
+    if len(gdf) != 1:
+        raise ValueError(
+            f"{len(gdf)} features found after buffer operation, should be just 1. "
+            "Was a Multi-Point/LineString/Polygon geometry passed in?"
+        )
+    buffer_feature = gdf.iloc[0].geometry.__geo_interface__
+
+    return Command(
+        update={
+            "search_area": buffer_feature,
+            "messages": [
+                ToolMessage(
+                    content=f"Created search area geometry buffer of {buffer_size_km} km around the place.",
+                    tool_call_id=tool_call_id,
+                )
+            ],
+        }
+    )