Skip to content

Commit 8fefb4b

Browse files
luigidemasiorpiske
luigidemasi
authored and
orpiske
committed
#185 : Add the ability to override global Service configuration per invocation using Tool definition
1 parent 0bd2fcc commit 8fefb4b

File tree

12 files changed

+213
-17
lines changed

12 files changed

+213
-17
lines changed

‎api/src/main/java/ai/wanaku/api/types/ToolReference.java

+65-3
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,10 @@ public static class Property {
197197

198198
private String type;
199199
private String description;
200+
private String target;
201+
private String scope;
202+
private String value;
203+
200204

201205
/**
202206
* Gets the type of the property.
@@ -234,19 +238,77 @@ public void setDescription(String description) {
234238
this.description = description;
235239
}
236240

241+
/**
242+
* Gets the target of the property.
243+
*
244+
* @return the description of the property
245+
*/
246+
public String getTarget() {
247+
return target;
248+
}
249+
250+
/**
251+
* Sets the target of the property.
252+
*
253+
* @param target the new target of the property
254+
*/
255+
public void setTarget(String target) {
256+
this.target = target;
257+
}
258+
259+
/**
260+
* Gets the scope of the property.
261+
*
262+
* @return the scope of the property
263+
*/
264+
public String getScope() {
265+
return scope;
266+
}
267+
268+
/**
269+
* Sets the scope of the property.
270+
*
271+
* @param scope the new scope of the property
272+
*/
273+
public void setScope(String scope) {
274+
this.scope = scope;
275+
}
276+
277+
/**
278+
* Gets the value of the property.
279+
*
280+
* @return the value of the property
281+
*/
282+
public String getValue() {
283+
return value;
284+
}
285+
286+
/**
287+
* Sets the value of the property.
288+
*
289+
* @param value the new value of the property
290+
*/
291+
public void setValue(String value) {
292+
this.value = value;
293+
}
294+
237295
@Override
238296
public boolean equals(Object o) {
239297
if (o == null || getClass() != o.getClass()) {
240298
return false;
241299
}
300+
242301
Property that = (Property) o;
243-
return Objects.equals(type, that.type)
244-
&& Objects.equals(description, that.description);
302+
return Objects.equals(type, that.type) &&
303+
Objects.equals(description, that.description) &&
304+
Objects.equals(target, that.target) &&
305+
Objects.equals(scope, that.scope) &&
306+
Objects.equals(value, that.value);
245307
}
246308

247309
@Override
248310
public int hashCode() {
249-
return Objects.hash(type, description);
311+
return Objects.hash(type, description,target,scope,value);
250312
}
251313
}
252314

‎core/core-exchange/src/main/java/ai/wanaku/core/exchange/ParsedToolInvokeRequest.java

+8-2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33

