googleapis · jaycee-li · Jun 17, 2024 · May 22, 2024
@@ -20,6 +20,10 @@
 public final class Constants {
   // Constants for VertexAI class
   public static final String USER_AGENT_HEADER = "model-builder";
+  static final String DEFAULT_LOCATION = "us-central1";
+  static final String GOOGLE_CLOUD_REGION = "GOOGLE_CLOUD_REGION";
+  static final String CLOUD_ML_REGION = "CLOUD_ML_REGION";
+  static final String GOOGLE_CLOUD_PROJECT = "GOOGLE_CLOUD_PROJECT";
 
   private Constants() {}
 }
@@ -27,6 +27,7 @@
 import com.google.api.gax.rpc.FixedHeaderProvider;
 import com.google.api.gax.rpc.HeaderProvider;
 import com.google.auth.Credentials;
+import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.vertexai.api.LlmUtilityServiceClient;
 import com.google.cloud.vertexai.api.LlmUtilityServiceSettings;
 import com.google.cloud.vertexai.api.PredictionServiceClient;
@@ -67,6 +68,11 @@ public class VertexAI implements AutoCloseable {
   private final transient Supplier<PredictionServiceClient> predictionClientSupplier;
   private final transient Supplier<LlmUtilityServiceClient> llmClientSupplier;
 
+  @InternalApi
+  static Optional<String> getEnvironmentVariable(String envKey) {
+    return Optional.ofNullable(System.getenv(envKey));
+  }
+
   /**
    * Constructs a VertexAI instance.
    *
@@ -85,6 +91,29 @@ public VertexAI(String projectId, String location) {
         /* llmClientSupplierOpt= */ Optional.empty());
   }
 
