initial

agent.py

@@ -1,8 +1,10 @@
-from omegaconf import OmegaConf
 import os
-
 from typing import Optional
 from pydantic import Field, BaseModel
+from omegaconf import OmegaConf
+
+from llama_index.core.utilities.sql_wrapper import SQLDatabase
+from sqlalchemy import create_engine
 
 from dotenv import load_dotenv
 load_dotenv(override=True)
@@ -71,11 +73,7 @@ def create_assistant_tools(cfg):
             tools_factory.guardrail_tools() +
             tools_factory.database_tools(
                 content_description = 'Electric Vehicles',
-                scheme = 'postgresql',
-                host = 'localhost', port = '5432', 
-                user = 'ofer', 
-                password = 'noanoa', 
-                dbname = 'ev_database'
+                sql_database = SQLDatabase(create_engine('sqlite:///ev_database.db')),
             ) +
             [ask_vehicles, ask_policies]
     )
@@ -83,9 +81,11 @@ def create_assistant_tools(cfg):
 def initialize_agent(_cfg, update_func):
     electric_vehicle_bot_instructions = """
     - You are a helpful research assistant, with expertise in electric vehicles, in conversation with a user.
+    - Before answering any user query, get sample data from each table in the database, so that you can understand NULL and unique values for each column.
     - For a query with multiple sub-questions, break down the query into the sub-questions, 
       and make separate calls to the ask_vehicles or ask_policies tool to answer each sub-question,
       then combine the answers to provide a complete response.
+    - Use the database tools to answer analytical questions.
     - Never discuss politics, and always respond politely.
     """
 
@@ -95,6 +95,7 @@ def initialize_agent(_cfg, update_func):
         custom_instructions=electric_vehicle_bot_instructions,
         update_func=update_func
     )
+    agent.report()
     return agent
 
 
@@ -104,7 +105,6 @@ def get_agent_config() -> OmegaConf:
         'corpus_ids': str(os.environ['VECTARA_CORPUS_IDS']).split(','),
         'api_keys': str(os.environ['VECTARA_API_KEYS']).split(','),
         'examples': os.environ.get('QUERY_EXAMPLES', None),
-        'title': "Electric Vehicles in the United States",
         'demo_welcome': "Welcome to the EV Assistant demo.",
         'demo_description': "This assistant can help you learn about electric vehicles in the United States, including how they work, the advantages of purchasing them, and reviews on the top choices.",
     })

app.py

@@ -4,6 +4,10 @@ import sys
 import streamlit as st
 from streamlit_pills import pills
 
+import sqlite3
+import pandas as pd
+from datasets import load_dataset
+
 from vectara_agent.agent import AgentStatusType
 from agent import initialize_agent, get_agent_config
 
@@ -47,7 +51,6 @@ def launch_bot():
         reset()
 
     cfg = st.session_state.cfg
-    st.set_page_config(page_title=cfg['title'], layout="wide")
 
     # left side content
     with st.sidebar:
@@ -126,5 +129,50 @@ def launch_bot():
 
     sys.stdout.flush()
 
+def setup_db():
+    db_path = 'ev_database.db'
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()        
+
+    with st.spinner("Loading data... Please wait..."):
+        def tables_populated() -> bool:
+            tables = ['ev_population', 'county_registrations', 'ev_registrations']        
+            for table in tables:
+                cursor.execute(f"SELECT name FROM sqlite_master WHERE type='table' AND name='{table}'")
+                result = cursor.fetchone()
+                if not result:
+                    return False
+            return True            
+
+        if tables_populated():
+            print("Database tables already populated, skipping setup")
+            conn.close()
+            return
+        else:
+            print("Populating database tables")
+
+        # Execute the SQL commands to create tables
+        with open('create_tables.sql', 'r') as sql_file:
+            sql_script = sql_file.read()
+            cursor.executescript(sql_script)
+
+        # Load data into ev_population table
+        df = load_dataset("vectara/ev-dataset", data_files="Electric_Vehicle_Population_Data.csv")['train'].to_pandas()
+        df.to_sql('ev_population', conn, if_exists='replace', index=False)
+
+        # Load data into county_registrations table
+        df = load_dataset("vectara/ev-dataset", data_files="Electric_Vehicle_Population_Size_History_By_County.csv")['train'].to_pandas()
+        df.to_sql('county_registrations', conn, if_exists='replace', index=False)
+
+        # Load data into ev_registrations table
+        df = load_dataset("vectara/ev-dataset", data_files="Electric_Vehicle_Title_and_Registration_Activity.csv")['train'].to_pandas()
+        df.to_sql('ev_registrations', conn, if_exists='replace', index=False)
+
+        # Commit changes and close connection
+        conn.commit()
+        conn.close()
+
 if __name__ == "__main__":
