Skip to content

Commit

Permalink
feat: [vertexai] allow setting ToolConfig and SystemInstruction in Ch…
Browse files Browse the repository at this point in the history 
…atSession

PiperOrigin-RevId: 641987014
  • Loading branch information
@jaycee-li @copybara-github
jaycee-li authored and copybara-github committed Jun 10, 2024
1 parent 4db0d1d commit 16c7781
Show file tree
Hide file tree
Showing 63 changed files with 6,825 additions and 3,795 deletions.
3 changes: 3 additions & 0 deletions ...e-cloud-vertexai/src/main/java/com/google/cloud/vertexai/api/PredictionServiceClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1410,6 +1410,7 @@ public final GenerateContentResponse generateContent(String model, List<Content>
* .addAllContents(new ArrayList<Content>())
* .setSystemInstruction(Content.newBuilder().build())
* .addAllTools(new ArrayList<Tool>())
* .setToolConfig(ToolConfig.newBuilder().build())
* .addAllSafetySettings(new ArrayList<SafetySetting>())
* .setGenerationConfig(GenerationConfig.newBuilder().build())
* .build();
Expand Down Expand Up @@ -1443,6 +1444,7 @@ public final GenerateContentResponse generateContent(GenerateContentRequest requ
* .addAllContents(new ArrayList<Content>())
* .setSystemInstruction(Content.newBuilder().build())
* .addAllTools(new ArrayList<Tool>())
* .setToolConfig(ToolConfig.newBuilder().build())
* .addAllSafetySettings(new ArrayList<SafetySetting>())
* .setGenerationConfig(GenerationConfig.newBuilder().build())
* .build();
Expand Down Expand Up @@ -1477,6 +1479,7 @@ public final GenerateContentResponse generateContent(GenerateContentRequest requ
* .addAllContents(new ArrayList<Content>())
* .setSystemInstruction(Content.newBuilder().build())
* .addAllTools(new ArrayList<Tool>())
* .setToolConfig(ToolConfig.newBuilder().build())
* .addAllSafetySettings(new ArrayList<SafetySetting>())
* .setGenerationConfig(GenerationConfig.newBuilder().build())
* .build();
Expand Down
18 changes: 0 additions & 18 deletions ...ertexai/src/main/java/com/google/cloud/vertexai/api/stub/EndpointServiceStubSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -390,15 +390,6 @@ public EndpointServiceStub createStub() throws IOException {
"Transport not supported: %s", getTransportChannelProvider().getTransportName()));
}

/** Returns the endpoint set by the user or the the service's default endpoint. */
@Override
public String getEndpoint() {
if (super.getEndpoint() != null) {
return super.getEndpoint();
}
return getDefaultEndpoint();
}

/** Returns the default service name. */
@Override
public String getServiceName() {
Expand Down Expand Up @@ -998,15 +989,6 @@ public UnaryCallSettings.Builder<GetIamPolicyRequest, Policy> getIamPolicySettin
return testIamPermissionsSettings;
}

/** Returns the endpoint set by the user or the the service's default endpoint. */
@Override
public String getEndpoint() {
if (super.getEndpoint() != null) {
return super.getEndpoint();
}
return getDefaultEndpoint();
}

@Override
public EndpointServiceStubSettings build() throws IOException {
return new EndpointServiceStubSettings(this);
Expand Down
14 changes: 7 additions & 7 deletions ...-vertexai/src/main/java/com/google/cloud/vertexai/api/stub/GrpcLlmUtilityServiceStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,14 +54,14 @@
@Generated("by gapic-generator-java")
public class GrpcLlmUtilityServiceStub extends LlmUtilityServiceStub {
private static final MethodDescriptor<CountTokensRequest, CountTokensResponse>
// TODO(b/317255628): switch back to the google.cloud.aiplatform.v1.LlmUtilityServiceClient
countTokensMethodDescriptor =
MethodDescriptor.<CountTokensRequest, CountTokensResponse>newBuilder()
.setType(MethodDescriptor.MethodType.UNARY)
.setFullMethodName("google.cloud.aiplatform.v1beta1.PredictionService/CountTokens")
.setRequestMarshaller(ProtoUtils.marshaller(CountTokensRequest.getDefaultInstance()))
.setResponseMarshaller(ProtoUtils.marshaller(CountTokensResponse.getDefaultInstance()))
.build();
MethodDescriptor.<CountTokensRequest, CountTokensResponse>newBuilder()
.setType(MethodDescriptor.MethodType.UNARY)
.setFullMethodName("google.cloud.aiplatform.v1.LlmUtilityService/CountTokens")
.setRequestMarshaller(ProtoUtils.marshaller(CountTokensRequest.getDefaultInstance()))
.setResponseMarshaller(
ProtoUtils.marshaller(CountTokensResponse.getDefaultInstance()))
.build();

private static final MethodDescriptor<ComputeTokensRequest, ComputeTokensResponse>
computeTokensMethodDescriptor =
Expand Down
12 changes: 12 additions & 0 deletions ...-vertexai/src/main/java/com/google/cloud/vertexai/api/stub/GrpcPredictionServiceStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -385,12 +385,24 @@ protected GrpcPredictionServiceStub(
streamDirectPredictTransportSettings =
GrpcCallSettings.<StreamDirectPredictRequest, StreamDirectPredictResponse>newBuilder()
.setMethodDescriptor(streamDirectPredictMethodDescriptor)
.setParamsExtractor(
request -> {
RequestParamsBuilder builder = RequestParamsBuilder.create();
builder.add("endpoint", String.valueOf(request.getEndpoint()));
return builder.build();
})
.build();
GrpcCallSettings<StreamDirectRawPredictRequest, StreamDirectRawPredictResponse>
streamDirectRawPredictTransportSettings =
GrpcCallSettings
.<StreamDirectRawPredictRequest, StreamDirectRawPredictResponse>newBuilder()
.setMethodDescriptor(streamDirectRawPredictMethodDescriptor)
.setParamsExtractor(
request -> {
RequestParamsBuilder builder = RequestParamsBuilder.create();
builder.add("endpoint", String.valueOf(request.getEndpoint()));
return builder.build();
})
.build();
GrpcCallSettings<StreamingPredictRequest, StreamingPredictResponse>
streamingPredictTransportSettings =
Expand Down
45 changes: 45 additions & 0 deletions ...ertexai/src/main/java/com/google/cloud/vertexai/api/stub/HttpJsonEndpointServiceStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -688,6 +688,16 @@ protected HttpJsonEndpointServiceStub(
"google.longrunning.Operations.CancelOperation",
HttpRule.newBuilder()
.setPost("/ui/{name=projects/*/locations/*/operations/*}:cancel")
.addAdditionalBindings(
HttpRule.newBuilder()
.setPost(
"/ui/{name=projects/*/locations/*/agents/*/operations/*}:cancel")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setPost(
"/ui/{name=projects/*/locations/*/apps/*/operations/*}:cancel")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setPost(
Expand Down Expand Up @@ -1067,6 +1077,15 @@ protected HttpJsonEndpointServiceStub(
"google.longrunning.Operations.DeleteOperation",
HttpRule.newBuilder()
.setDelete("/ui/{name=projects/*/locations/*/operations/*}")
.addAdditionalBindings(
HttpRule.newBuilder()
.setDelete(
"/ui/{name=projects/*/locations/*/agents/*/operations/*}")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setDelete("/ui/{name=projects/*/locations/*/apps/*/operations/*}")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setDelete(
Expand Down Expand Up @@ -1476,6 +1495,14 @@ protected HttpJsonEndpointServiceStub(
"google.longrunning.Operations.GetOperation",
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*/operations/*}")
.addAdditionalBindings(
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*/agents/*/operations/*}")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*/apps/*/operations/*}")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*/datasets/*/operations/*}")
Expand Down Expand Up @@ -1892,6 +1919,14 @@ protected HttpJsonEndpointServiceStub(
"google.longrunning.Operations.ListOperations",
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*}/operations")
.addAdditionalBindings(
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*/agents/*}/operations")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*/apps/*}/operations")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setGet("/ui/{name=projects/*/locations/*/datasets/*}/operations")
Expand Down Expand Up @@ -2294,6 +2329,16 @@ protected HttpJsonEndpointServiceStub(
"google.longrunning.Operations.WaitOperation",
HttpRule.newBuilder()
.setPost("/ui/{name=projects/*/locations/*/operations/*}:wait")
.addAdditionalBindings(
HttpRule.newBuilder()
.setPost(
"/ui/{name=projects/*/locations/*/agents/*/operations/*}:wait")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setPost(
"/ui/{name=projects/*/locations/*/apps/*/operations/*}:wait")
.build())
.addAdditionalBindings(
HttpRule.newBuilder()
.setPost(
Expand Down
71 changes: 35 additions & 36 deletions ...texai/src/main/java/com/google/cloud/vertexai/api/stub/HttpJsonLlmUtilityServiceStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,43 +63,42 @@ public class HttpJsonLlmUtilityServiceStub extends LlmUtilityServiceStub {
private static final TypeRegistry typeRegistry = TypeRegistry.newBuilder().build();

private static final ApiMethodDescriptor<CountTokensRequest, CountTokensResponse>
// TODO(b/317255628): switch back to the google.cloud.aiplatform.v1.LlmUtilityServiceClient
countTokensMethodDescriptor =
ApiMethodDescriptor.<CountTokensRequest, CountTokensResponse>newBuilder()
.setFullMethodName("google.cloud.aiplatform.v1beta1.PredictionService/CountTokens")
.setHttpMethod("POST")
.setType(ApiMethodDescriptor.MethodType.UNARY)
.setRequestFormatter(
ProtoMessageRequestFormatter.<CountTokensRequest>newBuilder()
.setPath(
"/v1/{endpoint=projects/*/locations/*/endpoints/*}:countTokens",
request -> {
Map<String, String> fields = new HashMap<>();
ProtoRestSerializer<CountTokensRequest> serializer =
ProtoRestSerializer.create();
serializer.putPathParam(fields, "endpoint", request.getEndpoint());
return fields;
})
.setAdditionalPaths(
"/v1/{endpoint=projects/*/locations/*/publishers/*/models/*}:countTokens")
.setQueryParamsExtractor(
request -> {
Map<String, List<String>> fields = new HashMap<>();
ProtoRestSerializer<CountTokensRequest> serializer =
ProtoRestSerializer.create();
return fields;
})
.setRequestBodyExtractor(
request ->
ProtoRestSerializer.create()
.toBody("*", request.toBuilder().clearEndpoint().build(), false))
.build())
.setResponseParser(
ProtoMessageResponseParser.<CountTokensResponse>newBuilder()
.setDefaultInstance(CountTokensResponse.getDefaultInstance())
.setDefaultTypeRegistry(typeRegistry)
.build())
.build();
ApiMethodDescriptor.<CountTokensRequest, CountTokensResponse>newBuilder()
.setFullMethodName("google.cloud.aiplatform.v1.LlmUtilityService/CountTokens")
.setHttpMethod("POST")
.setType(ApiMethodDescriptor.MethodType.UNARY)
.setRequestFormatter(
ProtoMessageRequestFormatter.<CountTokensRequest>newBuilder()
.setPath(
"/v1/{endpoint=projects/*/locations/*/endpoints/*}:countTokens",
request -> {
Map<String, String> fields = new HashMap<>();
ProtoRestSerializer<CountTokensRequest> serializer =
ProtoRestSerializer.create();
serializer.putPathParam(fields, "endpoint", request.getEndpoint());
return fields;
})
.setAdditionalPaths(
"/v1/{endpoint=projects/*/locations/*/publishers/*/models/*}:countTokens")
.setQueryParamsExtractor(
request -> {
Map<String, List<String>> fields = new HashMap<>();
ProtoRestSerializer<CountTokensRequest> serializer =
ProtoRestSerializer.create();
return fields;
})
.setRequestBodyExtractor(
request ->
ProtoRestSerializer.create()
.toBody("*", request.toBuilder().clearEndpoint().build(), false))
.build())
.setResponseParser(
ProtoMessageResponseParser.<CountTokensResponse>newBuilder()
.setDefaultInstance(CountTokensResponse.getDefaultInstance())
.setDefaultTypeRegistry(typeRegistry)
.build())
.build();

private static final ApiMethodDescriptor<ComputeTokensRequest, ComputeTokensResponse>
computeTokensMethodDescriptor =
Expand Down
4 changes: 4 additions & 0 deletions ...texai/src/main/java/com/google/cloud/vertexai/api/stub/HttpJsonPredictionServiceStub.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -211,6 +211,8 @@ public class HttpJsonPredictionServiceStub extends PredictionServiceStub {
serializer.putPathParam(fields, "endpoint", request.getEndpoint());
return fields;
})
.setAdditionalPaths(
"/v1/{endpoint=projects/*/locations/*/publishers/*/models/*}:directPredict")
.setQueryParamsExtractor(
request -> {
Map<String, List<String>> fields = new HashMap<>();
Expand Down Expand Up @@ -247,6 +249,8 @@ public class HttpJsonPredictionServiceStub extends PredictionServiceStub {
serializer.putPathParam(fields, "endpoint", request.getEndpoint());
return fields;
})
.setAdditionalPaths(
"/v1/{endpoint=projects/*/locations/*/publishers/*/models/*}:directRawPredict")
.setQueryParamsExtractor(
request -> {
Map<String, List<String>> fields = new HashMap<>();
Expand Down
18 changes: 0 additions & 18 deletions ...texai/src/main/java/com/google/cloud/vertexai/api/stub/LlmUtilityServiceStubSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,15 +226,6 @@ public LlmUtilityServiceStub createStub() throws IOException {
"Transport not supported: %s", getTransportChannelProvider().getTransportName()));
}

/** Returns the endpoint set by the user or the the service's default endpoint. */
@Override
public String getEndpoint() {
if (super.getEndpoint() != null) {
return super.getEndpoint();
}
return getDefaultEndpoint();
}

/** Returns the default service name. */
@Override
public String getServiceName() {
Expand Down Expand Up @@ -540,15 +531,6 @@ public UnaryCallSettings.Builder<GetIamPolicyRequest, Policy> getIamPolicySettin
return testIamPermissionsSettings;
}

/** Returns the endpoint set by the user or the the service's default endpoint. */
@Override
public String getEndpoint() {
if (super.getEndpoint() != null) {
return super.getEndpoint();
}
return getDefaultEndpoint();
}

@Override
public LlmUtilityServiceStubSettings build() throws IOException {
return new LlmUtilityServiceStubSettings(this);
Expand Down
23 changes: 4 additions & 19 deletions ...texai/src/main/java/com/google/cloud/vertexai/api/stub/PredictionServiceStubSettings.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,10 @@
public class PredictionServiceStubSettings extends StubSettings<PredictionServiceStubSettings> {
/** The default scopes of the service. */
private static final ImmutableList<String> DEFAULT_SERVICE_SCOPES =
ImmutableList.<String>builder().add("https://www.googleapis.com/auth/cloud-platform").build();
ImmutableList.<String>builder()
.add("https://www.googleapis.com/auth/cloud-platform")
.add("https://www.googleapis.com/auth/cloud-platform.read-only")
.build();

private final UnaryCallSettings<PredictRequest, PredictResponse> predictSettings;
private final UnaryCallSettings<RawPredictRequest, HttpBody> rawPredictSettings;
Expand Down Expand Up @@ -328,15 +331,6 @@ public PredictionServiceStub createStub() throws IOException {
"Transport not supported: %s", getTransportChannelProvider().getTransportName()));
}

/** Returns the endpoint set by the user or the the service's default endpoint. */
@Override
public String getEndpoint() {
if (super.getEndpoint() != null) {
return super.getEndpoint();
}
return getDefaultEndpoint();
}

/** Returns the default service name. */
@Override
public String getServiceName() {
Expand Down Expand Up @@ -806,15 +800,6 @@ public UnaryCallSettings.Builder<GetIamPolicyRequest, Policy> getIamPolicySettin
return testIamPermissionsSettings;
}

/** Returns the endpoint set by the user or the the service's default endpoint. */
@Override
public String getEndpoint() {
if (super.getEndpoint() != null) {
return super.getEndpoint();
}
return getDefaultEndpoint();
}

@Override
public PredictionServiceStubSettings build() throws IOException {
return new PredictionServiceStubSettings(this);
Expand Down
Loading

0 comments on commit 16c7781

Please sign in to comment.