Langchain sql agent not connecting to databricks unity catalog table

Question

I am trying to use the Langchain sql agent to connect to a table in databricks unity catalog. I am getting an error that InvalidRequestError: Resource not found

This is the code that I have written so far...

%pip install langchain==0.0.279 databricks-sql-connector openai==0.28.0

from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain import OpenAI
from langchain.chat_models import ChatOpenAI
import os
from langchain.agents.agent_types import AgentType
from langchain.chat_models import AzureChatOpenAI


db = SQLDatabase.from_databricks(catalog="dslakehouse", schema="default")
api_key = "XXX" 


MY_ENDPOINT = "https://XXX.azure.com/"
MY_DEPLOYMENT_NAME = "XXX"

llm = AzureChatOpenAI(
                  deployment_name=MY_DEPLOYMENT_NAME,
                  model_name='gpt-4o',
                  model_version='2024-05-13',
                  temperature = 0,
                  openai_api_key= api_key,
                  openai_api_base=MY_ENDPOINT,
                  openai_api_type='azure',
                  openai_api_version='2024-05-13',
                  max_retries=4,
)

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    engine="openai-ds-test-4o", 
)

The table i wan to interrogae is calle 'training_data' and is located in the default location of the dslakehouse.

When I try to interrogate the db using agent_executor.run("Question asked?")

I get a response that InvalidRequestError: Resource not found.

What am I doing wrong?

Answer 1

The below code works for me - connect data from data bricks unity catalog

from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_openai import AzureChatOpenAI
from langchain_openai import AzureOpenAIEmbeddings
        
    

    llm = AzureChatOpenAI(
                  openai_api_version= "",
                  azure_endpoint= "",
                  deployment_name= "",
                  api_key="",
                  temperature=0,
              )
            
            
            api_token = ""
            host = ""
            http_path = ""
            catalog = ""
            schema = ""
            
            uri = (
                f"databricks://token:{api_token}@{host}?"
                f"http_path={http_path}&catalog={catalog}&schema={schema}"
            )
            
            
            db = SQLDatabase.from_uri(
                database_uri=uri,
                include_tables=[""],
                sample_rows_in_table_info=5,
            )
        
        agent_executor = create_sql_agent(
            llm=llm,
            db=db,
            verbose=False,
            agent_executor_kwargs={
                "handle_parsing_errors": True,
            },
        )



agent_executor.invoke(
    {
        "input": "<quesion>"
    }
)