[FEATURE][UMBRELLA] Data Agent Engine — AI-Powered Autonomous Data Analysis

[wangzhigang1999]

Code of Conduct

• I agree to follow this project's Code of Conduct

Search before asking

• I have searched in the issues and found no similar issues.

Describe the feature

Umbrella issue tracking the implementation of DATA_AGENT engine type.

KPIP: #7373

The Data Agent enables users to perform data analysis through natural language — an AI agent autonomously explores schemas, generates SQL, executes queries via Kyuubi's existing multi-engine infrastructure, and self-corrects through multi-turn ReAct reasoning.

Architecture

Client (JDBC/REST/Web UI)
        │
        ▼
  Kyuubi Server (Gateway)  ◄───  JDBC (user creds)  ───┐
        │                                               │
        ▼                                               │
  Data Agent Engine                                     │
  ┌──────────────────────────┐                          │
  │  ReAct Loop              │                          │
  │  LLM ←→ Tools  ──────────┼──────────────────────────┘
  │       └─ sql_query (via Kyuubi JDBC)
  │  Middleware Pipeline     │
  │  ├─ ApprovalMiddleware   │
  │  └─ LoggingMiddleware    │
  └──────────────────────────┘

Sub-tasks

• PR 1: Module skeleton, configuration, and engine core — New module externals/kyuubi-data-agent-engine with engine fully runnable via Echo provider. Includes Thrift frontend, session/operation management, IncrementalFetchIterator for streaming, event system, and all kyuubi.engine.data.agent.* configuration entries.

• PR 2a: Tool system, data source, and prompt templates — SqlQueryTool with maxRows enforcement and output truncation, ToolRegistry with JSON schema generation for LLM function calling, data source abstraction with dialect auto-detection (Spark/SQLite/MySQL/Trino), and composable system prompt builder with per-dialect templates.

• PR 2b: Agent runtime, middleware, and OpenAI provider — ReAct loop agent with streaming LLM interaction, ConversationMemory for multi-turn context, middleware pipeline (ApprovalMiddleware with STRICT/NORMAL/AUTO_APPROVE modes, LoggingMiddleware), and OpenAI-compatible provider. Integration tests with MockLlmProvider validate the complete tool-call pipeline without a real LLM.

• PR 3: REST API and Web UI — SSE streaming chat endpoint (POST /api/v1/data-agent/{sessionHandle}/chat), tool approval endpoint (POST /api/v1/data-agent/{sessionHandle}/approve), and complete Vue web interface with session management, real-time message streaming, tool call visualization, and approval workflow UI.

Motivation

See KPIP-7373 for full motivation. In short: Kyuubi's existing Chat Engine is stateless with no data access. The Data Agent Engine bridges LLMs with Kyuubi's multi-engine SQL execution, enabling business users and analysts to query data warehouses through natural language without writing SQL.

Describe the solution

See KPIP-7373 for detailed design. Key decisions:

1. SQL routes through Kyuubi Server — The agent's sql_query tool connects back to Kyuubi Server via JDBC with the original user's credentials, inheriting AuthZ/Ranger policies, audit, and resource isolation.
2. Pluggable LLM providers — OpenAI-compatible API as default via official OpenAI Java SDK; extensible through DataAgentProvider interface.
3. Java for business logic, Scala for framework wrappers — Agent runtime, tools, providers, events, and middleware are all in Java; Scala is used only for thin integration with Kyuubi's Session/Operation/Engine infrastructure.
4. Streaming-first — IncrementalFetchIterator enables real-time event streaming to both JDBC and REST/SSE clients.
5. Human-in-the-loop approval — Configurable approval workflow (AUTO_APPROVE / NORMAL / STRICT) for controlling tool execution risk.

Additional context

Test strategy:

Layer	Approach	LLM Required?
Unit tests (Java)	JUnit 4 — tools, events, memory, middleware, prompts, data source	No
Integration tests (Scala)	MockLlmProvider drives full engine pipeline against SQLite	No
JDBC tests (Scala)	HiveJDBCTestHelper + Echo/Mock engine	No
Live tests (Java/Scala)	Real LLM API + SQLite test database	Yes (CI-optional)
Web UI tests (TypeScript)	Vitest — API client mocking	No

Are you willing to submit PR?

• Yes. I would be willing to submit a PR with guidance from the Kyuubi community to improve.

[wangzhigang1999]

Implementation progress update

The first PR is submitted: #7385 (module skeleton, configuration, and engine core).

The full implementation is split into 4 incremental PRs:

1. PR 1 (submitted): Module skeleton, engine core, Echo provider — [KYUUBI #7379][1/4] Data Agent Engine: module skeleton, configuration, and engine core #7385
2. PR 2a (ready): Tool system, data source, prompt templates
3. PR 2b (ready): Agent runtime, middleware, OpenAI provider
4. PR 3 (ready): REST API and Web UI

Here's a preview of the Web UI:

[wangzhigang1999]

Preliminary Benchmark Results

Note: This is a quick feasibility test to validate the approach, not a rigorous evaluation. The results below were obtained using an internal production data agent system. The same core idea (ReAct + tool-calling architecture, with data-agent-specific tools and mechanisms) has been contributed to Kyuubi as the Data Agent Engine, but these numbers were not benchmarked on the Kyuubi implementation itself. We plan to integrate these optimizations into Kyuubi incrementally.

We ran against the BIRD-SQL Mini-Dev set (500 examples), a lite version of the BIRD benchmark covering multiple professional domains.

Metric	Result
EX (Execution Accuracy, overall)	~70%
Soft-F1 (simple difficulty)	~83–84%

For reference, frontier LLM baselines on the BIRD full test set (not directly comparable due to different eval splits, but provides context):

Model	EX (test)
Claude Opus 4.6	70.15%
Qwen3-Coder-480B-A35B	68.14%
Claude 4.5 Sonnet	66.85%
GLM-4.7	62.94%
DeepSeek-R1	60.93%

And on the BIRD Mini-Dev leaderboard (SQLite, with oracle knowledge):

Model	Size	EX
TA + GPT-4	UNK	58.00%
GPT-4	UNK	47.80%
GPT-4-turbo	UNK	45.80%
Llama3-70b-instruct	70B	40.80%

Our ~70% EX is achieved with Qwen3.5/3.6-plus — a much smaller model compared to the frontier baselines above. Unlike vanilla LLM baselines that rely on single-pass prompting, our system uses a ReAct agent with data-agent-specific tools (schema exploration, SQL execution, result validation). These domain-specific optimizations enable a smaller model to match or exceed the accuracy of much larger frontier models. While these numbers come from an internal system rather than the Kyuubi engine, the architectural approach is identical and will be progressively integrated into the Kyuubi Data Agent Engine. This at least demonstrates that the agent-based approach is viable and effective.

[wangzhigang1999]

@pan3793 Could you reopen this issue? It's an umbrella tracking the full 4-PR Data Agent series — only PR 1/4 (#7385) has landed so far, with PRs 2a, 2b, and 3 still pending. I suspect the auto-close was triggered by the squash commit title referencing #7379. I'll switch the follow-up PRs to reference individual sub-task issues to avoid the same thing happening again. Thanks!