[Ferature Request] Support Google ModelGarden Models #567

README.md

@@ -7,7 +7,7 @@
 
 -----
 
-Google Gen AI Python SDK provides an interface for developers to integrate Google's generative models into their Python applications. It supports the [Gemini Developer API](https://ai.google.dev/gemini-api/docs) and [Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/overview) APIs.
+Google Gen AI Python SDK provides an interface for developers to integrate Google's generative models into their Python applications. It supports the [Gemini Developer API](https://ai.google.dev/gemini-api/docs), [Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/overview) APIs, and [ModelGarden](https://cloud.google.com/vertex-ai/docs/model-garden/get-started) models.
 
 ## Installation
 
@@ -25,7 +25,7 @@ from google.genai import types
 ## Create a client
 
 Please run one of the following code blocks to create a client for
-different services ([Gemini Developer API](https://ai.google.dev/gemini-api/docs) or [Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/overview)).
+different services ([Gemini Developer API](https://ai.google.dev/gemini-api/docs), [Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/overview), or [ModelGarden](https://cloud.google.com/vertex-ai/docs/model-garden/get-started)).
 
 ```python
 # Only run this block for Gemini Developer API
@@ -39,11 +39,18 @@ client = genai.Client(
 )
 ```
 
+```python
+# Only run this block for ModelGarden models
+client = genai.Client(
+    modelgarden=True, project='your-project-id', location='us-central1'
+)
+```
+
 **(Optional) Using environment variables:**
 
 You can create a client by configuring the necessary environment variables.
 Configuration setup instructions depends on whether you're using the Gemini
-Developer API or the Gemini API in Vertex AI.
+Developer API, the Gemini API in Vertex AI, or ModelGarden models.
 
 **Gemini Developer API:** Set `GOOGLE_API_KEY` as shown below:
 
@@ -60,6 +67,15 @@ export GOOGLE_CLOUD_PROJECT='your-project-id'
 export GOOGLE_CLOUD_LOCATION='us-central1'
 ```
 
+**ModelGarden models:** Set `GOOGLE_GENAI_USE_MODELGARDEN`, `GOOGLE_CLOUD_PROJECT`
+and `GOOGLE_CLOUD_LOCATION`, as shown below:
+
+```bash
+export GOOGLE_GENAI_USE_MODELGARDEN=true
+export GOOGLE_CLOUD_PROJECT='your-project-id'
+export GOOGLE_CLOUD_LOCATION='us-central1'
+```
+
 ```python
 client = genai.Client()
 ```

examples/modelgarden_example.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""
+Example demonstrating how to use Google ModelGarden models with the GenAI SDK.
+
+This example shows how to initialize the client with ModelGarden support and
+make requests to ModelGarden models.
+"""
+
+import os
+from google import genai
+from google.genai import types
+import argparse
+
+def main():
+    """Main function demonstrating ModelGarden usage"""
+    parser = argparse.ArgumentParser(description='Use Google ModelGarden models with the GenAI SDK')
+    parser.add_argument('--project', type=str, help='Google Cloud project ID')
+    parser.add_argument('--location', type=str, default='us-central1', 
+                        help='Google Cloud location (default: us-central1)')
+    parser.add_argument('--model', type=str, default='publishers/meta/models/llama3-8b-instruct',
+                       help='ModelGarden model to use (default: publishers/meta/models/llama3-8b-instruct)')
+    parser.add_argument('--prompt', type=str, default='Explain quantum computing in simple terms.',
+                       help='Prompt to send to the model')
+
+    args = parser.parse_args()
+
+    # Get project from args or environment variable
+    project = args.project or os.environ.get('GOOGLE_CLOUD_PROJECT')
+    if not project:
+        print("Error: Google Cloud project ID is required.")
+        print("Please provide it via --project argument or GOOGLE_CLOUD_PROJECT environment variable.")
+        return
+
+    # Get location from args or environment variable
+    location = args.location or os.environ.get('GOOGLE_CLOUD_LOCATION', 'us-central1')
+
+    # Create client with ModelGarden support
+    print(f"Initializing client with ModelGarden support (project: {project}, location: {location})")
+    client = genai.Client(
+        modelgarden=True,
+        project=project,
+        location=location
+    )
+
+    model = args.model
+    print(f"Using ModelGarden model: {model}")
+
+    # Create the prompt
+    prompt = args.prompt
+    print(f"Prompt: {prompt}")
+
+    # Generate content using the ModelGarden model
+    print("\nGenerating response...")
+    try:
+        response = client.models.generate_content(
+            model=model,
+            contents=prompt,
+            config=types.GenerateContentConfig(
+                temperature=0.2,
+                max_output_tokens=800,
+            )
+        )
+
+        print("\nResponse:")
+        print(response.text)
+
+    except Exception as e:
+        print(f"Error: {e}")
+        print("\nNote: Make sure you have access to the specified ModelGarden model and that")
+        print("your Google Cloud credentials have the necessary permissions.")
+
+if __name__ == "__main__":
+    main() 

google/genai/_api_client.py

@@ -291,24 +291,38 @@ def __del__(self) -> None:
 
 
 class BaseApiClient:
-  """Client for calling HTTP APIs sending and receiving JSON."""
+  """Base client for calling HTTP APIs sending and receiving JSON."""
 
   def __init__(
       self,
       vertexai: Optional[bool] = None,
+      modelgarden: Optional[bool] = None,
       api_key: Optional[str] = None,
       credentials: Optional[google.auth.credentials.Credentials] = None,
       project: Optional[str] = None,
       location: Optional[str] = None,
       http_options: Optional[HttpOptionsOrDict] = None,
   ):
     self.vertexai = vertexai
+    self.modelgarden = modelgarden
+
     if self.vertexai is None:
       if os.environ.get('GOOGLE_GENAI_USE_VERTEXAI', '0').lower() in [
           'true',
           '1',
       ]:
         self.vertexai = True
+
+    if self.modelgarden is None:
+      if os.environ.get('GOOGLE_GENAI_USE_MODELGARDEN', '0').lower() in [
+          'true',
+          '1',
+      ]:
+        self.modelgarden = True
+
+    # If modelgarden is True, then vertexai must also be True
+    if self.modelgarden:
+      self.vertexai = True
 
     # Validate explicitly set initializer values.
     if (project or location) and api_key:
@@ -400,7 +414,16 @@ def __init__(
         self._http_options.base_url = (
             f'https://{self.location}-aiplatform.googleapis.com/'
         )
-      self._http_options.api_version = 'v1beta1'
+
+      # If using modelgarden, use v1 API version
+      if self.modelgarden:
+        if not (self.project and self.location):
+          raise ValueError(
+              'Project and location must be set when using ModelGarden models.'
+          )
+        self._http_options.api_version = 'v1'
+      else:
+        self._http_options.api_version = 'v1beta1'
     else:  # Implicit initialization or missing arguments.
       if not self.api_key:
         raise ValueError(

google/genai/_replay_api_client.py

@@ -191,6 +191,7 @@ def __init__(
       replay_id: str,
       replays_directory: Optional[str] = None,
       vertexai: bool = False,
+      modelgarden: bool = False,
       api_key: Optional[str] = None,
       credentials: Optional[google.auth.credentials.Credentials] = None,
       project: Optional[str] = None,
@@ -199,6 +200,7 @@ def __init__(
   ):
     super().__init__(
         vertexai=vertexai,
+        modelgarden=modelgarden,
         api_key=api_key,
         credentials=credentials,
         project=project,

google/genai/client.py

@@ -110,12 +110,17 @@ class Client:
   `GOOGLE_GENAI_USE_VERTEXAI=true`, `GOOGLE_CLOUD_PROJECT` and
   `GOOGLE_CLOUD_LOCATION` environment variables.
 
+  ModelGarden users can provide `modelgarden=True` along with Vertex AI parameters
+  or by defining `GOOGLE_GENAI_USE_MODELGARDEN=true` as an environment variable.
+
   Attributes:
     api_key: The `API key <https://ai.google.dev/gemini-api/docs/api-key>`_ to
       use for authentication. Applies to the Gemini Developer API only.
     vertexai: Indicates whether the client should use the Vertex AI
       API endpoints. Defaults to False (uses Gemini Developer API endpoints).
       Applies to the Vertex AI API only.
+    modelgarden: Indicates whether the client should use ModelGarden models.
+      If True, vertexai is also set to True. Defaults to False.
     credentials: The credentials to use for authentication when calling the
       Vertex AI APIs. Credentials can be obtained from environment variables and
       default credentials. For more information, see
@@ -151,12 +156,23 @@ class Client:
     client = genai.Client(
         vertexai=True, project='my-project-id', location='us-central1'
     )
+
+  Usage for ModelGarden models:
+
+  .. code-block:: python
+
+    from google import genai
+
+    client = genai.Client(
+        modelgarden=True, project='my-project-id', location='us-central1'
+    )
   """
 
   def __init__(
       self,
       *,
       vertexai: Optional[bool] = None,
+      modelgarden: Optional[bool] = None,
       api_key: Optional[str] = None,
       credentials: Optional[google.auth.credentials.Credentials] = None,
       project: Optional[str] = None,
@@ -170,6 +186,8 @@ def __init__(
        vertexai (bool): Indicates whether the client should use the Vertex AI
          API endpoints. Defaults to False (uses Gemini Developer API endpoints).
          Applies to the Vertex AI API only.
+       modelgarden (bool): Indicates whether the client should use ModelGarden models.
+         If True, vertexai is also set to True. Defaults to False.
        api_key (str): The `API key
          <https://ai.google.dev/gemini-api/docs/api-key>`_ to use for
          authentication. Applies to the Gemini Developer API only.
@@ -199,6 +217,7 @@ def __init__(
 
     self._api_client = self._get_api_client(
         vertexai=vertexai,
+        modelgarden=modelgarden,
         api_key=api_key,
         credentials=credentials,
         project=project,
@@ -218,6 +237,7 @@ def __init__(
   @staticmethod
   def _get_api_client(
       vertexai: Optional[bool] = None,
+      modelgarden: Optional[bool] = None,
       api_key: Optional[str] = None,
       credentials: Optional[google.auth.credentials.Credentials] = None,
       project: Optional[str] = None,
@@ -235,6 +255,7 @@ def _get_api_client(
           replay_id=debug_config.replay_id,  # type: ignore[arg-type]
           replays_directory=debug_config.replays_directory,
           vertexai=vertexai,  # type: ignore[arg-type]
+          modelgarden=modelgarden,  # type: ignore[arg-type]
           api_key=api_key,
           credentials=credentials,
           project=project,
@@ -244,6 +265,7 @@ def _get_api_client(
 
     return BaseApiClient(
         vertexai=vertexai,
+        modelgarden=modelgarden,
         api_key=api_key,
         credentials=credentials,
         project=project,
@@ -287,3 +309,8 @@ def operations(self) -> Operations:
   def vertexai(self) -> bool:
     """Returns whether the client is using the Vertex AI API."""
     return self._api_client.vertexai or False
+
+  @property
+  def modelgarden(self) -> bool:
+    """Returns whether the client is using the ModelGarden API."""
+    return getattr(self._api_client, 'modelgarden', False) or False