DatabaseError (databricks.sql.exc.DatabaseError) Invalid SessionHandle in LangChain SQL chain

[JUNNIE_J]

When creating a SQL Chatbot with LangChain SQLDatabase chain, Databricks database, and OpenAI model, I got a "Invalid SessionHandle" Error.

Code of creating the SQL Chain:

<code>

llm = ChatOpenAI(temperature=0, openai_api_key=api_key, model_name='gpt-4-0125-preview')

def get_db(_=None😞

    db = SQLDatabase.from_databricks(catalog="_", schema="_", host="_", api_token="_", cluster_id="_", sample_rows_in_table_info=3)

    return db

db = get_db()

def get_chain():

    db = get_db()

    chain = SQLDatabaseChain.from_llm(llm=llm, db=db, verbose=True)

    return chain

chain = get_chain()

</code>

Code of deploying on MLflow:

<code>

mlflow.set_registry_uri("databricks-uc")

model_name = f"SQL_chain"

with mlflow.start_run(run_name="sql_chain") as run:

    user_query = "..."

    answer = chain.invoke(user_query)

    answer_json = json.dumps(answer)

    signature = infer_signature(user_query, answer_json)

    model_info = mlflow.langchain.log_model(

        chain,

        loader_fn=get_db,

        artifact_path="chain",

        registered_model_name=model_name,

        pip_requirements=[

            f"mlflow=={mlflow.__version__}",

            f"langchain=={langchain.__version__}",

            f"langchain_community=={langchain_community.__version__}",

            f"langchain_experimental=={langchain_experimental.__version__}",

            f"openai=={openai.__version__}",

            f"langchain_core=={langchain_core.__version__}",

            f"databricks",

            f"databricks-sql-connector==2.9.3",

            ],

        input_example=user_query,

        signature=signature

    )

</code>

[JUNNIE_J]

I have tried the solutions. But they don't work.

When i tried :

from pyspark.sql import SparkSession
from databricks import sql

def get_db(_=None):
db = SQLDatabase.from_databricks(...)

spark = SparkSession.builder.appName("NewSession").getOrCreate()
sql_context = sql.Context(spark.sparkContext)
sql_context.sessionState().newSession()

return db

It returned : AttributeError: module 'databricks.sql' has no attribute 'Context'.

When I tried :

def get_chain():
global db
if db is None or not db.is_valid():
db = get_db()
chain = SQLDatabaseChain.from_llm(llm=llm, db=db, verbose=True)
return chain

It returned: AttributeError: 'SQLDatabase' object has no attribute 'is_valid'.