Skip to content

Chat Open AI Databricks LLM Connector

ChatOpenAIDatabricksConnection

Bases: ConnectionInterface

The Chat Open AI(Chat GPT) Databricks LLM Connector enables you to connect to a Databricks SQL Warehouse and use the Chat Open AI(Chat GPT) LLM to generate SQL queries.

The connection class represents a connection to a database and uses the Databricks SQL Connector API's for Python to interact with cluster/jobs and langchain to connect to Chat Open AI(Chat GPT) LLM. To find details for SQL warehouses server_hostname and http_path location to the SQL Warehouse tab in the documentation.

Parameters:

Name Type Description Default
catalog str

Catalog name in Databricks

 required
schema str

Schema name in Databricks

 required
server_hostname str

Server hostname for the cluster or SQL Warehouse

 required
http_path str

Http path for the cluster or SQL Warehouse

 required
access_token str

Azure AD or Databricks PAT token

 required
openai_api_key str

OpenAI API key

 required
openai_model str

OpenAI model name

 'gpt-4'
sample_rows_in_table_info int

Number of rows to sample when getting table information

 3
verbose_logging bool

Whether to log verbose messages

 False
Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
class ChatOpenAIDatabricksConnection(ConnectionInterface):
    """
    The Chat Open AI(Chat GPT) Databricks LLM Connector enables you to connect to a Databricks SQL Warehouse and use the Chat Open AI(Chat GPT) LLM to generate SQL queries.

    The connection class represents a connection to a database and uses the Databricks SQL Connector API's for Python to interact with cluster/jobs and langchain to connect to Chat Open AI(Chat GPT) LLM.
    To find details for SQL warehouses server_hostname and http_path location to the SQL Warehouse tab in the documentation.

    Args:
        catalog: Catalog name in Databricks
        schema: Schema name in Databricks
        server_hostname: Server hostname for the cluster or SQL Warehouse
        http_path: Http path for the cluster or SQL Warehouse
        access_token: Azure AD or Databricks PAT token
        openai_api_key: OpenAI API key
        openai_model: OpenAI model name
        sample_rows_in_table_info: Number of rows to sample when getting table information
        verbose_logging: Whether to log verbose messages
    """

    def __init__(
        self,
        catalog: str,
        schema: str,
        server_hostname: str,
        http_path: str,
        access_token: str,
        openai_api_key: str,
        openai_model: str = "gpt-4",
        sample_rows_in_table_info: int = 3,
        verbose_logging: bool = False,
    ) -> None:
        _package_version_meets_minimum("langchain", "0.0.196")
        # connect to llm
        llm = ChatOpenAI(
            temperature=0, model_name=openai_model, openai_api_key=openai_api_key
        )

        # connect to Databricks
        db = SQLDatabase.from_databricks(
            catalog=catalog,
            schema=schema,
            api_token=access_token,
            host=server_hostname,
            warehouse_id=http_path.split("/")[-1],
            sample_rows_in_table_info=sample_rows_in_table_info,
        )

        prefix = """
    ...
    Always adhere to the format and don't return empty names or half responses.
    """
        toolkit = SQLDatabaseToolkit(db=db, llm=llm)

        model_agent_type = AgentType.ZERO_SHOT_REACT_DESCRIPTION
        if "0613" in openai_model:
            model_agent_type = AgentType.OPENAI_FUNCTIONS

        self.connection = create_sql_agent(
            llm=llm,
            prefix=prefix,
            toolkit=toolkit,
            verbose=verbose_logging,
            agent_type=model_agent_type,
        )

    def close(self) -> None:
        """Closes connection to database."""
        pass

    def cursor(self) -> object:
        """
        Initiates the cursor and returns it.

        Returns:
          ChatOpenAIDatabricksSQLCursor: Object to represent a connection to Databricks and Open AI with methods to interact with clusters/jobs and ChatGPT.
        """
        try:
            return ChatOpenAIDatabricksSQLCursor(self.connection)
        except Exception as e:
            logging.exception("error with cursor object")
            raise e

    def run(self, query: str) -> str:
        """
        Runs a query on the ChatGPT and the Databricks Cluster or SQL Warehouse.

        Returns:
            str: response from ChatGPT and the Databricks Cluster or SQL Warehouse
        """
        cursor = self.cursor()
        cursor.execute(query)
        return cursor.fetch_all()

