RDBC-959: sync with 7.1.3 -> 7.1.4 diff related to AI

Author: LielNagar

src/main/java/net/ravendb/client/documents/AI/AiAnswer.java

@@ -1,15 +1,46 @@
 package net.ravendb.client.documents.AI;
 
+import net.ravendb.client.documents.operations.AI.AiUsage;
+import java.time.Duration;
+
 public class AiAnswer<TAnswer> {
+    /**
+     * The answer content produced by the AI.
+     */
     private TAnswer answer;
+
+    /**
+     * The status of the conversation.
+     */
     private AiConversationResult status;
 
+    /**
+     * Token usage reported by the model for generating this answer (prompt/completion/total).
+     */
+    private AiUsage usage;
+
+    /**
+     * The total time elapsed to produce the answer(measured from the server's request to the LLM until the response was received).
+     */
+    private Duration elapsed;
+
     public AiAnswer() {
     }
 
-    public AiAnswer(TAnswer answer, AiConversationResult status) {
-        this.answer = answer;
-        this.status = status;
+    public AiUsage getUsage() {
+        return usage;
+    }
+
+    public void setUsage(AiUsage usage) {
+        this.usage = usage;
+    }
+
+    public Duration getElapsed() {
+        return elapsed;
+    }
+
+    public void setElapsed(Duration elapsed) {
+        this.elapsed = elapsed;
     }
 
     public TAnswer getAnswer() {

src/main/java/net/ravendb/client/documents/AI/AiConversation.java

@@ -4,6 +4,7 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import net.ravendb.client.documents.IDocumentStore;
 import net.ravendb.client.documents.operations.AI.agents.*;
+import net.ravendb.client.exceptions.ConcurrencyException;
 import net.ravendb.client.extensions.expressionExtension;
 import net.ravendb.client.util.SerializableFunction;
 import org.apache.commons.lang3.StringUtils;
@@ -298,40 +299,47 @@ public <TAnswer> CompletableFuture<AiAnswer<TAnswer>> run() {
     }
 
     private <TAnswer> CompletableFuture<AiAnswer<TAnswer>> runInternal(String streamPropertyPath, AiStreamCallback streamCallback) {
-        if (this.actionRequests != null && this.promptParts.isEmpty() && this.actionResponses.isEmpty()) {
-            AiAnswer<TAnswer> doneAnswer = new AiAnswer<>();
-            doneAnswer.setStatus(AiConversationResult.Done);
-            return CompletableFuture.completedFuture(doneAnswer);
-        }
-
-        RunConversationOperation<TAnswer> op = new RunConversationOperation<>(
-                this.agentId,
-                this.conversationId,
-                this.promptParts,
-                this.actionResponses,
-                this.options,
-                this.changeVector,
-                streamPropertyPath,
-                streamCallback
-        );
-
-        ConversationResult<TAnswer> result = this.store.maintenance()
-                .forDatabase(this.databaseName)
-                .send(op);
-
-        AiAnswer<TAnswer> answer = new AiAnswer<>();
-        answer.setAnswer(result.getResponse());
-        answer.setStatus(result.getActionRequests() == null || result.getActionRequests().isEmpty()
-                ? AiConversationResult.Done
-                : AiConversationResult.ActionRequired);
-
-        this.changeVector = result.getChangeVector();
-        this.conversationId = result.getConversationId();
-        this.actionRequests = result.getActionRequests() != null ? result.getActionRequests() : new ArrayList<>();
-        this.promptParts.clear();
-        this.actionResponses.clear();
+        try {
+            if (this.actionRequests != null && this.promptParts.isEmpty() && this.actionResponses.isEmpty()) {
+                AiAnswer<TAnswer> doneAnswer = new AiAnswer<>();
+                doneAnswer.setStatus(AiConversationResult.Done);
+                return CompletableFuture.completedFuture(doneAnswer);
+            }
 
-        return CompletableFuture.completedFuture(answer);
+            RunConversationOperation<TAnswer> op = new RunConversationOperation<>(
+                    this.agentId,
+                    this.conversationId,
+                    this.promptParts,
+                    this.actionResponses,
+                    this.options,
+                    this.changeVector,
+                    streamPropertyPath,
+                    streamCallback
+            );
+
+            ConversationResult<TAnswer> result = this.store.maintenance()
+                    .forDatabase(this.databaseName)
+                    .send(op);
+            this.changeVector = result.getChangeVector();
+            this.conversationId = result.getConversationId();
+            this.actionRequests = result.getActionRequests() != null ? result.getActionRequests() : new ArrayList<>();
+
+            AiAnswer<TAnswer> answer = new AiAnswer<>();
+            answer.setAnswer(result.getResponse());
+            answer.setStatus(result.getActionRequests() == null || result.getActionRequests().isEmpty()
+                    ? AiConversationResult.Done
+                    : AiConversationResult.ActionRequired);
+            answer.setUsage(result.getUsage());
+            answer.setElapsed(result.getElapsed());
+            return CompletableFuture.completedFuture(answer);
+        } catch (ConcurrencyException e) {
+            this.changeVector = e.getActualChangeVector();
+            throw e;
+        }
+        finally {
+            this.promptParts.clear();
+            this.actionResponses.clear();
+        }
     }
 
     private Object parseArgs(String argsJson) {

src/main/java/net/ravendb/client/documents/changes/AbstractDatabaseChanges.java

@@ -59,7 +59,7 @@ public abstract class AbstractDatabaseChanges<TDatabaseConnectionState extends A
     private final CancellationTokenSource _cts;
     private CompletableFuture<AbstractDatabaseChanges<TDatabaseConnectionState>> _tcs;
 
-    protected final ConcurrentMap<Integer, CompletableFuture<Void>> _confirmations = new ConcurrentHashMap<>();
+    protected final ConcurrentMap<Integer, CompletableFuture> _confirmations = new ConcurrentHashMap<>();
 
     protected final ConcurrentMap<DatabaseChangesOptions, DatabaseConnectionState> _states = new ConcurrentHashMap<>();
 

src/main/java/net/ravendb/client/documents/changes/AbstractDatabaseConnectionState.java

@@ -1,9 +1,9 @@
 package net.ravendb.client.documents.changes;
 
 import net.ravendb.client.primitives.EventHelper;
-
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Consumer;
 
@@ -14,9 +14,24 @@ public class AbstractDatabaseConnectionState {
     private final Runnable _onDisconnect;
     public final Runnable onConnect;
 
-    private final AtomicInteger _value = new AtomicInteger();
+    private final AtomicInteger _value = new AtomicInteger(0);
     public Exception lastException;
 
+    private final CompletableFuture<Void> firstSet = new CompletableFuture<>();
+    private CompletableFuture<Void> connected;
+
+    public void set(CompletableFuture<Void> connection) {
+        if (!firstSet.isDone()) {
+            connection.whenComplete((res, ex) -> {
+                if (ex != null) {
+                    firstSet.completeExceptionally(ex);
+                } else {
+                    firstSet.complete(null);
+                }
+            });
+        }
+        connected = connection;
+    }
 
     public void addOnError(Consumer<Exception> handler) {
         this.onError.add(handler);
@@ -32,7 +47,6 @@ protected AbstractDatabaseConnectionState(Runnable onConnect, Runnable onDisconn
         _value.set(0);
     }
 
-
     public void inc() {
         _value.incrementAndGet();
     }

src/main/java/net/ravendb/client/documents/operations/AI/AiUsage.java

@@ -1,5 +1,7 @@
 package net.ravendb.client.documents.operations.AI;
 
+import com.fasterxml.jackson.databind.JsonNode;
+
 public class AiUsage {
     private int promptTokens;
     private int completionTokens;
@@ -53,5 +55,40 @@ public int getCachedTokens() {
     public void setCachedTokens(int cachedTokens) {
         this.cachedTokens = cachedTokens;
     }
+
+    void updateFrom(JsonNode json) {
+        if (json.has("prompt_tokens")) {
+            this.promptTokens += json.get("prompt_tokens").asInt();
+        }
+        if (json.has("completion_tokens")) {
+            this.completionTokens += json.get("completion_tokens").asInt();
+        }
+        if (json.has("total_tokens")) {
+            this.totalTokens += json.get("total_tokens").asInt();
+        }
+
+        JsonNode promptDetails = json.get("prompt_tokens_details");
+        if (promptDetails != null && promptDetails.has("cached_tokens")) {
+            this.cachedTokens += promptDetails.get("cached_tokens").asInt();
+        }
+
+        JsonNode completionTokensDetails = json.get("completion_tokens_details");
+        if (completionTokensDetails != null && completionTokensDetails.has("reasoning_tokens")) {
+            this.reasoningTokens += completionTokensDetails.get("reasoning_tokens").asInt();
+        }
+    }
+
+    static AiUsage getUsageDifference(AiUsage current, AiUsage previous) {
+        int previousTotalWithoutReasoning = (previous.getCompletionTokens() - previous.getReasoningTokens() + previous.getPromptTokens());
+
+        AiUsage result = new AiUsage();
+        result.setPromptTokens(Math.max(current.getPromptTokens() - previousTotalWithoutReasoning, 0));
+        result.setTotalTokens(Math.max(current.getTotalTokens() - previousTotalWithoutReasoning, 0));
+        result.setCachedTokens(current.getCachedTokens());        // keep as-is
+        result.setCompletionTokens(current.getCompletionTokens()); // keep as-is
+        result.setReasoningTokens(current.getReasoningTokens());   // keep as-is
+
+        return result;
+    }
 }
 

src/main/java/net/ravendb/client/documents/operations/AI/agents/AiAgentParameter.java

@@ -1,5 +1,8 @@
 package net.ravendb.client.documents.operations.AI.agents;
 
+import java.util.HashMap;
+import java.util.Map;
+
 public class AiAgentParameter {
     private String name;
     private String description;
@@ -20,6 +23,19 @@ public AiAgentParameter(String name, String description) {
         this.description = description;
     }
 
+    /**
+     * Initializes a new agent parameter and controls whether its value should be sent to the LLM.
+     * Use this constructor when you need to explicitly hide sensitive values
+     * (e.g., userId, tenant, or company) from the model.
+     *
+     * @param name the parameter name; cannot be null or empty.
+     * @param description a human-readable description; may be null or empty when using this constructor.
+     * @param sendToModel when {@code false}, the parameter is hidden from the model
+     *                    (it will not be included in prompts/echo messages).
+     *                    When {@code true}, the parameter is exposed to the model.
+     *                    If this constructor is not called, the default is {@code null}
+     *                    (treated as exposed).
+     */
     public AiAgentParameter(String name, String description, Boolean sendToModel) {
         this(name, description);
         this.sendToModel = sendToModel;
@@ -48,5 +64,13 @@ public String getDescription() {
     public void setDescription(String description) {
         this.description = description;
     }
+
+    public Map<String, Object> toJson() {
+        Map<String, Object> json = new HashMap<>();
+        json.put("Name", this.name);
+        json.put("Description", this.description);
+        json.put("SendToModel", this.sendToModel);
+        return json;
+    }
 }
 

src/main/java/net/ravendb/client/documents/operations/AI/agents/ConversationResult.java

@@ -1,7 +1,10 @@
 package net.ravendb.client.documents.operations.AI.agents;
 
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import net.ravendb.client.documents.operations.AI.AiUsage;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.List;
 
 public class ConversationResult<TAnswer> {
@@ -16,6 +19,45 @@ public class ConversationResult<TAnswer> {
     public ConversationResult() {
     }
 
+    public static <TAnswer> ConversationResult<TAnswer> convert(
+            JsonNode response,
+            Class<TAnswer> answerClass
+    ) throws Exception {
+        ObjectMapper mapper = new ObjectMapper();
+
+        JsonNode totalUsageNode = response.get("TotalUsage");
+        JsonNode resultNode = response.get("Response");
+        String conversationId = response.has("ConversationId") ? response.get("ConversationId").asText() : null;
+        String changeVector = response.has("ChangeVector") ? response.get("ChangeVector").asText() : null;
+        JsonNode usageNode = response.get("Usage");
+        Duration elapsed = response.has("Elapsed") ? Duration.parse(response.get("Elapsed").asText()) : null;
+
+        List<AiAgentActionRequest> requests = null;
+        if (response.has("ActionRequests") && response.get("ActionRequests").isArray()) {
+            requests = new ArrayList<>();
+            for (JsonNode actionRequestNode : response.get("ActionRequests")) {
+                AiAgentActionRequest r = mapper.treeToValue(actionRequestNode, AiAgentActionRequest.class);
+                requests.add(r);
+            }
+        }
+
+        TAnswer answer = null;
+        if (resultNode != null && !resultNode.isNull()) {
+            answer = mapper.treeToValue(resultNode, answerClass);
+        }
+
+        ConversationResult<TAnswer> result = new ConversationResult<>();
+        result.setConversationId(conversationId);
+        result.setChangeVector(changeVector);
+        result.setActionRequests(requests);
+        result.setTotalUsage(totalUsageNode != null ? mapper.treeToValue(totalUsageNode, AiUsage.class) : null);
+        result.setResponse(answer);
+        result.setUsage(usageNode != null ? mapper.treeToValue(usageNode, AiUsage.class) : null);
+        result.setElapsed(elapsed);
+
+        return result;
+    }
+
     public String getConversationId() {
         return conversationId;
     }