feat: GenAI SDK client - add rag.create_corpus() method

agentplatform/_genai/_operations_utils.py

@@ -0,0 +1,94 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Utility functions for Operations."""
+
+import asyncio
+import datetime
+import time
+from typing import Any, Awaitable, Callable
+
+
+def await_operation(
+    *,
+    operation_name: str,
+    get_operation_fn: Callable[..., Any],
+    poll_interval: datetime.timedelta | float = 10.0,
+    timeout_seconds: float = 300.0,
+) -> Any:
+    """Waits for a long running operation to complete.
+
+    Args:
+        operation_name (str): Required. The name of the operation.
+        get_operation_fn (Callable): Required. Function to get the operation
+          status.
+        poll_interval (datetime.timedelta | float): The interval between polls.
+        timeout_seconds (float): The maximum wait duration in seconds.
+
+    Returns:
+        Any: The completed operation.
+    """
+    if isinstance(poll_interval, datetime.timedelta):
+        poll_seconds = poll_interval.total_seconds()
+    else:
+        poll_seconds = float(poll_interval)
+
+    start_time = time.time()
+    operation = get_operation_fn(operation_name=operation_name)
+    while not operation.done:
+        if (time.time() - start_time) > timeout_seconds:
+            raise TimeoutError(
+                f"Operation {operation_name} did not complete within the timeout "
+                f"of {timeout_seconds} seconds."
+            )
+        time.sleep(poll_seconds)
+        operation = get_operation_fn(operation_name=operation.name)
+    return operation
+
+
+async def await_operation_async(
+    *,
+    operation_name: str,
+    get_operation_fn: Callable[..., Awaitable[Any]],
+    poll_interval: datetime.timedelta | float = 10.0,
+    timeout_seconds: float = 300.0,
+) -> Any:
+    """Waits for a long running operation to complete asynchronously.
+
+    Args:
+        operation_name (str): Required. The name of the operation.
+        get_operation_fn (Callable): Required. Async function to get the operation
+          status.
+        poll_interval (datetime.timedelta | float): The interval between polls.
+        timeout_seconds (float): The maximum wait duration in seconds.
+
+    Returns:
+        Any: The completed operation.
+    """
+    if isinstance(poll_interval, datetime.timedelta):
+        poll_seconds = poll_interval.total_seconds()
+    else:
+        poll_seconds = float(poll_interval)
+
+    start_time = time.time()
+    operation = await get_operation_fn(operation_name=operation_name)
+    while not operation.done:
+        if (time.time() - start_time) > timeout_seconds:
+            raise TimeoutError(
+                f"Operation {operation_name} did not complete within the timeout "
+                f"of {timeout_seconds} seconds."
+            )
+        await asyncio.sleep(poll_seconds)
+        operation = await get_operation_fn(operation_name=operation.name)
+    return operation

agentplatform/_genai/_skills_utils.py

@@ -14,14 +14,10 @@
 #
 """Utility functions for Skills."""
 
-import asyncio
 import base64
-import datetime
 import io
 import os
 import pathlib
-import time
-from typing import Any, Awaitable, Callable
 import zipfile
 
 
@@ -71,77 +67,3 @@ def get_zipped_filesystem_payload(directory_path: pathlib.Path | str) -> str:
     """
     zip_bytes = zip_directory(directory_path)
     return base64.b64encode(zip_bytes).decode("utf-8")
-
-
-def await_operation(
-    *,
-    operation_name: str,
-    get_operation_fn: Callable[..., Any],
-    poll_interval: datetime.timedelta | float = 10.0,
-    timeout_seconds: float = 300.0,
-) -> Any:
-    """Waits for a long running operation to complete.
-
-    Args:
-        operation_name (str): Required. The name of the operation.
-        get_operation_fn (Callable): Required. Function to get the operation
-          status.
-        poll_interval (datetime.timedelta | float): The interval between polls.
-        timeout_seconds (float): The maximum wait duration in seconds.
-
-    Returns:
-        Any: The completed operation.
-    """
-    if isinstance(poll_interval, datetime.timedelta):
-        poll_seconds = poll_interval.total_seconds()
-    else:
-        poll_seconds = float(poll_interval)
-
-    start_time = time.time()
-    operation = get_operation_fn(operation_name=operation_name)
-    while not operation.done:
-        if (time.time() - start_time) > timeout_seconds:
-            raise TimeoutError(
-                f"Operation {operation_name} did not complete within the timeout "
-                f"of {timeout_seconds} seconds."
-            )
-        time.sleep(poll_seconds)
-        operation = get_operation_fn(operation_name=operation.name)
-    return operation
-
-
-async def await_operation_async(
-    *,
-    operation_name: str,
-    get_operation_fn: Callable[..., Awaitable[Any]],
-    poll_interval: datetime.timedelta | float = 10.0,
-    timeout_seconds: float = 300.0,
-) -> Any:
-    """Waits for a long running operation to complete asynchronously.
-
-    Args:
-        operation_name (str): Required. The name of the operation.
-        get_operation_fn (Callable): Required. Async function to get the operation
-          status.
-        poll_interval (datetime.timedelta | float): The interval between polls.
-        timeout_seconds (float): The maximum wait duration in seconds.
-
-    Returns:
-        Any: The completed operation.
-    """
-    if isinstance(poll_interval, datetime.timedelta):
-        poll_seconds = poll_interval.total_seconds()
-    else:
-        poll_seconds = float(poll_interval)
-
-    start_time = time.time()
-    operation = await get_operation_fn(operation_name=operation_name)
-    while not operation.done:
-        if (time.time() - start_time) > timeout_seconds:
-            raise TimeoutError(
-                f"Operation {operation_name} did not complete within the timeout "
-                f"of {timeout_seconds} seconds."
-            )
-        await asyncio.sleep(poll_seconds)
-        operation = await get_operation_fn(operation_name=operation.name)
-    return operation