close()

Closes connection to database.

Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
92
93
94
def close(self) -> None:
    """Closes connection to database."""
    pass

cursor()

Initiates the cursor and returns it.

Returns:

Name Type Description
ChatOpenAIDatabricksSQLCursor object

Object to represent a connection to Databricks and Open AI with methods to interact with clusters/jobs and ChatGPT.
Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def cursor(self) -> object:
    """
    Initiates the cursor and returns it.

    Returns:
      ChatOpenAIDatabricksSQLCursor: Object to represent a connection to Databricks and Open AI with methods to interact with clusters/jobs and ChatGPT.
    """
    try:
        return ChatOpenAIDatabricksSQLCursor(self.connection)
    except Exception as e:
        logging.exception("error with cursor object")
        raise e

run(query)

Runs a query on the ChatGPT and the Databricks Cluster or SQL Warehouse.

Returns:

Name Type Description
str str

response from ChatGPT and the Databricks Cluster or SQL Warehouse
Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
109
110
111
112
113
114
115
116
117
118
def run(self, query: str) -> str:
    """
    Runs a query on the ChatGPT and the Databricks Cluster or SQL Warehouse.

    Returns:
        str: response from ChatGPT and the Databricks Cluster or SQL Warehouse
    """
    cursor = self.cursor()
    cursor.execute(query)
    return cursor.fetch_all()

ChatOpenAIDatabricksSQLCursor

Bases: CursorInterface

Object to represent a connection to Databricks and Open AI with methods to interact with clusters/jobs and ChatGPT.

Parameters:

Name Type Description Default
cursor object

controls execution of commands on cluster or SQL Warehouse

 required
Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
class ChatOpenAIDatabricksSQLCursor(CursorInterface):
    """
    Object to represent a connection to Databricks and Open AI with methods to interact with clusters/jobs and ChatGPT.

    Args:
        cursor: controls execution of commands on cluster or SQL Warehouse
    """

    response = None

    def __init__(self, cursor: object) -> None:
        self.cursor = cursor

    def execute(self, query: str) -> None:
        """
        Prepares and runs a database query.

        Args:
            query: sql query to execute on ChatGPT and the Databricks Cluster or SQL Warehouse
        """
        try:
            self.response = self.cursor.run(query)
        except Exception as e:
            logging.exception("error while executing the query")
            raise e

    def fetch_all(
        self,
    ) -> list:
        """
        Gets all rows of a query.

        Returns:
            list: list of results
        """
        try:
            return self.response
        except Exception as e:
            logging.exception("error while fetching the rows of a query")
            raise e

    def close(self) -> None:
        """Closes the cursor."""
        pass

execute(query)

Prepares and runs a database query.

Parameters:

Name Type Description Default
query str

sql query to execute on ChatGPT and the Databricks Cluster or SQL Warehouse

 required
Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
134
135
136
137
138
139
140
141
142
143
144
145
def execute(self, query: str) -> None:
    """
    Prepares and runs a database query.

    Args:
        query: sql query to execute on ChatGPT and the Databricks Cluster or SQL Warehouse
    """
    try:
        self.response = self.cursor.run(query)
    except Exception as e:
        logging.exception("error while executing the query")
        raise e

fetch_all()

Gets all rows of a query.

Returns:

Name Type Description
list list

list of results
Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def fetch_all(
    self,
) -> list:
    """
    Gets all rows of a query.

    Returns:
        list: list of results
    """
    try:
        return self.response
    except Exception as e:
        logging.exception("error while fetching the rows of a query")
        raise e

close()

Closes the cursor.

Source code in src/sdk/python/rtdip_sdk/connectors/llm/chatopenai_databricks_connector.py 
162
163
164
def close(self) -> None:
    """Closes the cursor."""
    pass