Databricks SQL Agent is failing

SandipCoder
New Contributor II

Goal :

I want to build one Databricks SQL Agent using that any Databricks users can interact with a specified schema in Unity Catalog and generate insights on their data.

Followed the url as reference :https://docs.databricks.com/en/large-language-models/langchain.html

SampleCode:

from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.sql_database import SQLDatabase
from langchain import OpenAI

db = SQLDatabase.from_databricks(catalog="genai", schema="vector_db")
#llm = OpenAI(temperature=.7)
toolkit = SQLDatabaseToolkit(db=db, llm=llm_handler)
agent = create_sql_agent(llm=llm_handler, toolkit=toolkit, verbose=True)

Error :

AttributeError: module 'sqlalchemy.types' has no attribute 'Uuid'
File <command-4275727126297453>, line 6
3 from langchain.sql_database import SQLDatabase
4 from langchain import OpenAI
----> 6 db = SQLDatabase.from_databricks(catalog="genai", schema="vector_db")
7 #llm = OpenAI(temperature=.7)
8 toolkit = SQLDatabaseToolkit(db=db, llm=llm_handler)
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-21fa55bd-1c7c-42e6-8617-6775fdb15bae/lib/python3.11/site-packages/databricks/sqlalchemy/_types.py:35
13 """This method is supposed to accept a Python type and return a string representation of that type.
14 But due to some weirdness in the way SQLAlchemy's literal rendering works, we have to return
15 the value itself because, by the time it reaches our custom type code, it's already been converted

 

Sandip Bhowmick