feat: GenAI - Add model routing config to sdk

vertex-sdk-bot · copybara-github · commit c0626fe015e4 · 2024-09-19T06:40:05.000-07:00
PiperOrigin-RevId: 676391359
diff --git a/tests/system/vertexai/test_generative_models.py b/tests/system/vertexai/test_generative_models.py
@@ -36,6 +36,7 @@
 GEMINI_VISION_MODEL_NAME = "gemini-1.0-pro-vision"
 GEMINI_15_MODEL_NAME = "gemini-1.5-pro-preview-0409"
 GEMINI_15_PRO_MODEL_NAME = "gemini-1.5-pro-001"
+SMART_ROUTER_NAME = "smart-router-001"
 
 STAGING_API_ENDPOINT = "STAGING_ENDPOINT"
 PROD_API_ENDPOINT = "PROD_ENDPOINT"
@@ -494,6 +495,21 @@ def test_generate_content_function_calling(self, api_endpoint_env_name):
 
         assert summary
 
+    def test_generate_content_model_router(self, gapi_endpoint_env_name):
+        model = generative_models.GenerativeModel(SMART_ROUTER_NAME)
+        response = model.generate_content(
+            contents="Why is sky blue?",
+            generation_config=generative_models.GenerationConfig(
+                temperature=0,
+                routing_config=generative_models.GenerationConfig.RoutingConfig(
+                    routing_config=generative_models.GenerationConfig.RoutingConfig.AutoRoutingMode(
+                        model_routing_preference=generative_models.GenerationConfig.RoutingConfig.AutoRoutingMode.ModelRoutingPreference.BALANCED,
+                    ),
+                ),
+            ),
+        )
+        assert response.text
+
     def test_chat_automatic_function_calling(self, api_endpoint_env_name):
         get_current_weather_func = generative_models.FunctionDeclaration.from_func(
             get_current_weather
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -1576,6 +1576,7 @@ def __init__(
         response_mime_type: Optional[str] = None,
         response_schema: Optional[Dict[str, Any]] = None,
         seed: Optional[int] = None,
+        routing_config: Optional["RoutingConfig"] = None,
     ):
         r"""Constructs a GenerationConfig object.
 
@@ -1601,6 +1602,7 @@ def __init__(
                 response type, otherwise the behavior is undefined.
             response_schema: Output response schema of the genreated candidate text. Only valid when
                 response_mime_type is application/json.
+            routing_config: Model routing preference set in the request.
 
         Usage:
             ```
@@ -1636,6 +1638,10 @@ def __init__(
             response_schema=raw_schema,
             seed=seed,
         )
+        if routing_config is not None:
+            self._raw_generation_config.routing_config = (
+                routing_config._gapic_routing_config
+            )
 
     @classmethod
     def _from_gapic(
@@ -1659,6 +1665,109 @@ def to_dict(self) -> Dict[str, Any]:
     def __repr__(self) -> str:
         return self._raw_generation_config.__repr__()
 
+    class RoutingConfig:
+        r"""The configuration for model router requests.
+
+        The routing config is either one of the two nested classes:
+        - AutoRoutingMode: Automated routing.
+        - ManualRoutingMode: Manual routing.
+
+        Usage:
+        - AutoRoutingMode:
+
+            ```
+            routing_config=generative_models.RoutingConfig(
+                routing_config=generative_models.RoutingConfig.AutoRoutingMode(
+                    model_routing_preference=generative_models.RoutingConfig.AutoRoutingMode.ModelRoutingPreference.BALANCED,
+                ),
+            )
+            ```
+        - ManualRoutingMode:
+
+            ```
+            routing_config=generative_models.RoutingConfig(
+                routing_config=generative_models.RoutingConfig.ManutalRoutingMode(
+                    model_name="gemini-1.5-pro-001",
+                ),
+            )
+            ```
+        """
+
+        def __init__(
+            self,
+            *,
+            routing_config: Union[
+                "GenerationConfig.RoutingConfig.AutoRoutingMode",
+                "GenerationConfig.RoutingConfig.ManualRoutingMode",
+            ],
+        ):
+            if isinstance(routing_config, self.AutoRoutingMode):
+                self._gapic_routing_config = (
+                    gapic_content_types.GenerationConfig.RoutingConfig(
+                        auto_mode=routing_config._gapic_auto_mode
+                    )
+                )
+            else:
+                self._gapic_routing_config = (
+                    gapic_content_types.GenerationConfig.RoutingConfig(
+                        manual_mode=routing_config._gapic_manual_mode
+                    )
+                )
+
+        def __repr__(self):
+            return self._gapic_routing_config.__repr__()
+
+        class AutoRoutingMode:
+            r"""When automated routing is specified, the routing will be
+            determined by the routing model predicted quality and customer provided
+            model routing preference.
+            """
+
+            ModelRoutingPreference = (
+                gapic_content_types.GenerationConfig.RoutingConfig.AutoRoutingMode.ModelRoutingPreference
+            )
+
+            def __init__(
+                self,
+                *,
+                model_routing_preference: "GenerationConfig.RoutingConfig.AutoRoutingMode.ModelRoutingPreference",
+            ):
+                r"""AutoRouingMode constructor
+
+                Args:
+                    model_routing_preference: Model routing preference for the routing request.
+                """
+                self._gapic_auto_mode = (
+                    gapic_content_types.GenerationConfig.RoutingConfig.AutoRoutingMode(
+                        model_routing_preference=model_routing_preference
+                    )
+                )
+
+            def __repr__(self):
+                return self._gapic_auto_mode.__repr__()
+
+        class ManualRoutingMode:
+            r"""When manual routing is set, the specified model will be used
+            directly.
+            """
+
+            def __init__(
+                self,
+                *,
+                model_name: str,
+            ):
+                r"""ManualRoutingMode constructor
+
+                Args:
+                    model_name: The model to use. Only public LLM model names and those that are supported by the router are allowed.
+                """
+                self._gapic_manual_mode = gapic_content_types.GenerationConfig.RoutingConfig.ManualRoutingMode(
+                    model_name=model_name
+                )
+
+            def __repr__(self):
+                return self._gapic_manual_mode.__repr__()
+
 
 class Tool:
     r"""A collection of functions that the model may use to generate response.
