Databricks SQL 에이전트로 쿼리하기

Python SDK로 Databricks 다루기

Avi Steinberg

Senior Software Engineer

Databricks Unity 카탈로그

Databricks에서 데이터를 저장하는 기본 방법
여러 테이블을 포함할 수 있는 스키마로 구성
다양한 소스로부터 스키마에 테이블 생성 가능
samples 카탈로그에 다수의 오픈 소스 스키마 포함

Databricks Unity 카탈로그

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

LangChain

LLM 애플리케이션을 구축하고 배포할 수 있는 라이브러리
SQL을 사용하여 데이터를 쿼리하는 AI 에이전트 생성 가능

Python LangChain 종속성

Python 패키지 종속성 설치:

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

라이브러리 임포트:

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

Databricks SQL 웨어하우스 ID

Databricks SQL 웨어하우스 연결 정보

LangChain을 Databricks 워크스페이스에 인증하기

환경 변수 내보내기:

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

Databricks SQL 에이전트

Databricks 카탈로그 스키마의 데이터에 대해 AI 에이전트에 질문
AI 에이전트는 SQL이 필요한 데이터 질문에 응답 가능
verbose=True를 지정하면 AI 에이전트의 작업 과정 확인 가능

LangChain 쿼리 출력

Databricks SQL 에이전트로 카탈로그 쿼리하기

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)
agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)
# Query the Databricks SQL Agent
result = agent.run("What's the time and distance of the longest trip?")
print(result)

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

Databricks 카탈로그로 LangChain SQL 데이터베이스 생성하기

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

기반 모델로 LangChain LLM 생성하기

llm = ChatDatabricks(

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

    temperature=0.1,
    max_tokens=100,
)

temperature: 모델 응답의 무작위성을 지정하는 0~1 사이의 부동소수점 값
max_tokens: 응답의 최대 토큰 수 지정

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

LangChain SQL 데이터베이스 툴킷 생성하기

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

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

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

Databricks 데이터를 쿼리해 봅시다!

Python SDK로 Databricks 다루기