working examples with sqlalchemy2.0 using original design

graphs.py

@@ -0,0 +1,198 @@
+# This example does not use a langchain agent, 
+# The langchain sql chain has knowledge of the database, but doesn't interact with it becond intialization.
+# The output of the sql chain is parsed seperately and passed to `duckdb.sql()` by streamlit
+
+import os
+os.environ["WEBSOCKET_TIMEOUT_MS"] = "300000" # no effect
+
+import streamlit as st
+import geopandas as gpd
+import pandas as pd
+from shapely import wkb
+
+
+st.set_page_config(page_title="Protected Areas Database Chat", page_icon="🦜", layout="wide")
+st.title("Protected Areas Database Chat")
+
+## Database connection, reading directly from remote parquet file
+from sqlalchemy import create_engine
+from langchain.sql_database import SQLDatabase
+db_uri = "duckdb:///:memory:"
+parquet = "https://huggingface.co/datasets/boettiger-lab/pad-us-3/resolve/main/pad-stats.parquet"
+engine = create_engine(db_uri) #connect_args={'read_only': True})
+con = engine.connect()
+con.execute("install spatial; load spatial;")
+h = con.execute(f"create or replace view pad as select * from read_parquet('{parquet}');")
+h.fetchall()
+db = SQLDatabase(engine, view_support=True)
+
+
+@st.cache_data
+def query_database(response):
+    # con.sql(response).to_pandas().head(25) # uses ibis connection
+    # instead, use direct sqlAlchemy connection
+    z = con.execute(response).fetchall()
+    return pd.DataFrame(z).head(25)
+     
+
+query_database("select * from pad limit 1")
+
+
+@st.cache_data
+def get_geom(tbl):
+    tbl['geometry'] = tbl['geometry'].apply(wkb.loads)
+    gdf = gpd.GeoDataFrame(tbl, geometry='geometry')
+    return gdf
+
+
+
+# Helper plotting functions
+import pydeck as pdk
+def deck_map(gdf):
+    st.write(
+        pdk.Deck(
+            map_style="mapbox://styles/mapbox/light-v9",
+            initial_view_state={
+                "latitude": 35,
+                "longitude": -100,
+                "zoom": 3,
+                "pitch": 50,
+            },
+            layers=[
+                pdk.Layer(
+                    "GeoJsonLayer",
+                    gdf,
+                    pickable=True,
+                    stroked=True,
+                    filled=True,
+                    extruded=True,
+                    elevation_scale=10,
+                    get_fill_color=[2, 200, 100],
+                    get_line_color=[0,0,0],
+                    line_width_min_pixels=0,
+                ),
+            ],
+        )
+    )
+
+import leafmap.foliumap as leafmap
+def leaf_map(gdf):
+    m = leafmap.Map(center=[35, -100], zoom=4, layers_control=True)
+    m.add_gdf(gdf)
+    return m.to_streamlit()
+
+
+
+## ChatGPT Connection
+from langchain_openai import ChatOpenAI 
+# Requires ollama server running locally
+from langchain_community.llms import Ollama
+
+# # should we use ChatOllama instead?
+# from langchain_community.llms import ChatOllama
+
+models = {"chatgpt3.5": ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=st.secrets["OPENAI_API_KEY"])}
+
+other_models = {
+          "chatgpt4": ChatOpenAI(model="gpt-4", temperature=0, api_key=st.secrets["OPENAI_API_KEY"]),
+          "duckdb-nsql": Ollama(model="duckdb-nsql", temperature=0),
+          "command-r-plus": Ollama(model="command-r-plus", temperature=0),
+          "mixtral:8x22b":  Ollama(model="mixtral:8x22b", temperature=0),
+          "wizardlm2:8x22b":  Ollama(model="wizardlm2:8x22b", temperature=0),
+          "sqlcoder": Ollama(model="sqlcoder", temperature=0),
+          "zephyr": Ollama(model="zephyr", temperature=0),
+          "gemma:7b": Ollama(model="gemma:7b", temperature=0),
+          "codegemma": Ollama(model="codegemma", temperature=0),
+          "llama2": Ollama(model="llama2", temperature=0),
+         }
+
+
+
+
+map_tool = {"leafmap": leaf_map,
+            "deckgl": deck_map
+           }
+
+with st.sidebar:
+    choice = st.radio("Select an LLM:", models)
+    llm = models[choice]
+    map_choice = st.radio("Select mapping tool", map_tool)
+    mapper = map_tool[map_choice]
+## A SQL Chain
+from langchain.chains import create_sql_query_chain
+chain = create_sql_query_chain(llm, db)
+
+
+main = st.container()
+
+## Does not preserve history
+with main:
+
+    '''
+    The Protected Areas Database of the United States (PAD-US) is the official national inventory of
+    America’s parks and other protected lands, and is published by the USGS Gap Analysis Project,
+    [https://doi.org/10.5066/P9Q9LQ4B.](https://doi.org/10.5066/P9Q9LQ4B).  
+
+    This interactive tool allows users to explore the dataset, as well as a range of biodiversity
+    and climate indicators associated with each protected area. These indicators are integrated into 
+    a single table format shown below.  The chatbot assistant can turn natural language queries into
+    SQL queries based on the table schema.
+
+    See our [Protected Areas Explorer](https://huggingface.co/spaces/boettiger-lab/pad-us) for a companion non-chat-based tool.
+
+    ##### Example Queries returning summary tables
+
+    - What is the percent area in each gap code as a fraction of the total protected area?
+    - The manager_type column indicates whether a manager is federal, state, local, private, or NGO.  
+      the manager_name column indicates the responsible agency (National Park Service, Bureau of Land Management,
+      etc) in the case of federal manager types.  Which of the federal managers manage the most land in
+      gap_code 1 or 2, as a fraction of the total area?
+
+    When queries refer to specific managed areas, the chatbot can show those areas on an interactive map.
+    Do to software limitations, these maps will show no more than 25 polygons, even if more areas match the
+    requested search. The chatbot sometimes requires help identifying the right columns.  In order to create
+    a map, the SQL query must also return the geometry column.  Conisder the following examples:
+
+    ##### Example queries returning maps + tables
+
+    - Show me all the national monuments (designation_type) in Utah. Include the geometry column
+    - Show examples of Bureau of Land Management (manager_name) with the highest species richness? Include the geometry column
+    - Which site has the overall highest range-size-rarity? Include the geometry column, manager_name, and IUCN category.
+
+    '''
+
+    st.markdown("## 🦜 Chatbot:")
+    chatbox = st.container()
+    with chatbox:           
+        if prompt := st.chat_input(key="chain"):
+            st.chat_message("user").write(prompt)
+            with st.chat_message("assistant"):
+                response = chain.invoke({"question": prompt})
+                st.write(response)
+                tbl = query_database(response)
+                if 'geometry' in tbl:
+                    gdf = get_geom(tbl)
+                    mapper(gdf)
+                    n = len(gdf)
+                    st.write(f"matching features: {n}")
+                st.dataframe(tbl)
+
+
+st.divider()
+
+with st.container():
+    st.text("Database schema (top 3 rows)")
+    tbl = tbl = query_database("select * from pad limit 3")
+    st.dataframe(tbl)
+
+
+st.divider()
+
+'''
+Experimental prototype.  
+
+- Author: [Carl Boettiger](https://carlboettiger.info)
+- For data sources and processing, see: https://beta.source.coop/repositories/cboettig/pad-us-3/description/
+
+
+'''