CrewAI DatabricksQueryTool 详解让 AI Agent 用 SQL 查询 Databricks 工作区表【免费下载链接】crewAIFramework for orchestrating role-playing, autonomous AI agents. By fostering collaborative intelligence, CrewAI empowers agents to work together seamlessly, tackling complex tasks.项目地址: https://gitcode.com/GitHub_Trending/cr/crewAI本文以 CrewAI 工具库crewai-tools中的DatabricksQueryTool为对象系统讲解该工具如何让 AI Agent 通过 SQL 查询 Databricks 工作区表并拿到结果集包括两种认证方式的配置、实例化时的默认参数default_catalog/default_schema/default_warehouse_id、运行时五个输入参数含row_limit自动注入 LIMIT 的行为并结合 源码实现 剖析从语句提交、状态轮询到结果格式化的完整执行链路读完即可在自己的 Crew 中安全落地该工具。工具定位DatabricksQueryTool位于 databricks_query_tool 模块继承自crewai.tools.BaseTool其设计目标是让 Agent 能用 SQL 访问 Databricks 中存储的数据Agent 只需给出 SQL 语句以及可选的 catalog / schema / warehouse 参数工具就会通过databricks-sdk的 Statement Execution API 提交查询、等待执行完成、把结果集整理成对齐的文本表格返回给 Agent。该工具已通过 crewai_tools 包导出crewai_tools顶层与crewai_tools.tools子包均导出了DatabricksQueryTool因此可以直接from crewai_tools import DatabricksQueryTool使用。源码中工具自身的注册信息见 databricks_query_tool.pyname Databricks SQL Querydescription为 Execute SQL queries against Databricks workspace tables and return the results. Provide a query parameter with the SQL query to execute.——这段描述会直接作为提示词上下文交给 LLM 决策是否调用该工具args_schema DatabricksQueryToolSchema由 Pydantic 模型定义合法入参package_dependencies [databricks-sdk]用于在依赖缺失时给出安装提示。安装README 给出的 pip 安装方式为pip install crewai_tools databricks-sdk在 pyproject.toml 中crewai-tools提供了名为databricks-sdk的可选依赖组要求databricks-sdk0.46.0因此也可以按 官方工具文档 推荐的方式一次性安装uv add crewai-tools[databricks-sdk]两种方式等价区别只是是否把版本约束交给 crewai-tools 管理。注意databricks-sdk是运行时按需导入的见下文workspace_client属性如果环境中缺少该包工具在真正发起查询前就会抛出带安装提示的ImportErrordatabricks-sdk package not found, please run uv add databricks-sdk。认证两种提供凭据的方式工具不接收host/token等认证参数而是完全依赖环境变量。源码init在构造时即调用_validate_credentials()做前置校验L143-L154使用 Databricks CLI profile设置DATABRICKS_CONFIG_PROFILE环境变量指向你的 profile 名称适用于~/.databrickscfg中已配置 profile 的场景直接使用凭据同时设置DATABRICKS_HOST和DATABRICKS_TOKEN。export DATABRICKS_HOSThttps://your-workspace.cloud.databricks.com export DATABRICKS_TOKENdapi1234567890abcdef校验逻辑是二选一只要环境变量中存在DATABRICKS_CONFIG_PROFILE或者DATABRICKS_HOST与DATABRICKS_TOKEN同时存在即通过否则在工具实例化阶段就抛出ValueError而不是等到查询时才失败。源码中env_vars字段L102-L120也以结构化方式声明了这三个变量及用途描述供框架侧展示。个人访问令牌PAT与 workspace 地址的获取方式在 Databricks workspace 的 User Settings → Developer 中创建 PAT 并查看 host 信息。实例化三个默认参数from crewai_tools import DatabricksQueryTool # 基本用法不指定默认值 databricks_tool DatabricksQueryTool() # 指定 catalog、schema、warehouse 默认值 databricks_tool DatabricksQueryTool( default_catalogmy_catalog, default_schemamy_schema, default_warehouse_idwarehouse_id )三个构造参数L95-L98全部可选构造参数类型说明default_catalogstr \| None默认 catalog如maindefault_schemastr \| None默认 schema如defaultdefault_warehouse_idstr \| None默认 SQL warehouse ID它们的语义是运行时的兜底值_run()中按kwargs.get(catalog) or self.default_catalog的顺序取值L226-L230即每次调用显式传入的参数优先于实例化时的默认值。需要特别注意warehouse_id虽然可选但必须来自参数或默认值二者之一否则_run()会直接返回错误信息SQL warehouse ID must be provided either as a parameter or as a default.L245-L246。这是因为 Statement Execution API 必须明确指定在哪个 SQL warehouse 上执行语句Databricks 侧没有可推断的默认值。而 catalog / schema 可以缺失——缺失时 Databricks 会按会话自身的默认上下文解析表名。输入参数与 row_limit 的自动 LIMIT 注入工具对 LLM 暴露的入参由 Pydantic 模型DatabricksQueryToolSchemaL37-L70定义参数必填默认值说明query是—要执行的 SQL空字符串或纯空白会被校验器拒绝Query cannot be emptycatalog否实例default_catalogDatabricks catalog 名db_schema否实例default_schemaDatabricks schema 名warehouse_id否但必须可解析实例default_warehouse_idSQL warehouse IDrow_limit否1000返回的最大行数这里有一个容易踩坑的命名差异模块 README 将 schema 参数写作schema但源码中实际字段名是db_schemaschema是 Pydantic 模型的保留属性名不能直接用。官方工具文档 的 Parameters 一节同样写的是db_schema因此 Agent 调用时应使用db_schema。row_limit的行为值得单独说明。校验器中的validate_inputL59-L70在参数进入执行逻辑之前做了一次查询改写if self.row_limit and limit not in self.query.lower(): self.query f{self.query.rstrip(;)} LIMIT {self.row_limit};即当 SQL 中不包含大小写不敏感的limit字样时工具会自动去掉尾部分号并追加LIMIT {row_limit}默认 1000。这意味着未写 LIMIT 的SELECT *类查询天然被限制在 1000 行以内避免 Agent 一次拉回超大结果集如果你的 SQL 已经自带LIMIT则保持原样、不会被覆盖判断依据是字符串匹配limit因此查询文本中出现含该子串的其他内容也会跳过自动注入——对常规 DML/DDL 语句无影响。执行流程提交、轮询与状态机_run()L209 起的完整执行链路可以从源码结构看分为四段1. 惰性创建 WorkspaceClient。workspace_client属性L156-L168在首次访问时才from databricks.sdk import WorkspaceClient并实例化后续复用同一实例。WorkspaceClient()不传任何参数正是靠上文的三个环境变量完成认证。2. 提交语句。认证信息确认、参数校验通过后工具构造一个ExecutionContextTypedDict只含catalog/schema两个键调用execution self.workspace_client.statement_execution.execute_statement( warehouse_idwarehouse_id, statementquery, **context ) statement_id execution.statement_id提交失败时返回Error starting query execution: ...拿不到statement_id时返回Failed to retrieve statement ID after execution.。3. 轮询等待终态。Databricks 语句执行是异步的工具以固定节奏轮询statement.get_statement(statement_id)超时窗口timeout 300秒5 分钟每次间隔time.sleep(2)轮询一次状态判定基于result.status.state兼容字符串与枚举两种形态包含SUCCEEDED则跳出轮询包含FAILED则提取错误信息并返回Query execution failed: {error_info}包含CANCELED则返回Query was canceled错误信息提取做了双重兼容优先取status.error.message其次status.error.error_message再不行退化为str(error)轮询请求本身抛异常时不会立即失败连续 3 次以上才返回Error checking query status: ...容忍瞬时网络抖动若 5 分钟内未达终态返回Query timed out after 5 minutes (last state: ...)并附上最后观察到的状态。4. 结果解析。拿到SUCCEEDED的 statement 后工具从result.manifest.schema.columns取列名从result.result.data_array按 chunk 迭代行数据。从源码结构看这一段做了非常防御式的解析兼容data_array与data两种载荷、处理 chunk 内行结构异常的情况例如值被逐字符拆开时按启发式规则重建行边界、把超出列名的多余值归入动态列Column_{i}最后统一归一化——保证每行都含有全部列、缺失值补None。对 DDL 等无结果集的成功语句则直接返回Query executed successfully (no results to display)。结果格式化Agent 友好的文本表格解析出的行数据最终由_format_results()L170-L207渲染为固定宽度对齐的文本表格按每列数据的最大显示宽度计算列宽首行表头 --分隔线 数据行列之间以|分隔None值统一显示为NULL末尾附(N rows returned)行计数空结果的三种情况分别有明确文案无行、有行但无列、有列但无数据。这种纯文本表格对 LLM 是最友好的返回形态——结构清晰、token 开销可控、无需 Agent 再解析 JSON。所有异常路径也不会把原始堆栈抛给 Agent而是返回带traceback.format_exc()详情的错误文本如Error executing Databricks query: ...既保留了调试信息又保证工具调用本身不中断 Crew 流程。在 CrewAI Agent 中使用README 的最小集成方式是给 Agent 挂载工具from crewai_tools import DatabricksQueryTool databricks_tool DatabricksQueryTool( default_catalogmy_catalog, default_schemamy_schema, default_warehouse_idwarehouse_id ) # 在 Agent 定义中挂载 # Agent(config..., allow_delegationFalse, tools[databricks_tool])官方工具文档 给出了一个完整的 Crew 示例可直接复制运行前提是环境变量已按认证一节配置from crewai import Agent, Task, Crew from crewai_tools import DatabricksQueryTool tool DatabricksQueryTool( default_catalogmain, default_schemadefault, ) agent Agent( roleData Analyst, goalQuery Databricks, tools[tool], verboseTrue, ) task Task( descriptionSELECT * FROM my_table LIMIT 10, expected_output10 rows, agentagent, ) crew Crew( agents[agent], tasks[task], verboseTrue, ) result crew.kickoff() print(result)也可以不经过 Agent直接调用工具本身验证连通性源码 docstring 中的示例 tool DatabricksQueryTool(default_warehouse_idyour_warehouse_id) results tool.run(querySELECT * FROM my_table LIMIT 10)错误处理与实战建议结合 官方文档的 Error handling tips 一节 与源码行为整理如下认证错误确认DATABRICKS_HOST以https://开头且 token 有效使用 CLI profile 时确认DATABRICKS_CONFIG_PROFILE指向的 profile 存在。权限问题确保 token 对目标 SQL warehouse 与 schema 有访问权限warehouse_id必须与 token 所在工作区匹配。长查询限制工具端硬编码了 5 分钟轮询超时超时即返回 Query timed out。在 Agent 循环中应避免发起长耗时查询建议在 SQL 中主动加过滤条件和LIMITrow_limit的自动注入只是保底不是性能手段。db_schema而非schema给 Agent 的任务描述或自定义提示中如涉及该参数建议使用实际字段名db_schema。错误信息可读性所有失败路径都返回自然语言错误串SQL 执行失败、状态检查失败、结果处理失败均附详情Agent 通常能据此自行修正 SQL 重试如果你需要程序化判断失败可对这些返回串做前缀匹配如Query execution failed、Error executing Databricks query。关键文件索引内容路径工具模块 README本篇主体文档lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/README.md工具完整实现lib/crewai-tools/src/crewai_tools/tools/databricks_query_tool/databricks_query_tool.py官方工具文档Crew 示例与排错建议docs/edge/en/tools/search-research/databricks-query-tool.mdxdatabricks-sdk可选依赖声明0.46.0lib/crewai-tools/pyproject.toml包级导出lib/crewai-tools/src/crewai_tools/tools/init.py【免费下载链接】crewAIFramework for orchestrating role-playing, autonomous AI agents. By fostering collaborative intelligence, CrewAI empowers agents to work together seamlessly, tackling complex tasks.项目地址: https://gitcode.com/GitHub_Trending/cr/crewAI创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考