Version 1.2

examples/from_sql.py

@@ -0,0 +1,45 @@
+"""Example of using PandasAI with a CSV file."""
+
+from pandasai import SmartDatalake
+from pandasai.llm import OpenAI
+from pandasai.connectors import MySQLConnector, PostgreSQLConnector
+
+# With a MySQL database
+loan_connector = MySQLConnector(
+    config={
+        "host": "localhost",
+        "port": 3306,
+        "database": "mydb",
+        "username": "root",
+        "password": "root",
+        "table": "loans",
+        "where": [
+            # this is optional and filters the data to
+            # reduce the size of the dataframe
+            ["loan_status", "=", "PAIDOFF"],
+        ],
+    }
+)
+
+# With a PostgreSQL database
+payment_connector = PostgreSQLConnector(
+    config={
+        "host": "localhost",
+        "port": 5432,
+        "database": "mydb",
+        "username": "root",
+        "password": "root",
+        "table": "payments",
+        "where": [
+            # this is optional and filters the data to
+            # reduce the size of the dataframe
+            ["payment_status", "=", "PAIDOFF"],
+        ],
+    }
+)
+
+llm = OpenAI()
+df = SmartDatalake([loan_connector, payment_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.

pandasai/connectors/__init__.py

@@ -0,0 +1,17 @@
+"""
+Connectors are used to connect to databases, external APIs, and other data sources.
+
+The connectors package contains all the connectors that are used by the application.
+"""
+
+from .base import BaseConnector
+from .sql import SQLConnector, MySQLConnector, PostgreSQLConnector
+from .yahoo_finance import YahooFinanceConnector
+
+__all__ = [
+    "BaseConnector",
+    "SQLConnector",
+    "MySQLConnector",
+    "PostgreSQLConnector",
+    "YahooFinanceConnector",
+]

pandasai/connectors/base.py

@@ -0,0 +1,136 @@
+"""
+Base connector class to be extended by all connectors.
+"""
+
+from abc import ABC, abstractmethod
+from ..helpers.df_info import DataFrameType
+from ..helpers.logger import Logger
+from pydantic import BaseModel
+from typing import Optional
+
+
+class ConnectorConfig(BaseModel):
+    """
+    Connector configuration.
+    """
+
+    dialect: Optional[str] = None
+    driver: Optional[str] = None
+    username: str
+    password: str
+    host: str
+    port: int
+    database: str
+    table: str
+    where: list[list[str]] = None
+
+
+class BaseConnector(ABC):
+    """
+    Base connector class to be extended by all connectors.
+    """
+
+    _config = None
+    _logger: Logger = None
+    _additional_filters: list[list[str]] = None
+
+    def __init__(self, config):
+        """
+        Initialize the connector with the given configuration.
+
+        Args:
+            config (dict): The configuration for the connector.
+        """
+        self._config = config
+
+    @abstractmethod
+    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.
+        """
+        pass
+
+    @abstractmethod
+    def execute(self) -> DataFrameType:
+        """
+        Execute the given query on the data source that the connector is
+        connected to.
+        """
+        pass
+
+    def set_additional_filters(self, filters: dict):
+        """
+        Add additional filters to the connector.
+
+        Args:
+            filters (dict): The additional filters to add to the connector.
+        """
+        self._additional_filters = filters if filters else []
+
+    @property
+    def rows_count(self):
+        """
+        Return the number of rows in the data source that the connector is
+        connected to.
+        """
+        raise NotImplementedError
+
+    @property
+    def columns_count(self):
+        """
+        Return the number of columns in the data source that the connector is
+        connected to.
+        """
+        raise NotImplementedError
+
+    @property
+    def column_hash(self):
+        """
+        Return the hash code that is unique to the columns of the data source
+        that the connector is connected to.
+        """
+        raise NotImplementedError
+
+    @property
+    def path(self):
+        """
+        Return the path of the data source that the connector is connected to.
+        """
+        # JDBC string
+        return (
+            self.__class__.__name__
+            + "://"
+            + self._config.host
+            + ":"
+            + str(self._config.port)
+            + "/"
+            + self._config.database
+            + "/"
+            + self._config.table
+        )
+
+    @property
+    def logger(self):
+        """
+        Return the logger for the connector.
+        """
+        return self._logger
+
+    @logger.setter
+    def logger(self, logger: Logger):
+        """
+        Set the logger for the connector.
+
+        Args:
+            logger (Logger): The logger for the connector.
+        """
+        self._logger = logger
+
+    @property
+    def fallback_name(self):
+        """
+        Return the name of the table that the connector is connected to.
+        """
+        raise NotImplementedError