From b5429a2ca14bf3ef345b263fced339012bbbf441 Mon Sep 17 00:00:00 2001 From: Tanmaypatil123 Date: Tue, 24 Oct 2023 17:09:57 +0530 Subject: [PATCH] Added support for sqlite connectors --- docs/connectors.md | 22 +++++++++ examples/from_sql.py | 15 +++++- pandasai/connectors/__init__.py | 2 + pandasai/connectors/base.py | 9 ++++ pandasai/connectors/sql.py | 86 ++++++++++++++++++++++++++++++++- 5 files changed, 131 insertions(+), 3 deletions(-) diff --git a/docs/connectors.md b/docs/connectors.md index 1850b9a7c..58d11a970 100644 --- a/docs/connectors.md +++ b/docs/connectors.md @@ -89,6 +89,28 @@ df = SmartDataframe(mysql_connector) df.chat('What is the total amount of loans in the last year?') ``` +### Sqlite connector + +Similarly to the PostgreSQL and MySQL connectors, the Sqlite connector allows you to connect to a local Sqlite database file. It is designed to be easy to use, even if you are not familiar with Sqlite or with PandasAI. + +To use the Sqlite connector, you only need to import it into your Python code and pass it to a `SmartDataframe` or `SmartDatalake` object: + +```python +from pandasai.connectors import SqliteConnector + +connector = SqliteConnector(config={ + "database" : "PATH_TO_DB", + "table" : "actor", + "where" :[ + ["first_name","=","PENELOPE"] + ] +}) + +df = SmartDataframe(connector) +df.chat('How many records are there ?') +``` + + ### Generic SQL connector The generic SQL connector allows you to connect to any SQL database that is supported by SQLAlchemy. diff --git a/examples/from_sql.py b/examples/from_sql.py index c8d435559..3d1d36bd5 100644 --- a/examples/from_sql.py +++ b/examples/from_sql.py @@ -2,7 +2,7 @@ from pandasai import SmartDatalake from pandasai.llm import OpenAI -from pandasai.connectors import MySQLConnector, PostgreSQLConnector +from pandasai.connectors import MySQLConnector, PostgreSQLConnector,SqliteConnector # With a MySQL database loan_connector = MySQLConnector( @@ -38,8 +38,19 @@ } ) +# With a Sqlite databse + +invoice_connector = SqliteConnector( + config={ + "database" : "local_path_to_db", + "table" : "invoices", + "where" : [ + ["status","=","pending"] + ] + } +) llm = OpenAI() -df = SmartDatalake([loan_connector, payment_connector], config={"llm": llm}) +df = SmartDatalake([loan_connector, payment_connector,invoice_connector], config={"llm": llm}) response = df.chat("How many people from the United states?") print(response) # Output: 247 loans have been paid off by men. diff --git a/pandasai/connectors/__init__.py b/pandasai/connectors/__init__.py index fb80c8628..a484835ff 100644 --- a/pandasai/connectors/__init__.py +++ b/pandasai/connectors/__init__.py @@ -10,6 +10,7 @@ from .databricks import DatabricksConnector from .yahoo_finance import YahooFinanceConnector from .airtable import AirtableConnector +from .sql import SqliteConnector __all__ = [ "BaseConnector", @@ -20,4 +21,5 @@ "SnowFlakeConnector", "DatabricksConnector", "AirtableConnector", + "SqliteConnector", ] diff --git a/pandasai/connectors/base.py b/pandasai/connectors/base.py index 4de547486..1c487fc4d 100644 --- a/pandasai/connectors/base.py +++ b/pandasai/connectors/base.py @@ -39,6 +39,15 @@ class SQLBaseConnectorConfig(BaseConnectorConfig): dialect: Optional[str] = None +class SqliteConnectorConfig(SQLBaseConnectorConfig): + """ + Connector configurations for sqlite db. + """ + + table: str + database: str = "sqlite" + + class YahooFinanceConnectorConfig(BaseConnectorConfig): """ Connector configuration for Yahoo Finance. diff --git a/pandasai/connectors/sql.py b/pandasai/connectors/sql.py index 215da08f4..167a16236 100644 --- a/pandasai/connectors/sql.py +++ b/pandasai/connectors/sql.py @@ -5,7 +5,7 @@ import re import os import pandas as pd -from .base import BaseConnector, SQLConnectorConfig +from .base import BaseConnector, SQLConnectorConfig, SqliteConnectorConfig from .base import BaseConnectorConfig from sqlalchemy import create_engine, text, select, asc from sqlalchemy.engine import Connection @@ -364,6 +364,90 @@ def fallback_name(self): return self._config.table +class SqliteConnector(SQLConnector): + """ + Sqlite connector are used to connect to Sqlite databases. + """ + + def __init__(self, config: Union[SqliteConnectorConfig, dict]): + """ + Intialize the Sqlite connector with the given configuration. + + Args: + config (ConnectorConfig) : The configuration for the MySQL connector. + """ + config["dialect"] = "sqlite" + if isinstance(config, dict): + sqlite_env_vars = {"database": "SQLITE_DB_PATH", "table": "TABLENAME"} + config = self._populate_config_from_env(config, sqlite_env_vars) + + super().__init__(config) + + def _load_connector_config(self, config: Union[BaseConnectorConfig, dict]): + """ + Loads passed Configuration to object + + Args: + config (BaseConnectorConfig): Construct config in structure + + Returns: + config: BaseConenctorConfig + """ + return SqliteConnectorConfig(**config) + + def _init_connection(self, config: SqliteConnectorConfig): + """ + Initialize Database Connection + + Args: + config (SQLConnectorConfig): Configurations to load database + + """ + self._engine = create_engine(f"{config.dialect}:///{config.database}") + self._connection = self._engine.connect() + + def __del__(self): + """ + Close the connection to the SQL database. + """ + self._connection.close() + + @cache + def head(self): + """ + Return the head of the data source that the connector is connected to. + This information is passed to the LLM to provide the schema of the data source. + + Returns: + DataFrame: The head of the data source. + """ + + if self.logger: + self.logger.log( + f"Getting head of {self._config.table} " + f"using dialect {self._config.dialect}" + ) + + # Run a SQL query to get all the columns names and 5 random rows + query = self._build_query(limit=5, order="RANDOM()") + + # Return the head of the data source + return pd.read_sql(query, self._connection) + + def __repr__(self): + """ + Return the string representation of the SQL connector. + + Returns: + str: The string representation of the SQL connector. + """ + return ( + f"<{self.__class__.__name__} dialect={self._config.dialect}" + f"database={self._config.database} " + f"table={self._config.table}>" + ) + + class MySQLConnector(SQLConnector): """ MySQL connectors are used to connect to MySQL databases.