Skip to content

[KYUUBI #7379][2b/4] Data Agent Engine: agent runtime, middleware stack, and OpenAI provider #7417

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wangzhigang1999 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[KYUUBI #7379][2b/4] Data Agent Engine: agent runtime, middleware stack, and OpenAI provider #7417

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ce4eecc
[KYUUBI #7379][2b/4] Data Agent Engine: agent runtime, middleware sta…
wangzhigang1999 Apr 22, 2026
5223561
[KYUUBI #7379][2b/4] Move mysql-connector-j to test scope
wangzhigang1999 Apr 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions docs/configuration/settings.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co
| kyuubi.engine.chat.provider | ECHO | The provider for the Chat engine. Candidates: <ul> <li>ECHO: simply replies a welcome message.</li> <li>GPT: a.k.a ChatGPT, powered by OpenAI.</li> <li>ERNIE: ErnieBot, powered by Baidu.</li></ul> | string | 1.8.0 |
| kyuubi.engine.connection.url.use.hostname | true | (deprecated) When true, the engine registers with hostname to zookeeper. When Spark runs on K8s with cluster mode, set to false to ensure that server can connect to engine | boolean | 1.3.0 |
| kyuubi.engine.data.agent.approval.mode | NORMAL | Default approval mode for tool execution in the Data Agent engine. Candidates: <ul> <li>AUTO_APPROVE: all tools are auto-approved without user interaction.</li> <li>NORMAL: only destructive tools require explicit approval.</li> <li>STRICT: all tools require explicit user approval.</li></ul> | string | 1.12.0 |
| kyuubi.engine.data.agent.compaction.trigger.tokens | 128000 | The prompt-token threshold above which the Data Agent's compaction middleware summarizes older conversation history into a compact message. The check is made each turn as <code>real_prompt_tokens_of_previous_LLM_call + estimate_of_newly_appended_tail</code>; when this predicted prompt size reaches the configured value, older messages are replaced by a single summary message while the most recent exchanges are kept verbatim. Set to a very large value (e.g., <code>9223372036854775807</code>) to effectively disable compaction. | long | 1.12.0 |
| kyuubi.engine.data.agent.extra.classpath | &lt;undefined&gt; | The extra classpath for the Data Agent engine, for configuring the location of the LLM SDK and etc. | string | 1.12.0 |
| kyuubi.engine.data.agent.java.options | &lt;undefined&gt; | The extra Java options for the Data Agent engine | string | 1.12.0 |
| kyuubi.engine.data.agent.jdbc.url | &lt;undefined&gt; | The JDBC URL for the Data Agent engine to connect to the target database. If not set, the Data Agent will connect back to Kyuubi server via ZooKeeper service discovery. | string | 1.12.0 |
Expand Down
44 changes: 31 additions & 13 deletions externals/kyuubi-data-agent-engine/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,45 +50,63 @@
<version>${project.version}</version>
</dependency>

<!-- OpenAI official Java SDK -->
<dependency>
<groupId>com.openai</groupId>
<artifactId>openai-java</artifactId>
<version>${openai.sdk.version}</version>
</dependency>

<!-- JSON Schema generation from Jackson-annotated classes -->
<dependency>
<groupId>com.github.victools</groupId>
<artifactId>jsonschema-generator</artifactId>
<version>${victools.jsonschema.version}</version>
</dependency>

<dependency>
<groupId>com.github.victools</groupId>
<artifactId>jsonschema-module-jackson</artifactId>
<version>${victools.jsonschema.version}</version>
</dependency>

<!-- test dependencies -->
<!-- SQLite JDBC driver -->
<dependency>
<groupId>org.apache.kyuubi</groupId>
<artifactId>kyuubi-common_${scala.binary.version}</artifactId>
<version>${project.version}</version>
<type>test-jar</type>
<scope>test</scope>
<groupId>org.xerial</groupId>
<artifactId>sqlite-jdbc</artifactId>
<version>${sqlite.version}</version>
</dependency>

<!-- MySQL JDBC driver (test scope: GPL-licensed, cannot be bundled in Apache binary release) -->
<dependency>
<groupId>org.xerial</groupId>
<artifactId>sqlite-jdbc</artifactId>
<groupId>com.mysql</groupId>
<artifactId>mysql-connector-j</artifactId>
Comment thread
wangzhigang1999 marked this conversation as resolved.
<scope>test</scope>
</dependency>

<!-- Trino JDBC driver -->
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers-mysql</artifactId>
<groupId>io.trino</groupId>
<artifactId>trino-jdbc</artifactId>
</dependency>

<!-- Connection pool -->
<dependency>
<groupId>com.zaxxer</groupId>
<artifactId>HikariCP</artifactId>
</dependency>

<!-- test dependencies -->
<dependency>
<groupId>org.apache.kyuubi</groupId>
<artifactId>kyuubi-common_${scala.binary.version}</artifactId>
<version>${project.version}</version>
<type>test-jar</type>
<scope>test</scope>
</dependency>

<dependency>
<groupId>com.mysql</groupId>
<artifactId>mysql-connector-j</artifactId>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers-mysql</artifactId>
<scope>test</scope>
</dependency>

Expand Down
6 changes: 6 additions & 0 deletions ...agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/datasource/JdbcDialect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,12 @@

package org.apache.kyuubi.engine.dataagent.datasource;

import org.apache.kyuubi.engine.dataagent.datasource.dialect.GenericDialect;
import org.apache.kyuubi.engine.dataagent.datasource.dialect.MysqlDialect;
import org.apache.kyuubi.engine.dataagent.datasource.dialect.SparkDialect;
import org.apache.kyuubi.engine.dataagent.datasource.dialect.SqliteDialect;
import org.apache.kyuubi.engine.dataagent.datasource.dialect.TrinoDialect;

/**
* SQL dialect abstraction for datasource-specific SQL generation.
*
Expand Down
4 changes: 3 additions & 1 deletion .../dataagent/datasource/GenericDialect.java → ...nt/datasource/dialect/GenericDialect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,9 @@
* limitations under the License.
*/

package org.apache.kyuubi.engine.dataagent.datasource;
package org.apache.kyuubi.engine.dataagent.datasource.dialect;

import org.apache.kyuubi.engine.dataagent.datasource.JdbcDialect;

/**
* Fallback dialect for JDBC subprotocols that have no dedicated implementation. Carries the
Expand Down
6 changes: 4 additions & 2 deletions ...ne/dataagent/datasource/MysqlDialect.java → ...gent/datasource/dialect/MysqlDialect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,14 @@
* limitations under the License.
*/

package org.apache.kyuubi.engine.dataagent.datasource;
package org.apache.kyuubi.engine.dataagent.datasource.dialect;

import org.apache.kyuubi.engine.dataagent.datasource.JdbcDialect;

/** MySQL dialect. Uses backtick quoting for identifiers. */
public final class MysqlDialect implements JdbcDialect {

static final MysqlDialect INSTANCE = new MysqlDialect();
public static final MysqlDialect INSTANCE = new MysqlDialect();

private MysqlDialect() {}

Expand Down
6 changes: 4 additions & 2 deletions ...ne/dataagent/datasource/SparkDialect.java → ...gent/datasource/dialect/SparkDialect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,14 @@
* limitations under the License.
*/

package org.apache.kyuubi.engine.dataagent.datasource;
package org.apache.kyuubi.engine.dataagent.datasource.dialect;

import org.apache.kyuubi.engine.dataagent.datasource.JdbcDialect;

/** Spark SQL dialect. Uses backtick quoting for identifiers. */
public final class SparkDialect implements JdbcDialect {

static final SparkDialect INSTANCE = new SparkDialect();
public static final SparkDialect INSTANCE = new SparkDialect();

private SparkDialect() {}

Expand Down
6 changes: 4 additions & 2 deletions ...e/dataagent/datasource/SqliteDialect.java → ...ent/datasource/dialect/SqliteDialect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,14 @@
* limitations under the License.
*/

package org.apache.kyuubi.engine.dataagent.datasource;
package org.apache.kyuubi.engine.dataagent.datasource.dialect;

import org.apache.kyuubi.engine.dataagent.datasource.JdbcDialect;

/** SQLite dialect. Uses double-quote quoting for identifiers. */
public final class SqliteDialect implements JdbcDialect {

static final SqliteDialect INSTANCE = new SqliteDialect();
public static final SqliteDialect INSTANCE = new SqliteDialect();

private SqliteDialect() {}

Expand Down
6 changes: 4 additions & 2 deletions ...ne/dataagent/datasource/TrinoDialect.java → ...gent/datasource/dialect/TrinoDialect.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,14 @@
* limitations under the License.
*/

package org.apache.kyuubi.engine.dataagent.datasource;
package org.apache.kyuubi.engine.dataagent.datasource.dialect;

import org.apache.kyuubi.engine.dataagent.datasource.JdbcDialect;

/** Trino SQL dialect. Uses double-quote quoting for identifiers. */
public final class TrinoDialect implements JdbcDialect {

static final TrinoDialect INSTANCE = new TrinoDialect();
public static final TrinoDialect INSTANCE = new TrinoDialect();

private TrinoDialect() {}

Expand Down
26 changes: 24 additions & 2 deletions ...-engine/src/main/java/org/apache/kyuubi/engine/dataagent/provider/ProviderRunRequest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,23 @@

package org.apache.kyuubi.engine.dataagent.provider;

import org.apache.kyuubi.engine.dataagent.runtime.ApprovalMode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
* User-facing request parameters for a provider-level agent invocation. Only contains fields from
* the caller (question, model override, etc.). Adding new per-request options does not require
* changing the {@link DataAgentProvider} interface.
*
* <p>The approval mode is accepted as a raw string (natural for config-driven callers) and parsed
* into {@link ApprovalMode} by {@link #getApprovalMode()}. Unrecognised values fall back to {@link
* ApprovalMode#NORMAL} with a warning.
*/
public class ProviderRunRequest {

private static final Logger LOG = LoggerFactory.getLogger(ProviderRunRequest.class);

private final String question;
private String modelName;
private String approvalMode;
Expand All @@ -45,8 +55,20 @@ public ProviderRunRequest modelName(String modelName) {
return this;
}

public String getApprovalMode() {
return approvalMode;
/**
* Resolved approval mode. Returns {@link ApprovalMode#NORMAL} when the caller did not set one or
* supplied an unknown value.
*/
public ApprovalMode getApprovalMode() {
if (approvalMode == null || approvalMode.isEmpty()) {
return ApprovalMode.NORMAL;
}
try {
return ApprovalMode.valueOf(approvalMode.toUpperCase());
} catch (IllegalArgumentException e) {
LOG.warn("Unknown approval mode '{}', using default NORMAL", approvalMode);
return ApprovalMode.NORMAL;
}
}

public ProviderRunRequest approvalMode(String approvalMode) {
Expand Down
Loading
Loading