+    st.set_page_config(page_title="Electric Vehicles Assistant", layout="wide")
+    setup_db()
     launch_bot()

create_tables.sql

@@ -0,0 +1,67 @@
+CREATE TABLE ev_population (
+    vin VARCHAR(10), 
+    county VARCHAR(20), 
+    city VARCHAR(24), 
+    state VARCHAR(2), 
+    postal_code INTEGER, 
+    model_year INTEGER, 
+    make VARCHAR(20), 
+    model VARCHAR(24), 
+    ev_type VARCHAR(38), 
+    cafv_eligibility VARCHAR(60), 
+    electric_range INTEGER, 
+    base_msrp INTEGER, 
+    legislative_district INTEGER, 
+    dol_vehicle_id INTEGER, 
+    electric_utility VARCHAR(112)
+);
+
+CREATE TABLE county_registrations (
+    date DATE, 
+    county VARCHAR(20), 
+    state VARCHAR(2), 
+    primary_use VARCHAR(9), 
+    battery_evs INTEGER, 
+    plug_in_hybrids INTEGER, 
+    total_evs INTEGER, 
+    total_non_evs INTEGER, 
+    total_vehicles INTEGER, 
+    percent_evs REAL
+);
+
+CREATE TABLE ev_registrations (
+    ev_type VARCHAR(38), 
+    vin VARCHAR(10), 
+    dol_vehicle_id INTEGER, 
+    model_year INTEGER, 
+    make VARCHAR(20), 
+    model VARCHAR(24), 
+    primary_use VARCHAR(34), 
+    electric_range INTEGER, 
+    odometer_reading INTEGER, 
+    odometer_reading_description VARCHAR(55), 
+    new_or_used VARCHAR(4), 
+    sale_price INTEGER, 
+    sale_date DATE, 
+    base_msrp INTEGER, 
+    transaction_type VARCHAR(32), 
+    transaction_date DATE, 
+    year INTEGER, 
+    county VARCHAR(20), 
+    city VARCHAR(24), 
+    state VARCHAR(2), 
+    postal_code INTEGER, 
+    cafv_eligibility VARCHAR(40), 
+    meets_2019_hb_2042_electric_range_requirement BOOLEAN, 
+    meets_2019_hb_2042_sale_date_requirement BOOLEAN, 
+    meets_2019_hb_2042_sale_price_value_requirement BOOLEAN, 
+    battery_range_requirement_2019_hb_2042 VARCHAR(32), 
+    purchase_date_requirement_2019_hb_2042 VARCHAR(59), 
+    sale_price_value_requirement_2019_hb_2042 VARCHAR(59), 
+    ev_fee_paid VARCHAR(14), 
+    transportation_electrification_fee_paid VARCHAR(14), 
+    hybrid_vehicle_electrificatin_fee_paid VARCHAR(14), 
+    geoid_2020 INTEGER, 
+    legislative_district INTEGER, 
+    electric_utility VARCHAR(112)
+);

requirements.txt

@@ -3,5 +3,5 @@ pydantic==1.10.15
 python-dotenv==1.0.1
 streamlit==1.32.2
 streamlit_pills==0.3.0
-psycopg2-binary==2.9.9
+datasets==2.14.7
 git+https://{GITHUB_TOKEN}@github.com/vectara/vectara-agent.git