44
import ai.wanaku.core.uri.Parameter;
5+
import ai.wanaku.core.uri.URIHelper;
56
import ai.wanaku.core.uri.URIParser;
67
import java.util.Collections;
78
import java.util.HashMap;
@@ -10,7 +11,7 @@
1011
/**
1112
* Represents a parsed tool invocation request containing the URI and its body.
1213
*/
13-
public record ParsedToolInvokeRequest(String uri, String body) {
14+
public record ParsedToolInvokeRequest(String uri, String body, Map<String, String> headers) {
1415

1516
/**
1617
* Parses the URI provided by the router
@@ -41,8 +42,13 @@ public static ParsedToolInvokeRequest parseRequest(String uri, ToolInvokeRequest
4142

4243
String parsedUri = URIParser.parse(uri, map);
4344

45+
//Add additional configuration
46+
parsedUri = URIHelper.addQueryParameters(parsedUri, toolInvokeRequest.getServiceConfigurationsMap());
47+
4448
String body = toolInvokeRequest.getBody();
4549

46-
return new ParsedToolInvokeRequest(parsedUri, body);
50+
Map<String, String> headers = toolInvokeRequest.getHeadersMap();
51+
52+
return new ParsedToolInvokeRequest(parsedUri, body, headers);
4753
}
4854
}

‎core/core-exchange/src/main/proto/toolrequest.proto

+1
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ message ToolInvokeRequest {
1919
map<string, string> arguments = 3;
2020
map<string, string> serviceConfigurations = 4;
2121
map<string, string> credentialsConfigurations = 5;
22+
map<string, string> headers = 6;
2223
}
2324

2425
// The invocation response message

‎core/core-util/src/main/java/ai/wanaku/core/uri/URIHelper.java

+15-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package ai.wanaku.core.uri;
22

3+
34
import java.util.Map;
45

56
/**
@@ -16,7 +17,10 @@ private static String buildFromBaseUri(StringBuilder uri, Map<String, ?> paramet
1617
}
1718

1819
private static void buildQuery(StringBuilder uri, Map<String, ?> parameters) {
19-
boolean first = true;
20+
buildQuery(uri, parameters, true);
21+
}
22+
23+
private static void buildQuery(StringBuilder uri, Map<String, ?> parameters, boolean first) {
2024
for (Map.Entry<String, ?> entry : parameters.entrySet()) {
2125
if (first) {
2226
uri.append('?');
@@ -66,4 +70,14 @@ public static String buildUri(String scheme, String path, Map<String, ?> paramet
6670
uri.append(path);
6771
return buildFromBaseUri(uri, parameters);
6872
}
73+
74+
public static String addQueryParameters(String uri, Map<String, ?> parameters) {
75+
boolean first = true;
76+
if (uri.contains("?")) {
77+
first = false;
78+
}
79+
StringBuilder uriBuilder = new StringBuilder(uri);
80+
buildQuery(uriBuilder, parameters, first);
81+
return uriBuilder.toString();
82+
}
6983
}

‎core/core-util/src/main/java/ai/wanaku/core/util/ReservedPropertyNames.java

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package ai.wanaku.core.util;
2+
3+
/**
4+
* Reserved argument names
5+
*/
6+
public class ReservedPropertyNames {
7+
8+
/**
9+
* Indicate that the argument is intended to use as a tool/service configuration
10+
*/
11+
public static final String TARGET_CONFIGURATION = "configuration";
12+
13+
/**
14+
* Indicate that the argument is intended as header for services/tools that support them
15+
*/
16+
public static final String TARGET_HEADER = "header";
17+
18+
/**
19+
* Indicate that the argument is intended to be used on the service
20+
*/
21+
public static final String SCOPE_SERVICE = "service";
22+
23+
/**
24+
* Indicate that the argument is intended to be used on the service endpoint
25+
*/
26+
public static final String SCOPE_SERVICE_ENDPOINT = "service-endpoint";
27+
}

‎core/core-util/src/test/java/ai/wanaku/core/uri/URIHelperTest.java

+21
Original file line numberDiff line numberDiff line change
@@ -33,4 +33,25 @@ void testQueryTwoParams() {
3333
assertEquals("ftp://test?key1=value1&key2=value2", uri);
3434
}
3535

36+
@Test
37+
void testAddQueryParameterToBaseUri() {
38+
SortedMap <String, String> params = new TreeMap<>();
39+
params.put("key1", "value1");
40+
params.put("key2", "value2");
41+
String uri = URIHelper.addQueryParameters("ftp://test", params);
42+
assertEquals("ftp://test?key1=value1&key2=value2", uri);
43+
}
44+
45+
@Test
46+
void testAddQueryParametersToUriAlreadyContainsQueryParameters() {
47+
SortedMap <String, String> params = new TreeMap<>();
48+
params.put("key1", "value1");
49+
params.put("key2", "value2");
50+
String uri = URIHelper.addQueryParameters("ftp://test?foo=bar&dead=false", params);
51+
assertEquals("ftp://test?foo=bar&dead=false&key1=value1&key2=value2", uri);
52+
}
53+
54+
55+
56+
3657
}

‎routers/wanaku-router/src/main/java/ai/wanaku/routers/proxies/tools/InvokerProxy.java

+44-2
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import ai.wanaku.core.exchange.ToolInvokerGrpc;
1313
import ai.wanaku.core.mcp.providers.ServiceRegistry;
1414
import ai.wanaku.core.util.CollectionsHelper;
15-
import ai.wanaku.core.util.ReservedArgumentNames;
1615
import ai.wanaku.routers.proxies.ToolsProxy;
1716
import com.google.protobuf.ProtocolStringList;
1817
import io.grpc.ManagedChannel;
@@ -23,8 +22,16 @@
2322
import java.util.ArrayList;
2423
import java.util.List;
2524
import java.util.Map;
25+
import java.util.stream.Collectors;
26+
2627
import org.jboss.logging.Logger;
2728

29+
import static ai.wanaku.core.util.ReservedArgumentNames.BODY;
30+
import static ai.wanaku.core.util.ReservedPropertyNames.SCOPE_SERVICE;
31+
import static ai.wanaku.core.util.ReservedPropertyNames.SCOPE_SERVICE_ENDPOINT;
32+
import static ai.wanaku.core.util.ReservedPropertyNames.TARGET_CONFIGURATION;
33+
import static ai.wanaku.core.util.ReservedPropertyNames.TARGET_HEADER;
34+
2835
/**
2936
* A proxy class for invoking tools
3037
*/
@@ -72,12 +79,47 @@ private static ToolInvokeReply invokeRemotely(
7279
Map<String, String> serviceConfigurations = Configurations.toStringMap(configurations);
7380
Map<String, String> argumentsMap = CollectionsHelper.toStringStringMap(toolArguments.args());
7481

82+
Map<String, ToolReference.Property> inputSchema =toolReference
83+
.getInputSchema()
84+
.getProperties();
85+
86+
Map<String,String> serviceEndpointConfiguration = inputSchema.entrySet()
87+
.stream()
88+
.filter(entry -> {
89+
ToolReference.Property property = entry.getValue();
90+
return property !=null &&
91+
property.getTarget() != null &&
92+
property.getScope() != null &&
93+
property.getTarget().equals(TARGET_CONFIGURATION) &&
94+
property.getScope().equals(SCOPE_SERVICE_ENDPOINT);
95+
})
96+
.collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().getValue()));
97+
98+
//override the global configuration with the invocation specific one.
99+
serviceConfigurations.putAll(serviceEndpointConfiguration);
100+
101+
102+
// extract headers parameter
103+
Map<String, String> headers = inputSchema.entrySet()
104+
.stream()
105+
.filter(entry -> {
106+
ToolReference.Property property = entry.getValue();
107+
return property !=null &&
108+
property.getTarget() != null &&
109+
property.getScope() != null &&
110+
property.getTarget().equals(TARGET_HEADER) &&
111+
property.getScope().equals(SCOPE_SERVICE);
112+
})
113+
.collect(Collectors.toMap(Map.Entry::getKey, entry -> entry.getValue().getValue()));
114+
115+
75116
String body = extractBody(toolReference);
76117

