diff --git a/pyproject.toml b/pyproject.toml index 16c4247..bc7f064 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,8 @@ dependencies = [ "geopandas>=1.0.1", "geemap>=0.35.1", "langgraph>=0.2.56", + "pystac>=1.11.0", + "pystac-client>=0.8.5", ] [dependency-groups] diff --git a/tests/test_stac_tool.py b/tests/test_stac_tool.py new file mode 100644 index 0000000..926bddf --- /dev/null +++ b/tests/test_stac_tool.py @@ -0,0 +1,15 @@ +import datetime + +from zeno.tools.stac.stac_tool import stac_tool + + +def test_stac_tool(): + result = stac_tool.invoke( + input={ + "bbox": (73.88168, 15.45949, 73.88268, 15.46049), + "min_date": datetime.datetime(2024, 8, 1), + "max_date": datetime.datetime(2024, 9, 1), + } + ) + assert len(result) == 7 + assert result[0] == "S2A_43PCT_20240831_0_L2A" diff --git a/uv.lock b/uv.lock index 92be943..c9255e3 100644 --- a/uv.lock +++ b/uv.lock @@ -3201,6 +3201,8 @@ dependencies = [ { name = "langchain-openai" }, { name = "langfuse" }, { name = "langgraph" }, + { name = "pystac" }, + { name = "pystac-client" }, { name = "streamlit" }, { name = "streamlit-folium" }, { name = "thefuzz" }, @@ -3236,6 +3238,8 @@ requires-dist = [ { name = "langchain-openai", specifier = ">=0.2.6" }, { name = "langfuse", specifier = ">=2.53.9" }, { name = "langgraph", specifier = ">=0.2.56" }, + { name = "pystac", specifier = ">=1.11.0" }, + { name = "pystac-client", specifier = ">=0.8.5" }, { name = "streamlit", specifier = ">=1.40.0" }, { name = "streamlit-folium", specifier = ">=0.23.2" }, { name = "thefuzz", specifier = ">=0.22.1" }, @@ -3683,6 +3687,37 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/2f/68116db5b36b895c0450e3072b8cb6c2fac0359279b182ea97014d3c8ac0/pyshp-2.3.1-py2.py3-none-any.whl", hash = "sha256:67024c0ccdc352ba5db777c4e968483782dfa78f8e200672a90d2d30fd8b7b49", size = 46537 }, ] +[[package]] +name = "pystac" +version = "1.11.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/4f/f6f89956aabffd1211fa9c9130293ac67f774c66fab7944bfe33dc317f18/pystac-1.11.0.tar.gz", hash = "sha256:acb1e04be398a0cda2d8870ab5e90457783a8014a206590233171d8b2ae0d9e7", size = 141392 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/5b/60dc94cbf6af2fd3a3d3fae52de7e294819af2dfe7a1bea4d246beb7e0b6/pystac-1.11.0-py3-none-any.whl", hash = "sha256:10ac7c7b4ea6c5ec8333829a09ec1a33b596f02d1a97ffbbd72cd1b6c10598c1", size = 183925 }, +] + +[package.optional-dependencies] +validation = [ + { name = "jsonschema" }, +] + +[[package]] +name = "pystac-client" +version = "0.8.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pystac", extra = ["validation"] }, + { name = "python-dateutil" }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/0a/7f25bbad05b72602789f08b8ea444f362440e70e330b36c642ac1440e242/pystac_client-0.8.5.tar.gz", hash = "sha256:7fba8d4f3c641ff7e840084fc3a53c96443a227f8a5889ae500fc38183ccd994", size = 52367 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d8/25/1130b9916f7ff42fae3719c479ba4d7038659e3b6b681b736a3c936d5cc1/pystac_client-0.8.5-py3-none-any.whl", hash = "sha256:0da63244e46dae45f66b6bd6e40eb94c03cb6e4850536ef3ebd58a7efeb48f69", size = 41863 }, +] + [[package]] name = "pytest" version = "8.3.3" diff --git a/zeno/tools/stac/stac_tool.py b/zeno/tools/stac/stac_tool.py new file mode 100644 index 0000000..d6a4494 --- /dev/null +++ b/zeno/tools/stac/stac_tool.py @@ -0,0 +1,63 @@ +import datetime +from typing import Tuple + +from langchain_core.tools import tool +from pydantic import BaseModel, Field +from pystac_client import Client + + +class StacInput(BaseModel): + """Input schema for STAC search tool""" + + catalog: str = Field( + description="STAC catalog to use for search", + default="https://earth-search.aws.element84.com/v1", + ) + collection: str = Field( + description="STAC Clollection to use", default="sentinel-2-l2a" + ) + bbox: Tuple[float, float, float, float] = Field( + description="Bounding box for STAC search." + ) + min_date: datetime.datetime = Field( + description="Earliest date for retrieving STAC items.", + ) + max_date: datetime.datetime = Field( + description="Latest date for retrieving STAC items", + ) + + +@tool( + "stac-tool", + args_schema=StacInput, + response_format="content_and_artifact", +) +def stac_tool( + bbox: Tuple[float, float, float, float], + min_date: datetime.datetime, + max_date: datetime.datetime, + catalog: str = "https://earth-search.aws.element84.com/v1", + collection: str = "sentinel-2-l2a", +) -> dict: + """Find locations and their administrative hierarchies given a place name. + Returns a list of IDs with matches at different administrative levels + """ + print("---SENTINEL-TOOL---") + + catalog = Client.open(catalog) + + query = catalog.search( + collections=[collection], + datetime=[min_date, max_date], + max_items=10, + bbox=bbox, + ) + + items = list(query.items()) + print(f"Found: {len(items):d} datasets") + + # Convert STAC items into a GeoJSON FeatureCollection + stac_json = query.item_collection_as_dict() + stac_ids = [item.id for item in items] + + return stac_ids, stac_json