Skip to content

Commit

Permalink
Added support for sqlite connectors
Browse files Browse the repository at this point in the history
  • Loading branch information
Tanmaypatil123 committed Oct 24, 2023
1 parent 0d56e93 commit b5429a2
Show file tree
Hide file tree
Showing 5 changed files with 131 additions and 3 deletions.
22 changes: 22 additions & 0 deletions docs/connectors.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
15 changes: 13 additions & 2 deletions examples/from_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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.
2 changes: 2 additions & 0 deletions pandasai/connectors/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from .databricks import DatabricksConnector
from .yahoo_finance import YahooFinanceConnector
from .airtable import AirtableConnector
from .sql import SqliteConnector

__all__ = [
"BaseConnector",
Expand All @@ -20,4 +21,5 @@
"SnowFlakeConnector",
"DatabricksConnector",
"AirtableConnector",
"SqliteConnector",
]
9 changes: 9 additions & 0 deletions pandasai/connectors/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
86 changes: 85 additions & 1 deletion pandasai/connectors/sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit b5429a2

Please sign in to comment.