77118
ToolInvokeRequest toolInvokeRequest = ToolInvokeRequest.newBuilder()
78119
.setBody(body)
79120
.setUri(toolReference.getUri())
80121
.putAllServiceConfigurations(serviceConfigurations)
122+
.putAllHeaders(headers)
81123
.putAllArguments(argumentsMap)
82124
.build();
83125

@@ -87,7 +129,7 @@ private static ToolInvokeReply invokeRemotely(
87129

88130
private static String extractBody(ToolReference toolReference) {
89131
Map<String, ToolReference.Property> properties = toolReference.getInputSchema().getProperties();
90-
ToolReference.Property bodyProp = properties.get(ReservedArgumentNames.BODY);
132+
ToolReference.Property bodyProp = properties.get(BODY);
91133
if (bodyProp == null) {
92134
return EMPTY_BODY;
93135
}

‎routers/wanaku-router/src/main/webui/openapi.json

+9
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,15 @@
6363
},
6464
"description" : {
6565
"type" : "string"
66+
},
67+
"target" : {
68+
"type" : "string"
69+
},
70+
"scope" : {
71+
"type" : "string"
72+
},
73+
"value" : {
74+
"type" : "string"
6675
}
6776
}
6877
},

‎routers/wanaku-router/src/main/webui/openapi.yaml

+6
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,12 @@ components:
4343
type: string
4444
description:
4545
type: string
46+
target:
47+
type: string
48+
scope:
49+
type: string
50+
value:
51+
type: string
4652
ResourceReference:
4753
type: object
4854
properties:

‎services/tools/wanaku-routing-http-service/pom.xml

+2
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
<version>${project.version}</version>
3131
</dependency>
3232

33+
34+
3335
<dependency>
3436
<groupId>io.quarkus</groupId>
3537
<artifactId>quarkus-arc</artifactId>

‎services/tools/wanaku-routing-http-service/src/main/java/ai/wanaku/routing/service/HttpClient.java

+12-9
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,20 @@
11
package ai.wanaku.routing.service;
22

3+
import ai.wanaku.api.exceptions.WanakuException;
4+
import ai.wanaku.core.services.routing.Client;
35
import jakarta.enterprise.context.ApplicationScoped;
46

57
import ai.wanaku.core.exchange.ParsedToolInvokeRequest;
68
import ai.wanaku.core.exchange.ToolInvokeRequest;
7-
import ai.wanaku.core.services.routing.Client;
89
import org.apache.camel.ProducerTemplate;
910
import org.jboss.logging.Logger;
1011

12+
import java.util.HashMap;
13+
import java.util.Map;
14+
1115
@ApplicationScoped
1216
public class HttpClient implements Client {
17+
1318
private static final Logger LOG = Logger.getLogger(HttpClient.class);
1419

1520
private final ProducerTemplate producer;
@@ -18,20 +23,18 @@ public HttpClient(ProducerTemplate producer) {
1823
this.producer = producer;
1924
}
2025

26+
2127
@Override
22-
public Object exchange(ToolInvokeRequest request) {
28+
public Object exchange(ToolInvokeRequest request) throws WanakuException {
2329
producer.start();
2430

2531
ParsedToolInvokeRequest parsedRequest = ParsedToolInvokeRequest.parseRequest(request);
2632

2733
LOG.infof("Invoking tool at URI: %s", parsedRequest.uri());
2834

29-
String s;
30-
if (parsedRequest.body().isEmpty()) {
31-
s = producer.requestBody(parsedRequest.uri(), null, String.class);
32-
} else {
33-
s = producer.requestBody(parsedRequest.uri(), parsedRequest.body(), String.class);
34-
}
35-
return s;
35+
36+
Map<String,Object> headers = new HashMap<>(parsedRequest.headers());
37+
38+
return producer.requestBodyAndHeaders(parsedRequest.uri(), parsedRequest.body(), headers, String.class);
3639
}
3740
}

‎ui/src/models/property.ts

+3
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,7 @@
88
export interface Property {
99
type?: string;
1010
description?: string;
11+
target?: string;
12+
scope?: string;
13+
value?: string;
1114
}

0 commit comments

Comments
 (0)