fiboa/query.py

import re

import leafmap.maplibregl as leafmap
import pandas as pd
import streamlit as st
from ibis import _


def execute_prompt(con, chain, prompt):
    response = chain.invoke({"question": prompt})
    st.write(response)
    gdf = as_geopandas(con, response)

    if 'geometry' in gdf.columns:
        m = leafmap.Map()
        m.add_gdf(gdf)
        m.to_streamlit()
    else:
        gdf.drop(columns=['geometry', 'bbox', 'bbox.minx', 'bbox.maxx', 'bbox.miny', 'bbox.maxy'], errors='ignore', inplace=True)
        st.dataframe(pd.DataFrame(gdf), column_config={
            "area": st.column_config.NumberColumn("Area (ha)", format="%.5f"),
            "perimeter": st.column_config.NumberColumn("Perimeter (m)", format="%.3f"),
            "determination_datetime": st.column_config.DatetimeColumn("Determination Date"),
        })

def as_geopandas(con, response):
    response = re.sub(";$", "", response)
    sql_query = f"CREATE OR REPLACE VIEW testing AS ({response})"
    con.raw_sql(sql_query)
    gdf = con.table("testing")

    if 'geometry' in gdf.columns:
        gdf = (gdf
               .cast({"geometry": "geometry"})
               .mutate(geometry = _.geometry.convert("EPSG:4326", "EPSG:4326"))
               .to_pandas()
              ).set_crs(epsg=4326, inplace=True) 
    else:
        gdf = gdf.to_pandas()

    for col in gdf.columns:
        dtype = str(gdf[col].dtype)
        if dtype.startswith("datetime64"):
            gdf[col] = gdf[col].astype(str)
    
    return gdf