+  /**
+   * Constructs a VertexAI instance.
+   *
+   * <p><b>Note:</b> SDK infers location from runtime environment first. If there is no location
+   * inferred from runtime environment, SDK will default location to `us-central1`.
+   *
+   * <p>SDK will infer projectId from runtime environment and GoogleCredentials.
+   *
+   * @throws java.lang.IllegalArgumentException If there is not projectId inferred from either
+   *     runtime environment or GoogleCredentials
+   */
+  public VertexAI() {
+    this(
+        null,
+        null,
+        Transport.GRPC,
+        ImmutableList.of(),
+        /* credentials= */ Optional.empty(),
+        /* apiEndpoint= */ Optional.empty(),
+        /* predictionClientSupplierOpt= */ Optional.empty(),
+        /* llmClientSupplierOpt= */ Optional.empty());
+  }
+
   private VertexAI(
       String projectId,
       String location,
@@ -98,12 +127,8 @@ private VertexAI(
       throw new IllegalArgumentException(
           "At most one of Credentials and scopes should be specified.");
     }
-    checkArgument(!Strings.isNullOrEmpty(projectId), "projectId can't be null or empty");
-    checkArgument(!Strings.isNullOrEmpty(location), "location can't be null or empty");
     checkNotNull(transport, "transport can't be null");
-
-    this.projectId = projectId;
-    this.location = location;
+    this.location = Strings.isNullOrEmpty(location) ? inferLocation() : location;
     this.transport = transport;
 
     if (credentials.isPresent()) {
@@ -118,13 +143,15 @@ private VertexAI(
                   .build();
     }
 
+    this.projectId = Strings.isNullOrEmpty(projectId) ? inferProjectId() : projectId;
     this.predictionClientSupplier =
         Suppliers.memoize(predictionClientSupplierOpt.orElse(this::newPredictionServiceClient));
 
     this.llmClientSupplier =
         Suppliers.memoize(llmClientSupplierOpt.orElse(this::newLlmUtilityClient));
 
-    this.apiEndpoint = apiEndpoint.orElse(String.format("%s-aiplatform.googleapis.com", location));
+    this.apiEndpoint =
+        apiEndpoint.orElse(String.format("%s-aiplatform.googleapis.com", this.location));
   }
 
   /** Builder for {@link VertexAI}. */
@@ -141,8 +168,6 @@ public static class Builder {
     private Supplier<LlmUtilityServiceClient> llmClientSupplier;
 
     public VertexAI build() {
-      checkNotNull(projectId, "projectId must be set.");
-      checkNotNull(location, "location must be set.");
 
       return new VertexAI(
           projectId,
@@ -339,6 +364,32 @@ private LlmUtilityServiceClient newLlmUtilityClient() {
     }
   }
 
+  private String inferProjectId() {
+    final String projectNotFoundErrorMessage =
+        ("Unable to infer your project. Please provide a project Id by one of the following:"
+            + "\n- Passing a constructor argument by using new VertexAI(String projectId, String"
+            + " location)"
+            + "\n- Setting project using 'gcloud config set project my-project'");
+    final Optional<String> projectIdOptional =
+        getEnvironmentVariable(Constants.GOOGLE_CLOUD_PROJECT);
+    if (projectIdOptional.isPresent()) {
+      return projectIdOptional.get();
+    }
+    try {
+      return Optional.ofNullable((GoogleCredentials) this.credentialsProvider.getCredentials())
+          .map((credentials) -> credentials.getQuotaProjectId())
+          .orElseThrow(() -> new IllegalArgumentException(projectNotFoundErrorMessage));
+    } catch (IOException e) {
+      throw new IllegalArgumentException(projectNotFoundErrorMessage, e);
+    }
+  }
+
+  private String inferLocation() {
+    return getEnvironmentVariable(Constants.GOOGLE_CLOUD_REGION)
+        .orElse(
+            getEnvironmentVariable(Constants.CLOUD_ML_REGION).orElse(Constants.DEFAULT_LOCATION));
+  }
+
   private LlmUtilityServiceSettings getLlmUtilityServiceClientSettings() throws IOException {
     LlmUtilityServiceSettings.Builder settingsBuilder;
     if (transport == Transport.REST) {

@@ -29,19 +29,24 @@
 import com.google.cloud.vertexai.api.GenerationConfig;
 import com.google.cloud.vertexai.api.SafetySetting;
 import com.google.cloud.vertexai.api.Tool;
+import com.google.cloud.vertexai.api.ToolConfig;
 import com.google.common.collect.ImmutableList;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
-/** Represents a conversation between the user and the model */
+/**
+ * Represents a conversation between the user and the model.
+ *
+ * <p>Note: this class is NOT thread-safe.
+ */
 public final class ChatSession {
   private final GenerativeModel model;
   private final Optional<ChatSession> rootChatSession;
   private final Optional<AutomaticFunctionCallingResponder> automaticFunctionCallingResponder;
-  private List<Content> history = new ArrayList<>();
-  private int previousHistorySize = 0;
+  private List<Content> history;
+  private int previousHistorySize;
   private Optional<ResponseStream<GenerateContentResponse>> currentResponseStream;
   private Optional<GenerateContentResponse> currentResponse;
 
@@ -50,14 +55,17 @@ public final class ChatSession {
    * GenerationConfig) inherits from the model.
    */
   public ChatSession(GenerativeModel model) {
-    this(model, Optional.empty(), Optional.empty());
+    this(model, new ArrayList<>(), 0, Optional.empty(), Optional.empty());
   }
 
   /**
    * Creates a new chat session given a GenerativeModel instance and a root chat session.
    * Configurations of the chat (e.g., GenerationConfig) inherits from the model.
    *
    * @param model a {@link GenerativeModel} instance that generates contents in the chat.
+   * @param history a list of {@link Content} containing interleaving conversation between "user"
+   *     and "model".
+   * @param previousHistorySize the size of the previous history.
    * @param rootChatSession a root {@link ChatSession} instance. All the chat history in the current
    *     chat session will be merged to the root chat session.
    * @param automaticFunctionCallingResponder an {@link AutomaticFunctionCallingResponder} instance
@@ -66,10 +74,14 @@ public ChatSession(GenerativeModel model) {
    */
   private ChatSession(
       GenerativeModel model,
+      List<Content> history,
+      int previousHistorySize,
       Optional<ChatSession> rootChatSession,
       Optional<AutomaticFunctionCallingResponder> automaticFunctionCallingResponder) {
     checkNotNull(model, "model should not be null");
     this.model = model;
+    this.history = history;
+    this.previousHistorySize = previousHistorySize;
     this.rootChatSession = rootChatSession;
     this.automaticFunctionCallingResponder = automaticFunctionCallingResponder;
     currentResponseStream = Optional.empty();
@@ -84,15 +96,12 @@ private ChatSession(
    * @return a new {@link ChatSession} instance with the specified GenerationConfig.
    */
   public ChatSession withGenerationConfig(GenerationConfig generationConfig) {
-    ChatSession rootChat = rootChatSession.orElse(this);
-    ChatSession newChatSession =
-        new ChatSession(
-            model.withGenerationConfig(generationConfig),
-            Optional.of(rootChat),
-            automaticFunctionCallingResponder);
-    newChatSession.history = history;
-    newChatSession.previousHistorySize = previousHistorySize;
-    return newChatSession;
+    return new ChatSession(
+        model.withGenerationConfig(generationConfig),
+        history,
+        previousHistorySize,
+        Optional.of(rootChatSession.orElse(this)),
+        automaticFunctionCallingResponder);
   }
 
   /**
@@ -103,15 +112,12 @@ public ChatSession withGenerationConfig(GenerationConfig generationConfig) {
    * @return a new {@link ChatSession} instance with the specified SafetySettings.
    */
   public ChatSession withSafetySettings(List<SafetySetting> safetySettings) {
-    ChatSession rootChat = rootChatSession.orElse(this);
-    ChatSession newChatSession =
-        new ChatSession(
-            model.withSafetySettings(safetySettings),
-            Optional.of(rootChat),
-            automaticFunctionCallingResponder);
-    newChatSession.history = history;
-    newChatSession.previousHistorySize = previousHistorySize;
-    return newChatSession;
+    return new ChatSession(
+        model.withSafetySettings(safetySettings),
+        history,
+        previousHistorySize,
+        Optional.of(rootChatSession.orElse(this)),
+        automaticFunctionCallingResponder);
   }
 
   /**
@@ -122,13 +128,44 @@ public ChatSession withSafetySettings(List<SafetySetting> safetySettings) {
    * @return a new {@link ChatSession} instance with the specified Tools.
    */
   public ChatSession withTools(List<Tool> tools) {
-    ChatSession rootChat = rootChatSession.orElse(this);
-    ChatSession newChatSession =
-        new ChatSession(
-            model.withTools(tools), Optional.of(rootChat), automaticFunctionCallingResponder);
-    newChatSession.history = history;
-    newChatSession.previousHistorySize = previousHistorySize;
-    return newChatSession;
+    return new ChatSession(
+        model.withTools(tools),
+        history,
+        previousHistorySize,
+        Optional.of(rootChatSession.orElse(this)),
+        automaticFunctionCallingResponder);
+  }
+
+  /**
+   * Creates a copy of the current ChatSession with updated ToolConfig.
+   *
+   * @param toolConfig a {@link com.google.cloud.vertexai.api.ToolConfig} that will be used in the
+   *     new ChatSession.
+   * @return a new {@link ChatSession} instance with the specified ToolConfigs.
+   */
+  public ChatSession withToolConfig(ToolConfig toolConfig) {
+    return new ChatSession(
+        model.withToolConfig(toolConfig),
+        history,
+        previousHistorySize,
+        Optional.of(rootChatSession.orElse(this)),
+        automaticFunctionCallingResponder);
+  }
+
+  /**
+   * Creates a copy of the current ChatSession with updated SystemInstruction.
+   *
+   * @param systemInstruction a {@link com.google.cloud.vertexai.api.Content} containing system
+   *     instructions.
+   * @return a new {@link ChatSession} instance with the specified ToolConfigs.
+   */
+  public ChatSession withSystemInstruction(Content systemInstruction) {
+    return new ChatSession(
+        model.withSystemInstruction(systemInstruction),
+        history,
+        previousHistorySize,
+        Optional.of(rootChatSession.orElse(this)),
+        automaticFunctionCallingResponder);
   }
 
   /**
@@ -141,13 +178,12 @@ public ChatSession withTools(List<Tool> tools) {
    */
   public ChatSession withAutomaticFunctionCallingResponder(
       AutomaticFunctionCallingResponder automaticFunctionCallingResponder) {
-    ChatSession rootChat = rootChatSession.orElse(this);
-    ChatSession newChatSession =
-        new ChatSession(
-            model, Optional.of(rootChat), Optional.of(automaticFunctionCallingResponder));
-    newChatSession.history = history;
-    newChatSession.previousHistorySize = previousHistorySize;
-    return newChatSession;
+    return new ChatSession(
+        model,
+        history,
+        previousHistorySize,
+        Optional.of(rootChatSession.orElse(this)),
+        Optional.of(automaticFunctionCallingResponder));
   }
 
   /**

@@ -39,7 +39,13 @@
 import java.util.List;
 import java.util.Optional;
 
-/** This class holds a generative model that can complete what you provided. */
+/**
+ * This class holds a generative model that can complete what you provided. This class is
+ * thread-safe.
+ *
+ * <p>Note: The instances of {@link ChatSession} returned by {@link GenerativeModel#startChat()} are
+ * NOT thread-safe.
+ */
 public final class GenerativeModel {
   private final String modelName;
   private final String resourceName;
@@ -645,6 +651,11 @@ public Optional<Content> getSystemInstruction() {
     return systemInstruction;
   }
 
+  /**
+   * Returns a new {@link ChatSession} instance that can be used to start a chat with this model.
+   *
+   * <p>Note: the returned {@link ChatSession} instance is NOT thread-safe.
+   */
   public ChatSession startChat() {
     return new ChatSession(this);
   }