mirror of
https://github.com/dataforcanada/d4c-service-geo-assistant.git
synced 2026-06-13 22:41:01 +02:00
Martha
89 lines
2.4 KiB
Python
89 lines
2.4 KiB
Python
from typing import Annotated
|
|
import duckdb
|
|
import json
|
|
from langchain_core.tools import tool
|
|
from langgraph.types import Command
|
|
from langchain_core.messages import ToolMessage
|
|
from langchain_core.tools.base import InjectedToolCallId
|
|
|
|
|
|
def create_database_connection():
|
|
"""Create and configure a DuckDB connection with necessary extensions.
|
|
|
|
Args:
|
|
database_path: Path to the DuckDB database file
|
|
|
|
Returns:
|
|
Configured DuckDB connection
|
|
"""
|
|
connection = duckdb.connect()
|
|
connection.execute("INSTALL spatial;")
|
|
connection.execute("INSTALL httpfs;")
|
|
connection.load_extension("spatial")
|
|
connection.load_extension("httpfs")
|
|
return connection
|
|
|
|
|
|
@tool
|
|
def get_place(
|
|
place_name: str, tool_call_id: Annotated[str, InjectedToolCallId] = ""
|
|
) -> Command:
|
|
"""Get place location from Overture Maps based on user input place name."""
|
|
|
|
db_connection = create_database_connection()
|
|
|
|
location_results = db_connection.execute(
|
|
f"""
|
|
LOAD spatial;
|
|
|
|
SET s3_region='us-west-2';
|
|
|
|
SELECT
|
|
id,
|
|
jaro_winkler_similarity(LOWER(names.primary), LOWER('{place_name}')) AS similarity_score,
|
|
names.primary AS name,
|
|
confidence,
|
|
CAST(socials AS JSON) AS socials,
|
|
ST_AsGeoJSON(geometry) AS geometry,
|
|
FROM read_parquet(
|
|
's3://overturemaps-us-west-2/release/2025-11-19.0/theme=places/type=place/*',
|
|
filename=true,
|
|
hive_partitioning=1
|
|
)
|
|
WHERE jaro_winkler_similarity(LOWER(names.primary), LOWER('{place_name}')) > 0.5
|
|
ORDER BY similarity_score DESC
|
|
LIMIT 1;
|
|
"""
|
|
).fetchall()
|
|
|
|
db_connection.close()
|
|
|
|
geometry = json.loads(location_results[0][-1])
|
|
|
|
# Create FeatureCollection
|
|
feature_collection = {
|
|
"type": "FeatureCollection",
|
|
"features": [
|
|
{
|
|
"type": "Feature",
|
|
"geometry": geometry,
|
|
"properties": {
|
|
"name": location_results[0][2],
|
|
"overture_id": location_results[0][0],
|
|
},
|
|
}
|
|
],
|
|
}
|
|
|
|
return Command(
|
|
update={
|
|
"place": feature_collection,
|
|
"messages": [
|
|
ToolMessage(
|
|
content=f"Found place with Overture name: {location_results[0][2]} based on user query",
|
|
tool_call_id=tool_call_id,
|
|
)
|
|
],
|
|
},
|
|
)
|