feat: add config options, simplification settings, geometry retrieval by id (#11)

* feat: add config options, simplification settings, geometry retrieval by id

* fix: update sql param and model key selection

Author: ividito

geodini/agents/geocoder_agent.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+import os
 import time
 import traceback
 from concurrent.futures import ThreadPoolExecutor
@@ -15,6 +16,7 @@
 from shapely.ops import transform
 
 from geodini import hookspecs, lib
+from geodini.models import MODEL_LIGHT
 from geodini.agents.utils.postgis_exec import (
     clear_geometries_table,
     create_geometries_table,
@@ -72,9 +74,12 @@ class RerankingResult:
     most_probable: str
 
 
+_GEOMETRY_PRECISION = int(os.getenv("GEOMETRY_DECIMAL_PRECISION", "6"))
+
+
 class RoundedFloat(float):
     def __repr__(self):
-        return f"{self:.2f}"
+        return f"{self:.{_GEOMETRY_PRECISION}f}"
 
 
 def recursively_convert(obj):
@@ -93,8 +98,11 @@ def clip_coordinates_with_rounding(geojson: dict[str, Any]) -> dict[str, Any]:
 
 
 def simplify_geometry(
-    geometry: dict[str, Any], tolerance_m: float = 10000
+    geometry: dict[str, Any],
+    tolerance_m: float | None = None,
 ) -> dict[str, Any]:
+    if tolerance_m is None:
+        tolerance_m = float(os.getenv("GEOMETRY_AGENT_SIMPLIFY_TOLERANCE", "10000"))
     to_meters = Transformer.from_crs("EPSG:4326", "EPSG:3857", always_xy=True).transform
     to_degrees = Transformer.from_crs(
         "EPSG:3857", "EPSG:4326", always_xy=True
@@ -111,8 +119,28 @@ def simplify_geometry(
     return geojson
 
 
+def simplify_geometry_to_size(
+    geometry: dict[str, Any],
+    max_bytes: int,
+    initial_tolerance_m: float = 1000,
+    max_iterations: int = 10,
+) -> dict[str, Any]:
+    """Iteratively simplify geometry until JSON representation fits within max_bytes."""
+    tolerance = initial_tolerance_m
+    result = simplify_geometry(geometry, tolerance_m=tolerance)
+
+    for _ in range(max_iterations):
+        serialized = json.dumps(result)
+        if len(serialized.encode("utf-8")) <= max_bytes:
+            return result
+        tolerance *= 2
+        result = simplify_geometry(geometry, tolerance_m=tolerance)
+
+    return result
+
+
 rephrase_agent = Agent(
-    "openai:gpt-4.1-mini",
+    MODEL_LIGHT,
     output_type=RephrasedQuery,
     system_prompt="""
         Given the search query, rephrase it to be more specific and accurate. We will be using this query to search for places in the overture database. So it helps to make the query be full formal name of the place.
@@ -131,7 +159,7 @@ def simplify_geometry(
 
 
 routing_agent = Agent(
-    "openai:gpt-4.1-mini",
+    MODEL_LIGHT,
     output_type=RoutingResult,
     system_prompt="""
         Given the search query, determine if it is a simple or complex query.
@@ -143,7 +171,7 @@ def simplify_geometry(
 
 
 complex_geocode_query_agent = Agent(
-    "openai:gpt-4.1-mini",
+    MODEL_LIGHT,
     output_type=ComplexGeocodeResult,
     system_prompt="""
         Given the search query, return ALL relevant places to search for in the query as queries.
@@ -170,8 +198,7 @@ def simplify_geometry(
 
 
 rerank_agent = Agent(
-    # 4o-mini is smarter than 3.5-turbo. And does better in edge cases.
-    "openai:gpt-4.1-mini",
+    MODEL_LIGHT,
     output_type=RerankingResult,
     system_prompt="""
         Given the search query and results, rank them in order of 
@@ -312,6 +339,7 @@ async def simple_geocode(query: str, simplify_geometry: bool = True) -> dict:
         "query": query,
         "results": [
             {
+                "id": most_probable["id"] if most_probable else None,
                 "geometry": most_probable["geometry"] if most_probable else None,
                 "country": most_probable["country"] if most_probable else None,
                 "name": most_probable["name"] if most_probable else query,

geodini/agents/utils/geocoder.py

@@ -107,7 +107,8 @@ def build_postgis_query(simplify_geometry: bool = True) -> str:
     """Build PostgreSQL query for searching overture unified data using trigram similarity"""
 
     # Choose geometry function based on whether to simplify
-    geometry_func = "ST_AsGeoJSON(ST_Simplify(geometry, 0.001))" if simplify_geometry else "ST_AsGeoJSON(geometry)"
+    db_tolerance = float(os.getenv("GEOMETRY_DB_SIMPLIFY_TOLERANCE", "0.001"))
+    geometry_func = f"ST_AsGeoJSON(ST_Simplify(geometry, {db_tolerance}))" if simplify_geometry else "ST_AsGeoJSON(geometry)"
 
     sql_query = f"""
         SELECT 
@@ -162,6 +163,48 @@ def build_postgis_query(simplify_geometry: bool = True) -> str:
     return sql_query
 
 
+def get_geometry_by_id(
+    division_id: str, simplify: bool = True
+) -> dict[str, Any] | None:
+    """Retrieve geometry for a specific division by ID."""
+    engine = get_postgis_engine()
+    db_tolerance = float(os.getenv("GEOMETRY_DB_SIMPLIFY_TOLERANCE", "0.001"))
+
+    geometry_func = (
+        f"ST_AsGeoJSON(ST_Simplify(geometry, {db_tolerance}))"
+        if simplify
+        else "ST_AsGeoJSON(geometry)"
+    )
+
+    sql_query = f"""
+        SELECT
+            id,
+            COALESCE(common_en_name, primary_name) as name,
+            subtype, country,
+            {geometry_func} as geometry
+        FROM all_geometries
+        WHERE id = :id
+        LIMIT 1
+    """
+
+    try:
+        with engine.begin() as conn:
+            result = conn.execute(text(sql_query), {"id": division_id})
+            row = result.fetchone()
+            if row:
+                return {
+                    "id": row.id,
+                    "name": row.name,
+                    "subtype": row.subtype,
+                    "country": row.country,
+                    "geometry": json.loads(row.geometry) if row.geometry else None,
+                }
+    except Exception as e:
+        logger.error(f"Error retrieving geometry by ID: {e}")
+
+    return None
+
+
 if __name__ == "__main__":
     import time
 

geodini/agents/utils/postgis_exec.py

@@ -5,6 +5,8 @@
 import psycopg2
 from pydantic_ai import Agent
 
+from geodini.models import MODEL_HEAVY
+
 
 @dataclass
 class PostGISResult:
@@ -96,33 +98,34 @@ def search_subtype_within_aoi(subtype: str, aoi: dict) -> list[dict]:
     # aoi is the geojson geometry as dict as returned from run_postgis_query
     # return a list of dictionaries of the form:
     # { "geometry": result_geometry_as_json_dict, "country": country_name }
+    db_tolerance = float(os.getenv("GEOMETRY_DB_SIMPLIFY_TOLERANCE", "0.001"))
     conn = get_postgis_connection()
     try:
         with conn.cursor() as cur:
             # Convert AOI dict to GeoJSON string
             aoi_geojson = json.dumps(aoi)
-            
+
             # SQL query to find places of given subtype within the AOI
             sql_query = """
-            SELECT 
-                ST_AsGeoJSON(ST_Simplify(geometry, 0.001)) as geometry,
+            SELECT
+                ST_AsGeoJSON(ST_Simplify(geometry, %s)) as geometry,
                 country,
                 COALESCE(common_en_name, primary_name) as name
             FROM all_geometries
-            WHERE 
+            WHERE
                 source_type = 'division'
                 AND subtype = %s
                 AND geometry IS NOT NULL
                 AND ST_Within(
                     geometry,
                     ST_GeomFromGeoJSON(%s)
                 )
-            ORDER BY 
+            ORDER BY
                 ST_Area(geometry) DESC
             LIMIT 100
             """
-            
-            cur.execute(sql_query, (subtype, aoi_geojson))
+
+            cur.execute(sql_query, (db_tolerance, subtype, aoi_geojson))
             results = cur.fetchall()
 
             # Convert results to expected format
@@ -141,7 +144,7 @@ def search_subtype_within_aoi(subtype: str, aoi: dict) -> list[dict]:
 
 
 postgis_agent = Agent(
-    "openai:gpt-4.1",
+    MODEL_HEAVY,
     output_type=PostGISResult,
     system_prompt="""
     You are a helpful assistant that can help with PostGIS queries.
@@ -185,7 +188,7 @@ def search_subtype_within_aoi(subtype: str, aoi: dict) -> list[dict]:
 
 
 postgis_query_judgement_agent = Agent(
-    "openai:gpt-4o",
+    MODEL_HEAVY,
     output_type=PostGISResult,
     system_prompt="""
     You are a helpful assistant that can help with PostGIS queries.

geodini/api/api.py

@@ -7,7 +7,8 @@
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.middleware.cors import CORSMiddleware
 
-from geodini.agents.geocoder_agent import search
+from geodini.agents.geocoder_agent import search, simplify_geometry_to_size
+from geodini.agents.utils.geocoder import get_geometry_by_id
 from geodini.agents.utils.postgis_exec import get_postgis_connection
 from geodini.cache import init_cache
 
@@ -50,10 +51,11 @@ async def root():
 @app.get("/search")
 async def search_endpoint(
     query: str = Query(..., description="The search query string"),
+    max_bytes: int = Query(0, description="Max geometry size in bytes (0 = use server default)"),
 ) -> dict[str, Any]:
     """
     Unified search endpoint that handles both simple and complex queries.
-    
+
     Simple queries: "New York City", "London in Canada", "India"
     Complex queries: "India and Sri Lanka", "Within 100km of Mumbai", "France north of Paris"
 
@@ -65,6 +67,15 @@ async def search_endpoint(
         # Get result from unified search
         result = await search(query)
 
+        # Apply byte-threshold simplification if requested
+        effective_max_bytes = max_bytes or int(os.getenv("GEOMETRY_MAX_BYTES", "0"))
+        if effective_max_bytes > 0 and result.get("results"):
+            for r in result["results"]:
+                if r.get("geometry"):
+                    r["geometry"] = simplify_geometry_to_size(
+                        r["geometry"], max_bytes=effective_max_bytes
+                    )
+
         return result
 
     except Exception as e:
@@ -74,6 +85,25 @@ async def search_endpoint(
         )
 
 
+@app.get("/geometry/{geometry_id}")
+async def get_geometry_endpoint(
+    geometry_id: str,
+    simplify: bool = Query(True, description="Whether to simplify the geometry"),
+    max_bytes: int = Query(0, description="Max geometry size in bytes (0 = no limit)"),
+) -> dict[str, Any]:
+    """Retrieve full geometry by division ID."""
+    result = get_geometry_by_id(geometry_id, simplify=simplify)
+    if result is None:
+        raise HTTPException(status_code=404, detail="Geometry not found")
+
+    if max_bytes > 0 and result.get("geometry"):
+        result["geometry"] = simplify_geometry_to_size(
+            result["geometry"], max_bytes=max_bytes
+        )
+
+    return result
+
+
 @app.get("/health")
 async def health_check():
     """Health check endpoint to verify the API is running."""

geodini/api/mcp_server.py

@@ -1,3 +1,5 @@
+import os
+
 from mcp.server.fastmcp import FastMCP
 
 from geodini.agents.geocoder_agent import search, simplify_geometry
@@ -10,7 +12,8 @@ async def geocode(query: str) -> str:
     """Geocode a query and download the geojson geometry"""
     result = await search(query)
     if "result" in result and "geometry" in result["result"]:
-        return simplify_geometry(result["result"]["geometry"], tolerance_m=1000)
+        mcp_tolerance = float(os.getenv("GEOMETRY_MCP_SIMPLIFY_TOLERANCE", "1000"))
+        return simplify_geometry(result["result"]["geometry"], tolerance_m=mcp_tolerance)
     return "No geometry found for query"
 
 

geodini/models.py

@@ -0,0 +1,33 @@
+"""Centralized LLM model configuration.
+
+Two tiers control all agent models:
+- MODEL_LIGHT: fast/cheap tasks (rephrase, routing, complex parse, rerank)
+- MODEL_HEAVY: SQL generation and error correction
+
+Each model's provider is extracted from its prefix (e.g. "anthropic"
+from "anthropic:claude-sonnet-4-6"). The provider-specific API key env
+var (e.g. ANTHROPIC_API_KEY) takes priority; LLM_API_KEY is used as a
+shared fallback for any provider that doesn't have its own key set.
+
+Supported provider prefixes: openai, anthropic, groq, mistral.
+Bedrock uses AWS creds (AWS_ACCESS_KEY_ID, etc.).
+Ollama uses OLLAMA_BASE_URL (no API key needed).
+"""
+
+import os
+
+MODEL_LIGHT = os.getenv("MODEL_LIGHT", "anthropic:claude-haiku-4-5")
+MODEL_HEAVY = os.getenv("MODEL_HEAVY", "anthropic:claude-sonnet-4-6")
+
+KEY_MAP = {
+    "openai": "OPENAI_API_KEY",
+    "anthropic": "ANTHROPIC_API_KEY",
+    "groq": "GROQ_API_KEY",
+    "mistral": "MISTRAL_API_KEY",
+}
+
+llm_fallback_key = os.getenv("LLM_API_KEY", "")
+for model in (MODEL_LIGHT, MODEL_HEAVY):
+    provider = model.split(":")[0]
+    if provider in KEY_MAP and llm_fallback_key:
+        os.environ.setdefault(KEY_MAP[provider], llm_fallback_key)

helm/geodini/templates/_helpers.tpl

@@ -61,6 +61,17 @@ Create the name of the service account to use
 {{- end -}}
 {{- end -}}
 
+{{/*
+Return the secret name to use for credentials.
+*/}}
+{{- define "geodini.secretName" -}}
+{{- if .Values.secrets.existingSecret -}}
+{{- .Values.secrets.existingSecret -}}
+{{- else -}}
+{{- include "geodini.fullname" . }}-geodini-secret
+{{- end -}}
+{{- end -}}
+
 {{/*
 Return the appropriate apiVersion for deployment.
 */}}