cboettig commited on
Commit
02d986c
·
1 Parent(s): 83df5bc
Files changed (1) hide show
  1. minimal-example.py +58 -0
minimal-example.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # This example does not use a langchain agent,
2
+ # The langchain sql chain has knowledge of the database, but doesn't interact with it becond intialization.
3
+ # The output of the sql chain is parsed seperately and passed to `duckdb.sql()` by streamlit
4
+
5
+ import streamlit as st
6
+
7
+ ## Database connection
8
+ from sqlalchemy import create_engine
9
+ from langchain.sql_database import SQLDatabase
10
+ db_uri = "duckdb:///pad.duckdb"
11
+ engine = create_engine(db_uri, connect_args={'read_only': True})
12
+ db = SQLDatabase(engine, view_support=True)
13
+
14
+ import duckdb
15
+
16
+ con = duckdb.connect("pad.duckdb", read_only=True)
17
+ con.install_extension("spatial")
18
+ con.load_extension("spatial")
19
+
20
+ ## ChatGPT Connection
21
+ from langchain_openai import ChatOpenAI
22
+ chatgpt_llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=st.secrets["OPENAI_API_KEY"])
23
+ chatgpt4_llm = ChatOpenAI(model="gpt-4", temperature=0, api_key=st.secrets["OPENAI_API_KEY"])
24
+
25
+
26
+ # Requires ollama server running locally
27
+ from langchain_community.llms import Ollama
28
+ ## # from langchain_community.llms import ChatOllama
29
+ ollama_llm = Ollama(model="duckdb-nsql", temperature=0)
30
+
31
+ models = {"ollama": ollama_llm, "chatgpt3.5": chatgpt_llm, "chatgpt4": chatgpt4_llm}
32
+ with st.sidebar:
33
+ choice = st.radio("Select an LLM:", models)
34
+ llm = models[choice]
35
+
36
+ ## A SQL Chain
37
+ from langchain.chains import create_sql_query_chain
38
+ chain = create_sql_query_chain(llm, db)
39
+
40
+ # agent does not work
41
+ # agent = create_sql_agent(llm, db=db, verbose=True)
42
+
43
+ if prompt := st.chat_input():
44
+ st.chat_message("user").write(prompt)
45
+ with st.chat_message("assistant"):
46
+ response = chain.invoke({"question": prompt})
47
+ st.write(response)
48
+
49
+ tbl = con.sql(response).to_df()
50
+ st.dataframe(tbl)
51
+
52
+
53
+ # duckdb_sql fails but chatgpt3.5 succeeds with a query like:
54
+ # use the st_area function and st_GeomFromWKB functions to compute the area of the Shape column in the fee table, and then use that to compute the total area under each GAP_Sts category
55
+
56
+
57
+ # Federal agencies are identified as 'FED' in the Mang_Type column in the 'combined' data table. The Mang_Name column indicates the different agencies. Which federal agencies manage the greatest area of GAP_Sts 1 or 2 land?
58
+