feat(agentplatform): Support ExampleStores module for few-shot example management

agentplatform/_genai/_transformers.py

@@ -576,28 +576,52 @@ def t_endpoint(endpoint: str) -> str:
 
 
 def t_strict_endpoint(endpoint: str) -> str:
-    """Validates a name that must address an Endpoint resource.
+  """Validates a name that must address an Endpoint resource.
 
     A publisher model is served by the shared prediction endpoint and has no
     Endpoint resource behind it, so it can be predicted against but never
     fetched, undeployed or deleted. Methods that address the resource itself
     must reject it here rather than issue a request against a publishers/... URL.
     """
-    if not endpoint:
-        raise ValueError("endpoint is required.")
+  if not endpoint:
+    raise ValueError("endpoint is required.")
 
-    if is_endpoint_resource_name(endpoint):
-        return endpoint
+  if is_endpoint_resource_name(endpoint):
+    return endpoint
 
-    if _is_publisher_model_resource_name(endpoint):
-        raise ValueError(
+  if _is_publisher_model_resource_name(endpoint):
+    raise ValueError(
             f"{endpoint} is a publisher model, which does not support this"
             " method. Only endpoints in the format of"
             " projects/.../locations/.../endpoints/... or endpoints/... are"
             " supported."
         )
 
-    raise ValueError(
+  raise ValueError(
         f"Invalid endpoint format: {endpoint}. Must be in the format of"
         " projects/.../locations/.../endpoints/... or endpoints/..."
     )
+
+_EXAMPLE_STORE_RES_NAME_RES = (
+    re.compile(r"^projects/[^/]+/locations/[^/]+/exampleStores/[^/]+$"),
+    re.compile(r"^exampleStores/[^/]+$"),
+)
+
+
+def is_example_store_resource_name(name: str) -> bool:
+  """Returns whether the name addresses an ExampleStore resource."""
+  return any(pattern.match(name) for pattern in _EXAMPLE_STORE_RES_NAME_RES)
+
+
+def t_example_store(example_store: str) -> str:
+  """Validates a name that must address an ExampleStore resource."""
+  if not example_store:
+    raise ValueError("example_store is required.")
+
+  if is_example_store_resource_name(example_store):
+    return example_store
+
+  raise ValueError(
+      f"Invalid example store format: {example_store}. Must be in the format"
+      " of projects/.../locations/.../exampleStores/... or exampleStores/..."
+  )

agentplatform/_genai/client.py

@@ -50,6 +50,9 @@
     from agentplatform._genai import (
         endpoints as endpoints_module,
     )
+    from agentplatform._genai import (
+        example_stores as example_stores_module,
+    )
 
 _GENAI_MODULES_TELEMETRY_HEADER = "vertex-genai-modules"
 
@@ -96,6 +99,7 @@ def __init__(self, api_client: genai_client.BaseApiClient):  # type: ignore[name
         self._model_garden: Optional[ModuleType] = None
         self._feedback_entries: Optional[ModuleType] = None
         self._endpoints: Optional[ModuleType] = None
+        self._example_stores: Optional[ModuleType] = None
 
     @property
     @_common.experimental_warning(
@@ -195,6 +199,15 @@ def endpoints(self) -> "endpoints_module.AsyncEndpoints":
             )
         return self._endpoints.AsyncEndpoints(self._api_client)  # type: ignore[no-any-return]
 
+    @property
+    def example_stores(self) -> "example_stores_module.AsyncExampleStores":
+        if self._example_stores is None:
+            self._example_stores = importlib.import_module(
+                ".example_stores",
+                __package__,
+            )
+        return self._example_stores.AsyncExampleStores(self._api_client)  # type: ignore[no-any-return]
+
     @property
     @_common.experimental_warning(
         "The Vertex SDK GenAI async rag module is experimental, "
@@ -328,6 +341,7 @@ def __init__(
         self._model_garden: Optional[ModuleType] = None
         self._feedback_entries: Optional[ModuleType] = None
         self._endpoints: Optional[ModuleType] = None
+        self._example_stores: Optional[ModuleType] = None
 
     @property
     def evals(self) -> "evals_module.Evals":
@@ -452,6 +466,15 @@ def endpoints(self) -> "endpoints_module.Endpoints":
             )
         return self._endpoints.Endpoints(self._api_client)  # type: ignore[no-any-return]
 
+    @property
+    def example_stores(self) -> "example_stores_module.ExampleStores":
+        if self._example_stores is None:
+            self._example_stores = importlib.import_module(
+                ".example_stores",
+                __package__,
+            )
+        return self._example_stores.ExampleStores(self._api_client)  # type: ignore[no-any-return]
+
     @property
     @_common.experimental_warning(
         "The Vertex SDK GenAI rag module is experimental, "