Skip to content

Commit ee65ff8

Browse files
LielNagarLielNagar
committed
RDBC-959: added real agent basic api test
1 parent f331779 commit ee65ff8

File tree

9 files changed

+155
-2
lines changed

9 files changed

+155
-2
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -242,7 +242,7 @@ public <TAnswer> CompletableFuture<AiAnswer<TAnswer>> run() {
242242
@SuppressWarnings("unchecked")
243243
AiAnswer<TAnswer> result = (AiAnswer<TAnswer>) runInternal(null, null).join();
244244

245-
if ("Done".equals(result.getStatus())) {
245+
if (result.getStatus() == AiConversationResult.Done) {
246246
return result;
247247
}
248248

‎src/main/java/net/ravendb/client/documents/AI/AiConversationCreationOptions.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package net.ravendb.client.documents.AI;
22

3+
import java.util.HashMap;
34
import java.util.Map;
45

56
public class AiConversationCreationOptions {
@@ -22,6 +23,14 @@ public void setParameters(Map<String, Object> parameters) {
2223
this.parameters = parameters;
2324
}
2425

26+
public AiConversationCreationOptions addParameter(String name, Object value) {
27+
if (this.parameters == null) {
28+
this.parameters = new HashMap<>();
29+
}
30+
this.parameters.put(name, value);
31+
return this;
32+
}
33+
2534
public Integer getExpirationInSec() {
2635
return expirationInSec;
2736
}

‎src/main/java/net/ravendb/client/documents/commands/RunConversationCommand.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.fasterxml.jackson.core.JsonGenerator;
44
import com.fasterxml.jackson.core.type.TypeReference;
5+
import com.fasterxml.jackson.databind.MapperFeature;
56
import com.fasterxml.jackson.databind.ObjectMapper;
67
import com.fasterxml.jackson.databind.node.ObjectNode;
78
import net.ravendb.client.documents.conventions.DocumentConventions;
@@ -126,6 +127,13 @@ public CompletableFuture<String> setResponseAsync(InputStream bodyStream, boolea
126127
return this.parseResponseDefaultAsync(bodyStream);
127128
}
128129

130+
@Override
131+
public void setResponse(String response, boolean fromCache) throws IOException {
132+
ObjectMapper mapper = new ObjectMapper();
133+
mapper.configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true);
134+
this.result = mapper.readValue(response, ConversationResult.class);
135+
}
136+
129137
private CompletableFuture<String> parseResponseDefaultAsync(InputStream bodyStream) {
130138
return CompletableFuture.supplyAsync(() -> {
131139
try {

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,23 @@ public class AiUsage {
55
private int completionTokens;
66
private int totalTokens;
77
private int cachedTokens;
8+
private int reasoningTokens;
89

910
public AiUsage() {
1011
}
1112

12-
public AiUsage(int promptTokens, int completionTokens, int totalTokens, int cachedTokens) {
13+
public AiUsage(int promptTokens, int completionTokens, int totalTokens, int cachedTokens, int reasoningTokens) {
1314
this.promptTokens = promptTokens;
1415
this.completionTokens = completionTokens;
1516
this.totalTokens = totalTokens;
1617
this.cachedTokens = cachedTokens;
18+
this.reasoningTokens = reasoningTokens;
1719
}
1820

21+
public int getReasoningTokens() { return reasoningTokens; }
22+
23+
public void setReasoningTokens(int reasoningTokens) { this.reasoningTokens = reasoningTokens; }
24+
1925
public int getPromptTokens() {
2026
return promptTokens;
2127
}

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package net.ravendb.client.documents.operations.AI.agents;
22

3+
import com.fasterxml.jackson.core.JsonProcessingException;
4+
import com.fasterxml.jackson.databind.ObjectMapper;
35
import net.ravendb.client.documents.conventions.DocumentConventions;
46
import net.ravendb.client.documents.operations.IMaintenanceOperation;
57
import net.ravendb.client.http.RavenCommand;
@@ -32,11 +34,21 @@ public TcpConnectionHeaderMessage.OperationResultType getResultType() {
3234

3335
@Override
3436
public RavenCommand<AiAgentConfigurationResult> getCommand(DocumentConventions conventions) {
37+
String json = toJson(sampleObject);
38+
configuration.setSampleObject(json);
3539
return new AddOrUpdateAiAgentCommand(this.configuration, this.sampleObject, conventions);
3640
}
3741

3842
public static boolean hasNoSampleObjectAndScheme(AiAgentConfiguration configuration) {
3943
return (configuration.getOutputSchema() == null || configuration.getOutputSchema().trim().isEmpty())
4044
&& (configuration.getSampleObject() == null || configuration.getSampleObject().trim().isEmpty());
4145
}
46+
47+
private String toJson(Object obj){
48+
try {
49+
return new ObjectMapper().writeValueAsString(obj);
50+
} catch (JsonProcessingException e) {
51+
throw new RuntimeException("Failed to serialize object to JSON", e);
52+
}
53+
}
4254
}

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package net.ravendb.client.documents.operations.AI.agents;
22

3+
import java.util.ArrayList;
34
import java.util.List;
45

56
public class AiAgentConfiguration {
@@ -18,6 +19,22 @@ public class AiAgentConfiguration {
1819
public AiAgentConfiguration() {
1920
}
2021

22+
public AiAgentConfiguration(String name, String connectionStringName, String systemPrompt) {
23+
if (name == null || name.isEmpty()) {
24+
throw new IllegalArgumentException("Name cannot be null or empty");
25+
}
26+
if (connectionStringName == null || connectionStringName.isEmpty()) {
27+
throw new IllegalArgumentException("connectionStringName cannot be null or empty");
28+
}
29+
if (systemPrompt == null || systemPrompt.isEmpty()) {
30+
throw new IllegalArgumentException("systemPrompt cannot be null or empty");
31+
}
32+
33+
this.name = name;
34+
this.connectionStringName = connectionStringName;
35+
this.systemPrompt = systemPrompt;
36+
}
37+
2138
public AiAgentConfiguration(String identifier, String name, String connectionStringName, String systemPrompt,
2239
String sampleObject, String outputSchema, List<AiAgentToolQuery> queries,
2340
List<AiAgentToolAction> actions, List<AiAgentParameter> parameters,
@@ -100,6 +117,9 @@ public void setActions(List<AiAgentToolAction> actions) {
100117
}
101118

102119
public List<AiAgentParameter> getParameters() {
120+
if (parameters == null) {
121+
parameters = new ArrayList<>();
122+
}
103123
return parameters;
104124
}
105125

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,10 @@ public class AiAgentParameter {
77
public AiAgentParameter() {
88
}
99

10+
public AiAgentParameter(String name) {
11+
this.name = name;
12+
}
13+
1014
public AiAgentParameter(String name, String description) {
1115
this.name = name;
1216
this.description = description;

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,14 @@ public class AiAgentToolQuery {
1111
public AiAgentToolQuery() {
1212
}
1313

14+
public AiAgentToolQuery(String name, String description, String query,
15+
String parametersSampleObject) {
16+
this.name = name;
17+
this.description = description;
18+
this.query = query;
19+
this.parametersSampleObject = parametersSampleObject;
20+
}
21+
1422
public AiAgentToolQuery(String name, String description, String query,
1523
String parametersSampleObject, String parametersSchema) {
1624
this.name = name;

‎src/test/java/net/ravendb/client/test/client/documents/AI/AiAgentTests.java‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,95 @@
11
package net.ravendb.client.test.client.documents.AI;
22

33
import net.ravendb.client.RemoteTestBase;
4+
import net.ravendb.client.documents.AI.AiAnswer;
5+
import net.ravendb.client.documents.AI.AiConversation;
6+
import net.ravendb.client.documents.AI.AiConversationCreationOptions;
7+
import net.ravendb.client.documents.AI.AiConversationResult;
48
import net.ravendb.client.documents.IDocumentStore;
9+
import net.ravendb.client.documents.operations.AI.AiConnectionString;
10+
import net.ravendb.client.documents.operations.AI.AiModelType;
11+
import net.ravendb.client.documents.operations.AI.OpenAiSettings;
512
import net.ravendb.client.documents.operations.AI.agents.*;
613
import net.ravendb.client.documents.operations.AI.agents.AiAgentConfiguration;
14+
import net.ravendb.client.documents.operations.IMaintenanceOperation;
15+
import net.ravendb.client.documents.operations.connectionStrings.GetConnectionStringsOperation;
16+
import net.ravendb.client.documents.operations.connectionStrings.GetConnectionStringsResult;
717
import net.ravendb.client.documents.operations.connectionStrings.PutConnectionStringOperation;
18+
import net.ravendb.client.documents.operations.connectionStrings.PutConnectionStringResult;
819
import net.ravendb.client.documents.operations.etl.RavenConnectionString;
20+
import net.ravendb.client.documents.session.IDocumentSession;
921
import net.ravendb.client.infrastructure.EnableOnServer;
1022
import org.junit.jupiter.api.Assertions;
23+
import org.junit.jupiter.api.Disabled;
1124
import org.junit.jupiter.api.Test;
1225
import static org.assertj.core.api.Assertions.assertThat;
1326
import static org.junit.jupiter.api.Assertions.assertEquals;
1427
import static org.junit.jupiter.api.Assertions.assertNotNull;
1528
import java.util.ArrayList;
29+
import java.util.Arrays;
1630
import java.util.Collections;
31+
import java.util.List;
1732

1833
@EnableOnServer(thresholdVersion = "7.1")
1934
public class AiAgentTests extends RemoteTestBase {
2035

36+
@Disabled
37+
@Test
38+
public void AiAgentClientApiBasicTest(){
39+
String apiKey = System.getenv("RAVENDB_JAVA_TESTS_OPENAI_API_KEY");
40+
assertNotNull(apiKey, "OpenAI API key is not set in environment variable RAVENDB_JAVA_TESTS_OPENAI_API_KEY");
41+
try (IDocumentStore store = getDocumentStore()) {
42+
try (IDocumentSession session = store.openSession()) {
43+
OpenAiSettings ai = new OpenAiSettings(apiKey,
44+
"https://api.openai.com/",
45+
"gpt-4o-mini",
46+
null,
47+
null,
48+
null,
49+
0.0);
50+
51+
AiConnectionString openAiCs = new AiConnectionString();
52+
openAiCs.setModelType(AiModelType.Chat);
53+
openAiCs.setName("openai");
54+
openAiCs.setOpenAiSettings(ai);
55+
56+
IMaintenanceOperation<PutConnectionStringResult> putOpenAi = new PutConnectionStringOperation(openAiCs);
57+
store.maintenance().send(putOpenAi);
58+
59+
GetConnectionStringsResult connectionStrings = store.maintenance().send(new GetConnectionStringsOperation());
60+
assertThat(connectionStrings.getAiConnectionStrings()).hasSize(1);
61+
assertThat(connectionStrings.getAiConnectionStrings().get("openai")).isNotNull();
62+
63+
AiAgentConfiguration agent = new AiAgentConfiguration("shopping assistant", "openai", "You are an AI agent of an online shop, helping customers answer queries about that topic only. When talking about orders or products, include the ids as well.");
64+
agent.getParameters().add(new AiAgentParameter("company"));
65+
List<AiAgentToolQuery> queries = new ArrayList<>();
66+
queries.add(new AiAgentToolQuery("ProductSearch", "semantic search the store product catalog", "from Products where vector.search(embedding.text(Name), $query)", "{\"query\": [\"term or phrase to search in the catalog\"]}"));
67+
queries.add(new AiAgentToolQuery("RecentOrder", "Get the recent orders of the current user", "from Orders where Company = $company order by OrderedAt desc limit 10", "{}"));
68+
agent.setQueries(queries);
69+
String identifier = store.ai().createAgent(agent,AnswerSchema.INSTANCE).getIdentifier();
70+
AiConversation chat = store.ai().conversation(identifier,"chats/", new AiConversationCreationOptions().addParameter("company", "companies/90-A"));
71+
72+
chat.setUserPrompt("what goes well with my cheese?");
73+
AiAnswer<AnswerSchema> result = chat.<AnswerSchema>run().get();
74+
assertEquals(AiConversationResult.Done, result.getStatus());
75+
assertNotNull(result.getAnswer());
76+
assertNotNull(chat.getId());
77+
78+
chat.setUserPrompt("what goes well with my cheese?");
79+
result = chat.<AnswerSchema>run().get();
80+
assertEquals(AiConversationResult.Done, result.getStatus());
81+
assertNotNull(result.getAnswer());
82+
83+
chat.setUserPrompt("what cheese goes well with italian food?");
84+
result = chat.<AnswerSchema>run().get();
85+
assertEquals(AiConversationResult.Done, result.getStatus());
86+
assertNotNull(result.getAnswer());
87+
}
88+
} catch (Exception e){
89+
throw new RuntimeException(e);
90+
}
91+
}
92+
2193
@Test
2294
public void canCreateAiAgent() {
2395
try (IDocumentStore store = getDocumentStore()) {
@@ -177,4 +249,18 @@ public void cannotCreateAgentWithoutSchemaOrSampleObject() {
177249
throw new RuntimeException(e);
178250
}
179251
}
252+
253+
private static class AnswerSchema {
254+
public static final AnswerSchema INSTANCE = new AnswerSchema();
255+
public String answer = "Answer to the user question";
256+
public boolean relevant = true;
257+
public List<String> relevantOrdersId = new ArrayList<>(
258+
Arrays.asList("The order ids relevant to the query or response")
259+
);
260+
public List<String> matchingProductsId = new ArrayList<>(
261+
Arrays.asList("All the product ids referenced either by the user or the system")
262+
);
263+
private AnswerSchema() {
264+
}
265+
}
180266
}

0 commit comments

Comments
 (0)