Change the default dialect to SQLite (#13)

Author: narasimhard

Makefile

@@ -9,6 +9,7 @@ setup: download_models ## Setup
 	./.sidekickvenv/bin/python3 -m pip install --upgrade pip
 	./.sidekickvenv/bin/python3 -m pip install wheel
 	./.sidekickvenv/bin/python3 -m pip install -r requirements.txt
+	mkdir -p ./db/sqlite
 
 download_models:
 	mkdir -p ./models/sentence_transformers/sentence-transformers_all-MiniLM-L6-v2

sidekick/configs/.env.toml

@@ -4,7 +4,7 @@ MODEL_NAME = "gpt-3.5-turbo-0301" # Others: e.g. gpt-4, gpt-4-32k, text-davinci-
 
 [LOCAL_DB_CONFIG]
 HOST_NAME = "localhost"
-USER_NAME = "postgres"
+USER_NAME = "sqlite"
 PASSWORD = "abc"
 DB_NAME = "querydb"
 PORT = "5432"
@@ -13,4 +13,4 @@ PORT = "5432"
 LOG-LEVEL = "INFO"
 
 [DB-DIALECT]
-DB_TYPE = "postgresql"
+DB_TYPE = "sqlite"

sidekick/configs/prompt_template.py

@@ -50,7 +50,7 @@
             """
 
 DEBUGGING_PROMPT = {
-    "system_prompt": "Act as a SQL expert for PostgreSQL code",
+    "system_prompt": "Act as a SQL expert for {_dialect} code",
     "user_prompt": """
                 ### Fix syntax errors for provided incorrect SQL Query.
                 # Add ``` as prefix and ``` as suffix to generated SQL

sidekick/db_config.py

@@ -23,7 +23,7 @@ def __init__(
         base_path,
         schema_info_path=None,
         schema_info=None,
-        dialect="postgresql",
+        dialect="sqlite",
     ) -> None:
         self.db_name = db_name
         self.hostname = hostname
@@ -36,7 +36,10 @@ def __init__(
         self._engine = None
         self.dialect = dialect
         self.base_path = base_path
-        self._url = f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/"
+        if dialect == "sqlite":
+            self._url = f"sqlite:///{base_path}/db/sqlite/{db_name}.db"
+        else:
+            self._url = f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/"
 
     @property
     def table_name(self):
@@ -51,18 +54,26 @@ def engine(self):
         return self._engine
 
     def db_exists(self):
-        engine = create_engine(f"{self._url}{self.db_name}", echo=True)
+        if self.dialect == "sqlite":
+            engine = create_engine(f"{self._url}", echo=True)
+        else:
+            engine = create_engine(f"{self._url}{self.db_name}", echo=True)
         return database_exists(f"{engine.url}")
 
     def create_db(self):
         engine = create_engine(self._url)
         self._engine = engine
 
         with engine.connect() as conn:
-            conn.execute("commit")
+            # conn.execute("commit")
             # Do not substitute user-supplied database names here.
-            res = conn.execute(f"CREATE DATABASE {self.db_name}")
-        return res
+            if self.dialect != "sqlite":
+                conn.execute("commit")
+                res = conn.execute(f"CREATE DATABASE {self.db_name}")
+                self._url = f"{self._url}{self.db_name}"
+                return res
+            else:
+                logger.debug("SQLite DB is created when 'engine.connect()' is called")
 
     def _extract_schema_info(self, schema_info_path=None):
         # From jsonl format
@@ -91,7 +102,7 @@ def _extract_schema_info(self, schema_info_path=None):
 
     def create_table(self, schema_info_path=None, schema_info=None):
         engine = create_engine(
-            f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/{self.db_name}"
+            self._url, isolation_level="AUTOCOMMIT"
         )
         self._engine = engine
         if self.schema_info is None:
@@ -119,20 +130,22 @@ def create_table(self, schema_info_path=None, schema_info=None):
                 )
                 """
         with engine.connect() as conn:
-            conn.execute("commit")
+            if self.dialect != "sqlite":
+                conn.execute("commit")
             conn.execute(create_syntax)
         return
 
     def has_table(self):
         engine = create_engine(
-            f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/{self.db_name}"
+            self._url
         )
+
         return sqlalchemy.inspect(engine).has_table(self.table_name)
 
     def add_samples(self, data_csv_path=None):
-        conn_str = f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/{self.db_name}"
+        conn_str = self._url
         try:
-            df = pd.read_csv(data_csv_path, infer_datetime_format=True)
+            df = pd.read_csv(data_csv_path)
             engine = create_engine(conn_str, isolation_level="AUTOCOMMIT")
 
             sample_query = f"SELECT COUNT(*) AS ROWS FROM {self.table_name} LIMIT 1"
@@ -153,12 +166,14 @@ def add_samples(self, data_csv_path=None):
 
     def execute_query_db(self, query=None, n_rows=100):
         output = []
+        if self.dialect != "sqlite":
+            conn_str = f"{self._url}{self.db_name}"
+        else:
+            conn_str = self._url
+
         try:
             if query:
                 # Create an engine
-                conn_str = (
-                    f"{self.dialect}://{self.user_name}:{self.password}@{self.hostname}:{self.port}/{self.db_name}"
-                )
                 engine = create_engine(conn_str)
 
                 # Create a connection

sidekick/prompter.py

@@ -112,10 +112,10 @@ def update_table_info(cache_path: str, table_info_path: str = None, table_name:
         json.dump(table_metadata, outfile, indent=4, sort_keys=False)
 
 
-@configure.command("db-setup", help="Enter information to configure postgres database locally")
+@configure.command("db-setup", help=f"Enter information to configure {db_dialect} database locally")
 @click.option("--db_name", "-n", default="querydb", help="Database name", prompt="Enter Database name")
 @click.option("--hostname", "-h", default="localhost", help="Database hostname", prompt="Enter hostname name")
-@click.option("--user_name", "-u", default="postgres", help="Database username", prompt="Enter username name")
+@click.option("--user_name", "-u", default=f"{db_dialect}", help="Database username", prompt="Enter username name")
 @click.option(
     "--password",
     "-p",
@@ -141,8 +141,11 @@ def db_setup(db_name: str, hostname: str, user_name: str, password: str, port: i
         f.close()
         path = f"{base_path}/var/lib/tmp/data"
         # For current session
-        db_obj = DBConfig(db_name, hostname, user_name, password, port, base_path=base_path)
-        if not db_obj.db_exists():
+        db_obj = DBConfig(db_name, hostname, user_name, password, port, base_path=base_path, dialect=db_dialect)
+        if db_obj.dialect == 'sqlite' and not os.path.isfile(f"{base_path}/db/sqlite/{db_name}.db"):
+            db_obj.create_db()
+            click.echo("Database created successfully!")
+        elif not db_obj.db_exists():
             db_obj.create_db()
             click.echo("Database created successfully!")
         else:
@@ -293,9 +296,12 @@ def query(question: str, table_info_path: str, sample_queries: str):
     passwd = env_settings["LOCAL_DB_CONFIG"]["PASSWORD"]
     db_name = env_settings["LOCAL_DB_CONFIG"]["DB_NAME"]
 
-    db_url = f"{db_dialect}+psycopg2://{user_name}:{passwd}@{host_name}/{db_name}".format(
-        user_name, passwd, host_name, db_name
-    )
+    if db_dialect == "sqlite":
+        db_url = f"sqlite:///{base_path}/db/sqlite/{db_name}.db"
+    else:
+        db_url = f"{db_dialect}+psycopg2://{user_name}:{passwd}@{host_name}/{db_name}".format(
+            user_name, passwd, host_name, db_name
+        )
 
     if table_info_path is None:
         table_info_path = _get_table_info(path)
@@ -318,7 +324,7 @@ def query(question: str, table_info_path: str, sample_queries: str):
         sql_g._tasks = updated_tasks
 
     model_name = env_settings["OPENAI"]["MODEL_NAME"]
-    res = sql_g.generate_sql(table_names, question, model_name=model_name)
+    res = sql_g.generate_sql(table_names, question, model_name=model_name, _dialect=db_dialect)
     logger.info(f"Input query: {question}")
     logger.info(f"Generated response:\n\n{res}")
 
@@ -335,7 +341,7 @@ def query(question: str, table_info_path: str, sample_queries: str):
                 click.echo(f"Updated SQL:\n {updated_sql}")
             elif res_val.lower() == "r" or res_val.lower() == "regenerate":
                 click.echo("Attempting to regenerate...")
-                res = sql_g.generate_sql(table_names, question, model_name=model_name)
+                res = sql_g.generate_sql(table_names, question, model_name=model_name, _dialect=db_dialect)
                 logger.info(f"Input query: {question}")
                 logger.info(f"Generated response:\n\n{res}")
 
@@ -351,7 +357,8 @@ def query(question: str, table_info_path: str, sample_queries: str):
                 port = env_settings["LOCAL_DB_CONFIG"]["PORT"]
                 db_name = env_settings["LOCAL_DB_CONFIG"]["DB_NAME"]
 
-                db_obj = DBConfig(db_name, hostname, user_name, password, port, base_path=base_path)
+                db_obj = DBConfig(db_name, hostname, user_name, password, port, base_path=base_path, dialect=db_dialect)
+
                 output_res = db_obj.execute_query_db(query=_val)
                 click.echo(f"The query results are:\n {output_res}")
             elif option == "pandas":

sidekick/query.py

@@ -123,7 +123,7 @@ def _query_tasks(self, question_str, data_info, sample_queries, table_name: list
             res = ex_value.statement if ex_value.statement else None
             return res
 
-    def generate_response(self, context_container, sql_index, input_prompt, attempt_fix_on_error: bool = True):
+    def generate_response(self, context_container, sql_index, input_prompt, attempt_fix_on_error: bool = True, _dialect: str = "sqlite"):
         try:
             response = sql_index.query(input_prompt, sql_context_container=context_container)
             res = response.extra_info["sql_query"]
@@ -137,7 +137,7 @@ def generate_response(self, context_container, sql_index, input_prompt, attempt_
                     # Attempt to heal with simple feedback
                     # Reference: Teaching Large Language Models to Self-Debug, https://arxiv.org/abs/2304.05128
                     logger.info(f"Attempting to fix syntax error ...,\n {se}")
-                    system_prompt = DEBUGGING_PROMPT["system_prompt"]
+                    system_prompt = DEBUGGING_PROMPT["system_prompt"].format(_dialect=_dialect)
                     user_prompt = DEBUGGING_PROMPT["user_prompt"].format(ex_traceback=ex_traceback, qry_txt=qry_txt)
                     # Role and content
                     query_msg = [{"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}]
@@ -192,7 +192,7 @@ def generate_tasks(self, table_names: list, input_question: str):
             raise se
 
     def generate_sql(
-        self, table_name: list, input_question: str, _dialect: str = "postgres", model_name: str = "gpt-3.5-turbo-0301"
+        self, table_name: list, input_question: str, _dialect: str = "sqlite", model_name: str = "gpt-3.5-turbo-0301"
     ):
         _tasks = self.task_formatter(self._tasks)
         context_file = f"{self.path}/var/lib/tmp/data/context.json"
@@ -223,7 +223,7 @@ def generate_sql(
         index = GPTSQLStructStoreIndex(
             [], sql_database=self.sql_database, table_name=table_name, service_context=service_context_gpt3
         )
-        res = self.generate_response(context_container, sql_index=index, input_prompt=query_str)
+        res = self.generate_response(context_container, sql_index=index, input_prompt=query_str, _dialect = _dialect)
         try:
             # Check if `SQL` is formatted ---> ``` SQL_text ```
             if "```" in str(res):