[KYUUBI #7379][4/4] Data Agent Engine: REST API, SSE streaming, and Web UI

docs/quick_start/index.rst

@@ -26,3 +26,4 @@ Quick Start
     quick_start
     quick_start_with_helm
     quick_start_with_jdbc
+    quick_start_with_data_agent

externals/kyuubi-data-agent-engine/pom.xml

@@ -125,13 +125,19 @@
             <version>${victools.jsonschema.version}</version>
         </dependency>
 
-        <!-- JDBC drivers are not bundled; users provide them at runtime via extra classpath.
-             SQLite and MySQL drivers are kept here in test scope only. -->
+        <!-- JDBC drivers. SQLite and PostgreSQL are bundled (Apache-compatible licenses, common
+             in AI demos). MySQL is GPL — kept in test scope only; production users must supply
+             the driver via kyuubi.engine.data.agent.extra.classpath. Other drivers (ClickHouse,
+             Snowflake, Doris, etc.) follow the same extra-classpath path. -->
         <dependency>
             <groupId>org.xerial</groupId>
             <artifactId>sqlite-jdbc</artifactId>
             <version>${sqlite.version}</version>
-            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>org.postgresql</groupId>
+            <artifactId>postgresql</artifactId>
         </dependency>
 
         <dependency>

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/datasource/JdbcDialect.java

@@ -83,6 +83,7 @@ static JdbcDialect fromUrl(String jdbcUrl) {
       return null;
     }
     switch (name) {
+      case "kyuubi":
       case "hive2":
       case "spark":
         return SparkDialect.INSTANCE;

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/provider/chatcompletion/ChatCompletionProvider.java

@@ -118,18 +118,46 @@ private static DataSource attachJdbcDataSource(
     if (jdbcUrl == null) {
       return null;
     }
-    LOG.info("Data Agent JDBC URL configured ({})", jdbcUrl.replaceAll("//.*@", "//<redacted>@"));
-
-    String sessionUser =
-        ConfUtils.optionalString(conf, KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY());
+    JdbcDialect dialect = JdbcDialect.fromUrl(jdbcUrl);
+    if (dialect == null) {
+      throw new IllegalArgumentException(
+          "Invalid " + KyuubiConf.ENGINE_DATA_AGENT_JDBC_URL().key());
+    }
+    LOG.info("Data Agent JDBC datasource configured ({})", dialect.datasourceName());
+
+    // For Kyuubi/HiveServer2 backends, pass the session user for proxy-user impersonation —
+    // but only when the URL does not already carry explicit credentials, because URL-supplied
+    // credentials should win (some users pass credentials via session variables).
+    String subprotocol = JdbcDialect.extractSubprotocol(jdbcUrl);
+    boolean isKyuubiBackend = "kyuubi".equals(subprotocol) || "hive2".equals(subprotocol);
+    String hikariUser = null;
+    if (isKyuubiBackend && !hasKyuubiUrlCredentials(jdbcUrl)) {
+      hikariUser = ConfUtils.optionalString(conf, KyuubiReservedKeys.KYUUBI_SESSION_USER_KEY());
+    }
 
-    DataSource ds = DataSourceFactory.create(jdbcUrl, sessionUser);
+    DataSource ds = DataSourceFactory.create(jdbcUrl, hikariUser);
     registry.register(new RunSelectQueryTool(ds, queryTimeoutSeconds));
     registry.register(new RunMutationQueryTool(ds, queryTimeoutSeconds));
-    promptBuilder.datasource(JdbcDialect.fromUrl(jdbcUrl).datasourceName());
+    promptBuilder.datasource(dialect.datasourceName());
     return ds;
   }
 
+  /**
+   * Return true when a Kyuubi/HiveServer2 JDBC URL carries an explicit user identity as the
+   * semicolon-delimited {@code ;user=...} session variable. Only semicolon parameters are checked
+   * because the Kyuubi/Hive JDBC driver reads them as auth session vars — query parameters ({@code
+   * ?key=val}) are parsed as Hive configuration, not authentication. A URL that only carries {@code
+   * ;password=} without a user is not considered credentialed, so the session user is still passed
+   * for proxy-user impersonation.
+   */
+  static boolean hasKyuubiUrlCredentials(String jdbcUrl) {
+    if (jdbcUrl == null) {
+      return false;
+    }
+    String lower = jdbcUrl.toLowerCase();
+    return lower.contains(";user=");
+  }
+
   @Override
   public void open(String sessionId, String user) {
     sessions.put(sessionId, new ConversationMemory());

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/provider/echo/EchoProvider.java

@@ -54,7 +54,7 @@ public void run(String sessionId, ProviderRunRequest request, Consumer<AgentEven
 
     onEvent.accept(new ContentComplete(reply));
     onEvent.accept(new StepEnd(1));
-    onEvent.accept(new AgentFinish(1, 0, 0, 0));
+    onEvent.accept(new AgentFinish(1, 0, 0, 0, 0));
   }
 
   @Override

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/runtime/AgentRunContext.java

@@ -30,9 +30,14 @@ public class AgentRunContext {
   private final String sessionId;
   private Consumer<AgentEvent> eventEmitter;
   private int iteration;
-  private long promptTokens;
-  private long completionTokens;
-  private long totalTokens;
+  // accumulatedXxx: summed across every LLM call in this run (billing).
+  // lastXxx: most recent LLM call only; UIs use lastPromptTokens as "current context size".
+  // total is forwarded from the provider verbatim (not prompt+completion) because reasoning /
+  // cached tokens don't show up in the split, and CompactionMiddleware needs the real total.
+  private long accumulatedPromptTokens;
+  private long accumulatedCompletionTokens;
+  private long lastPromptTokens;
+  private long lastCompletionTokens;
   private ApprovalMode approvalMode;
 
   public AgentRunContext(ConversationMemory memory, ApprovalMode approvalMode) {
@@ -67,27 +72,34 @@ public void setIteration(int iteration) {
     this.iteration = iteration;
   }
 
-  public long getPromptTokens() {
-    return promptTokens;
+  public long getAccumulatedPromptTokens() {
+    return accumulatedPromptTokens;
   }
 
-  public long getCompletionTokens() {
-    return completionTokens;
+  public long getAccumulatedCompletionTokens() {
+    return accumulatedCompletionTokens;
   }
 
-  public long getTotalTokens() {
-    return totalTokens;
+  public long getLastPromptTokens() {
+    return lastPromptTokens;
+  }
+
+  public long getLastCompletionTokens() {
+    return lastCompletionTokens;
   }
 
   /**
    * Record one LLM call's usage. Updates both the per-run counters on this context and the
    * session-level cumulative on the underlying {@link ConversationMemory}, so middlewares that need
    * a session-wide picture can read it directly from memory without keeping their own bookkeeping.
+   * The provider's {@code total} is forwarded as-is and may exceed {@code prompt + completion} when
+   * the provider counts cached or reasoning tokens separately.
    */
   public void addTokenUsage(long prompt, long completion, long total) {
-    this.promptTokens += prompt;
-    this.completionTokens += completion;
-    this.totalTokens += total;
+    this.accumulatedPromptTokens += prompt;
+    this.accumulatedCompletionTokens += completion;
+    this.lastPromptTokens = prompt;
+    this.lastCompletionTokens = completion;
     memory.addCumulativeTokens(prompt, completion, total);
   }
 

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/runtime/LlmStreamClient.java

@@ -18,6 +18,7 @@
 package org.apache.kyuubi.engine.dataagent.runtime;
 
 import com.openai.client.OpenAIClient;
+import com.openai.core.JsonValue;
 import com.openai.core.http.StreamResponse;
 import com.openai.models.chat.completions.ChatCompletionAssistantMessageParam;
 import com.openai.models.chat.completions.ChatCompletionChunk;
@@ -31,6 +32,7 @@
 import java.util.List;
 import java.util.Map;
 import org.apache.kyuubi.engine.dataagent.runtime.event.ContentDelta;
+import org.apache.kyuubi.engine.dataagent.runtime.event.ReasoningDelta;
 import org.apache.kyuubi.engine.dataagent.tool.ToolRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -90,10 +92,28 @@ private void consumeChunk(AgentRunContext ctx, ChatCompletionChunk chunk, Stream
                 acc.content.append(text);
                 ctx.emit(new ContentDelta(text));
               });
+      // qwen3 / DeepSeek-R1 stream chain-of-thought as a non-standard `reasoning_content`
+      // field on the delta. openai-java exposes it via _additionalProperties().
+      String reasoning = extractReasoningDelta(c.delta()._additionalProperties());
+      if (reasoning != null && !reasoning.isEmpty()) {
+        ctx.emit(new ReasoningDelta(reasoning));
+      }
       c.delta().toolCalls().ifPresent(acc::mergeToolCallDeltas);
     }
   }
 
+  private static String extractReasoningDelta(Map<String, JsonValue> extras) {
+    if (extras == null || extras.isEmpty()) return null;
+    JsonValue v = extras.get("reasoning_content");
+    if (v == null) v = extras.get("reasoning");
+    if (v == null) return null;
+    try {
+      return v.convert(String.class);
+    } catch (RuntimeException e) {
+      return null;
+    }
+  }
+
   /**
    * Mutable accumulator for a single streaming LLM turn. Tool call fields are keyed by the chunk's
    * {@code index} because provider SDKs may deliver a single logical call across multiple chunks

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/runtime/ReactAgent.java

@@ -194,9 +194,10 @@ private static void emitFinish(AgentRunContext ctx) {
     ctx.emit(
         new AgentFinish(
             ctx.getIteration(),
-            ctx.getPromptTokens(),
-            ctx.getCompletionTokens(),
-            ctx.getTotalTokens()));
+            ctx.getAccumulatedPromptTokens(),
+            ctx.getAccumulatedCompletionTokens(),
+            ctx.getLastPromptTokens(),
+            ctx.getLastCompletionTokens()));
   }
 
   /**

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/runtime/event/AgentFinish.java

@@ -17,47 +17,64 @@
 
 package org.apache.kyuubi.engine.dataagent.runtime.event;
 
-/** The agent has finished its analysis. */
+/**
+ * The agent has finished its analysis. Carries two distinct usage views: {@code accumulated*} sums
+ * every LLM call in this run (billing / total cost), {@code last*} reflects only the final call
+ * (current context size, last response size).
+ */
 public final class AgentFinish extends AgentEvent {
   private final int totalSteps;
-  private final long promptTokens;
-  private final long completionTokens;
-  private final long totalTokens;
+  private final long accumulatedPromptTokens;
+  private final long accumulatedCompletionTokens;
+  private final long lastPromptTokens;
+  private final long lastCompletionTokens;
 
-  public AgentFinish(int totalSteps, long promptTokens, long completionTokens, long totalTokens) {
+  public AgentFinish(
+      int totalSteps,
+      long accumulatedPromptTokens,
+      long accumulatedCompletionTokens,
+      long lastPromptTokens,
+      long lastCompletionTokens) {
     super(EventType.AGENT_FINISH);
     this.totalSteps = totalSteps;
-    this.promptTokens = promptTokens;
-    this.completionTokens = completionTokens;
-    this.totalTokens = totalTokens;
+    this.accumulatedPromptTokens = accumulatedPromptTokens;
+    this.accumulatedCompletionTokens = accumulatedCompletionTokens;
+    this.lastPromptTokens = lastPromptTokens;
+    this.lastCompletionTokens = lastCompletionTokens;
   }
 
   public int totalSteps() {
     return totalSteps;
   }
 
-  public long promptTokens() {
-    return promptTokens;
+  public long accumulatedPromptTokens() {
+    return accumulatedPromptTokens;
+  }
+
+  public long accumulatedCompletionTokens() {
+    return accumulatedCompletionTokens;
   }
 
-  public long completionTokens() {
-    return completionTokens;
+  public long lastPromptTokens() {
+    return lastPromptTokens;
   }
 
-  public long totalTokens() {
-    return totalTokens;
+  public long lastCompletionTokens() {
+    return lastCompletionTokens;
   }
 
   @Override
   public String toString() {
     return "AgentFinish{totalSteps="
         + totalSteps
-        + ", promptTokens="
-        + promptTokens
-        + ", completionTokens="
-        + completionTokens
-        + ", totalTokens="
-        + totalTokens
+        + ", accumulatedPromptTokens="
+        + accumulatedPromptTokens
+        + ", accumulatedCompletionTokens="
+        + accumulatedCompletionTokens
+        + ", lastPromptTokens="
+        + lastPromptTokens
+        + ", lastCompletionTokens="
+        + lastCompletionTokens
         + "}";
   }
 }

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/runtime/event/EventType.java

@@ -35,6 +35,9 @@ public enum EventType {
   /** A single token or chunk from the LLM streaming response. */
   CONTENT_DELTA("content_delta"),
 
+  /** A chunk of the LLM's reasoning / chain-of-thought stream (provider-specific). */
+  REASONING_DELTA("reasoning_delta"),
+
   /** The complete LLM output for one reasoning step. */
   CONTENT_COMPLETE("content_complete"),
 

externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/runtime/event/ReasoningDelta.java

@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.engine.dataagent.runtime.event;
+
+/**
+ * A chunk of the LLM's hidden reasoning / chain-of-thought stream. Emitted alongside {@link
+ * ContentDelta} when the provider exposes a {@code reasoning_content} delta field (e.g. qwen3,
+ * DeepSeek-R1). UI clients typically render these in a collapsible "thinking" panel.
+ */
+public final class ReasoningDelta extends AgentEvent {
+  private final String text;
+
+  public ReasoningDelta(String text) {
+    super(EventType.REASONING_DELTA);
+    this.text = text;
+  }
+
+  public String text() {
+    return text;
+  }
+
+  @Override
+  public String toString() {
+    String preview = text != null && text.length() > 200 ? text.substring(0, 200) + "..." : text;
+    return "ReasoningDelta{text='" + preview + "'}";
+  }
+}