Consultar con el agente Databricks SQL

Databricks con el SDK de Python

Avi Steinberg

Senior Software Engineer

Databricks Unity Catalog

Método principal para almacenar datos en Databricks
Contiene esquemas, que pueden incluir varias tablas
Crea tablas en un esquema desde muchas fuentes
Varios esquemas open source en el catálogo samples

Catálogo Unity de Databricks

¹ https://docs.databricks.com/aws/en/catalogs/

LangChain

Librería para crear y desplegar apps con LLM
Permite crear agentes de IA que usan SQL para consultar datos

Dependencias de Python para LangChain

Instala dependencias de Python:

pip install --upgrade databricks-langchain langchain-community langchain 
                      databricks-sql-connector databricks-sqlalchemy

Importa librerías:

from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from databricks_langchain import ChatDatabricks

ID de Databricks SQL Warehouse

Detalles de conexión del Databricks SQL Warehouse

Autenticar LangChain con el Workspace de Databricks

Exporta variables de entorno:

os.environ["DATABRICKS_TOKEN"] = "<Your-Access-Token>"
os.environ["DATABRICKS_HOST"] = "<your-workspace-id>.cloud.databricks.com"

Agente Databricks SQL

Haz preguntas a un agente de IA sobre datos en un esquema del catálogo de Databricks
El agente puede responder preguntas sobre los datos que requieren SQL
Especifica verbose=True para que el agente muestre su proceso

Salida de consulta de LangChain

Usar el agente Databricks SQL para consultar el Catálogo

db = SQLDatabase.from_databricks(
    catalog="samples", 
    schema="nyctaxi",
    warehouse_id=warehouse_id)
llm = ChatDatabricks(
    endpoint="databricks-meta-llama-3-3-70b-instruct",
    temperature=0.1,
    max_tokens=100)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agente = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)
# Consultar el agente Databricks SQL
resultado = agente.run("What's the time and distance of the longest trip?")
print(resultado)

¹ https://docs.databricks.com/aws/en/large-language-models/langchain#databricks-sql-agent

Crear una base de datos SQL de LangChain desde un Catálogo de Databricks

db = SQLDatabase.from_databricks(
    catalog="samples", 
    schema="nyctaxi",
    warehouse_id="<your-warehouse-id>"
)

¹ api.python.langchain.com/en/latest/utilities/langchain_community.utilities.sql_database.SQLDatabase.html

Crear un LLM de LangChain con un modelo base

llm = ChatDatabricks(

    endpoint="databricks-meta-llama-3-3-70b-instruct",

    temperature=0.1,
    max_tokens=100,
)

temperature: número entre 0 y 1 que controla la aleatoriedad de las respuestas
max_tokens: máximo de tokens en la respuesta

¹ https://docs.databricks.com/aws/en/large-language-models/langchain

Crear un toolkit de base de datos SQL de LangChain

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

¹ https://docs.databricks.com/aws/en/large-language-models/langchain

Crear un agente Databricks SQL

agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)

result = agent.run("What's the time and distance of the longest trip?")
display(result)

The average trip takes approximately 15 minutes.

¹ https://docs.databricks.com/aws/en/large-language-models/langchain

¡Consultemos datos de Databricks!

Databricks con el SDK de Python