google-cloud-agentplatform 1.165.1.dev0__py3-none-any.whl

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/_prompt_management_utils.py

@@ -0,0 +1,147 @@
+# Copyright 2025 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 prompt management."""
+
+from typing import Optional
+
+from google.genai import types as genai_types
+
+from . import types
+
+
+DEFAULT_API_SCHEMA_VERSION = "1.0.0"
+PROMPT_SCHEMA_URI = (
+    "gs://google-cloud-aiplatform/schema/dataset/metadata/text_prompt_1.0.0.yaml"
+)
+PROMPT_TYPE = "multimodal_freeform"
+
+
+def _create_dataset_metadata_from_prompt(
+    prompt: types.Prompt,
+    variables: Optional[list[dict[str, genai_types.Part]]] = None,
+) -> types.SchemaTextPromptDatasetMetadata:
+    """Convert a types.Prompt into types.SchemaTextPromptDatasetMetadata."""
+
+    prompt_metadata = types.SchemaTextPromptDatasetMetadata()
+
+    prompt_api_schema = types.SchemaPromptApiSchema()
+    prompt_api_schema.multimodal_prompt = types.SchemaPromptSpecMultimodalPrompt(
+        prompt_message=prompt.prompt_data
+    )
+
+    prompt_api_schema.api_schema_version = DEFAULT_API_SCHEMA_VERSION
+
+    prompt_metadata.has_prompt_variable = bool(variables)
+
+    if variables:
+        prompt_execution_list = []
+        for prompt_var in variables:
+            prompt_instance_execution = types.SchemaPromptInstancePromptExecution()
+            prompt_instance_execution.arguments = {}
+            for key, val in prompt_var.items():
+                prompt_instance_execution.arguments[key] = (
+                    types.SchemaPromptInstanceVariableValue(
+                        part_list=types.SchemaPromptSpecPartList(parts=[val])
+                    )
+                )
+            prompt_execution_list.append(prompt_instance_execution)
+        prompt_api_schema.executions = prompt_execution_list
+
+    # Need to exclude variables from the prompt message as it is a client side
+    # only field
+    if prompt_api_schema.multimodal_prompt.prompt_message:
+        prompt_message_dict = (
+            prompt_api_schema.multimodal_prompt.prompt_message.model_dump(
+                exclude=["variables"], exclude_none=True
+            )
+        )
+        prompt_api_schema.multimodal_prompt.prompt_message = (
+            types.SchemaPromptSpecPromptMessage(**prompt_message_dict)
+        )
+    prompt_metadata.prompt_api_schema = prompt_api_schema
+
+    prompt_metadata.prompt_type = PROMPT_TYPE
+
+    return prompt_metadata
+
+
+def _create_prompt_from_dataset_metadata(
+    dataset: types.Dataset,
+) -> types.Prompt:
+    """Constructs a types.Prompt from a types.Dataset resource returned from the API.
+
+    Args:
+      dataset: The types.Dataset object containing the prompt metadata.
+
+    Returns:
+      A types.Prompt object reconstructed from the dataset metadata.
+    """
+    if (
+        not hasattr(dataset, "metadata")
+        or dataset.metadata is None
+        or not isinstance(dataset.metadata, types.SchemaTextPromptDatasetMetadata)
+    ):
+        raise ValueError(
+            "Error retrieving prompt: prompt dataset resource is missing 'metadata'."
+        )
+    api_schema = dataset.metadata.prompt_api_schema
+    prompt = types.Prompt()
+
+    if api_schema is None:
+        return prompt
+
+    if api_schema.multimodal_prompt:
+
+        prompt_message = api_schema.multimodal_prompt.prompt_message
+        prompt.prompt_data = prompt_message
+
+        if api_schema.executions:
+            executions = api_schema.executions
+            if executions and prompt.prompt_data is not None:
+                prompt.prompt_data.variables = []
+                for execution in executions:
+                    if execution.arguments:
+                        args = execution.arguments
+                        var_map = {}
+                        for key, val in args.items():
+                            if (
+                                val.part_list is not None
+                                and val.part_list.parts is not None
+                            ):
+                                part_list = val.part_list.parts
+                                if part_list and part_list[0].text:
+                                    var_map[key] = part_list[0]
+                        if var_map and prompt.prompt_data.variables is not None:
+                            prompt.prompt_data.variables.append(var_map)
+
+    return prompt
+
+
+def _raise_for_invalid_prompt(
+    prompt: types.Prompt,
+) -> None:
+
+    if not prompt.prompt_data:
+        raise ValueError("Prompt data must be provided.")
+    if not prompt.prompt_data.contents:
+        raise ValueError("Prompt contents must be provided.")
+    if not prompt.prompt_data.model:
+        raise ValueError("Model name must be provided.")
+    if (
+        prompt.prompt_data
+        and prompt.prompt_data.contents
+        and len(prompt.prompt_data.contents) > 1
+    ):
+        raise ValueError("Multi-turn prompts are not currently supported.")

agentplatform/_genai/_prompt_optimizer_utils.py

@@ -0,0 +1,215 @@
+# Copyright 2025 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 prompt optimizer."""
+
+import json
+from typing import Any, Optional, Union
+from typing_extensions import TypeAlias
+
+from pydantic import ValidationError
+
+from . import types
+
+try:
+    import pandas as pd  # pylint: disable=g-import-not-at-top
+
+    PandasDataFrame: TypeAlias = pd.DataFrame
+except ImportError:
+    pd = None
+    PandasDataFrame = Any  # type: ignore[misc]
+
+
+def _construct_input_prompt(
+    example_df: PandasDataFrame,
+    *,
+    prompt_col_name: str,
+    model_response_col_name: str,
+    rubrics_col_name: str,
+    rubrics_evaluations_col_name: str,
+    target_response_col_name: str,
+    system_instruction: Optional[str] = None,
+) -> str:
+    """Construct the input prompt for the few shot prompt optimizer."""
+
+    all_prompts = []
+    for row in example_df.to_dict(orient="records"):
+        example_data = {
+            "prompt": row[prompt_col_name],
+            "model_response": row[model_response_col_name],
+        }
+        if rubrics_col_name:
+            example_data["rubrics"] = row[rubrics_col_name]
+        if rubrics_evaluations_col_name:
+            example_data["rubrics_evaluations"] = row[rubrics_evaluations_col_name]
+        if target_response_col_name:
+            example_data["target_response"] = row[target_response_col_name]
+
+        json_str = json.dumps(example_data, indent=2)
+        all_prompts.append(f"```JSON\n{json_str}\n```")
+
+    all_prompts_str = "\n\n".join(all_prompts)
+
+    if system_instruction is None:
+        system_instruction = ""
+
+    return "\n".join(
+        [
+            "Original System Instructions:\n",
+            system_instruction,
+            "Examples:\n",
+            all_prompts_str,
+            "\nNew Output:\n",
+        ]
+    )
+
+
+def _get_few_shot_prompt(
+    system_instruction: str,
+    config: types.OptimizeConfig,
+) -> str:
+    """Builds the few shot prompt."""
+
+    if config.examples_dataframe is None:
+        raise ValueError("The 'examples_dataframe' is required in the config.")
+
+    if "prompt" not in config.examples_dataframe.columns:
+        raise ValueError("'prompt' is required in the examples_dataframe.")
+
+    if "prompt" not in config.examples_dataframe.columns:
+        raise ValueError("'prompt' is required in the examples_dataframe.")
+    prompt_col_name = "prompt"
+
+    if "model_response" not in config.examples_dataframe.columns:
+        raise ValueError("'model_response' is required in the example_df.")
+    model_response_col_name = "model_response"
+
+    target_response_col_name = ""
+    rubrics_col_name = ""
+    rubrics_evaluations_col_name = ""
+
+    if (
+        config.optimization_target
+        == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE
+    ):
+        if "target_response" not in config.examples_dataframe.columns:
+            raise ValueError("'target_response' is required in the examples_dataframe.")
+        target_response_col_name = "target_response"
+        if "rubrics" in config.examples_dataframe.columns:
+            raise ValueError(
+                "Only 'target_response' should be provided "
+                "for OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE "
+                "but 'rubrics' was provided."
+            )
+
+    elif (
+        config.optimization_target
+        == types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS
+    ):
+        if not {"rubrics", "rubrics_evaluations"}.issubset(
+            config.examples_dataframe.columns
+        ):
+            raise ValueError(
+                "rubrics and rubrics_evaluations is required in the"
+                "examples_dataframe when rubrics is set."
+            )
+
+        rubrics_col_name = "rubrics"
+        rubrics_evaluations_col_name = "rubrics_evaluations"
+        if "target_response" in config.examples_dataframe.columns:
+            raise ValueError(
+                "Only 'rubrics' and 'rubrics_evaluations' should be provided "
+                "for OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS "
+                "but target_response was provided."
+            )
+    else:
+        raise ValueError("One of 'target_response' or 'rubrics' must be provided.")
+
+    return _construct_input_prompt(
+        config.examples_dataframe,
+        prompt_col_name=prompt_col_name,
+        model_response_col_name=model_response_col_name,
+        rubrics_col_name=rubrics_col_name,
+        rubrics_evaluations_col_name=rubrics_evaluations_col_name,
+        target_response_col_name=target_response_col_name,
+        system_instruction=system_instruction,
+    )
+
+
+def _get_service_account(
+    config: types.PromptOptimizerConfigOrDict,
+) -> str:
+    """Get the service account from the config for the custom job."""
+    if isinstance(config, dict):
+        config = types.PromptOptimizerConfig.model_validate(config)
+
+    if (
+        config.service_account and config.service_account_project_number
+    ):  # pytype: disable=attribute-error
+        raise ValueError(
+            "Only one of service_account or "
+            "service_account_project_number can be provided."
+        )
+    elif config.service_account:  # pytype: disable=attribute-error
+        return config.service_account  # pytype: disable=attribute-error
+    elif config.service_account_project_number:  # pytype: disable=attribute-error
+        return f"{config.service_account_project_number}-compute@developer.gserviceaccount.com"  # pytype: disable=attribute-error
+    else:
+        raise ValueError(
+            "Either service_account or service_account_project_number " "is required."
+        )
+
+
+def _clean_and_parse_optimized_prompt(output_str: str) -> Optional[Any]:
+    """Cleans a string response returned from the prompt optimizer endpoint.
+
+    Args:
+        output_str: The optimized prompt string containing the JSON data,
+          potentially with markdown formatting like ```json ... ```.
+
+    Returns:
+        The parsed JSON data, or None if parsing fails.
+    """
+    lines = output_str.strip().split("\n")
+    # Remove markdown delimiters
+    if lines and lines[0].strip().startswith("```"):
+        cleaned_string = "\n".join(lines[1:-1])
+    else:
+        cleaned_string = output_str
+
+    # remove any 'json' labels if they exist on the first line.
+    if cleaned_string.strip().startswith("json"):
+        cleaned_string = cleaned_string.strip()[4:].strip()
+
+    try:
+        return json.loads(cleaned_string)
+    except json.JSONDecodeError as e:
+        # TODO(b/437144880): raise errors.ClientError here instead
+        raise ValueError(
+            f"Failed to parse the response from prompt optimizer endpoint. {e}"
+        ) from e
+
+
+def _parse(
+    output_str: str,
+) -> Union[
+    types.prompts.ParsedResponse,
+    types.prompts.ParsedResponseFewShot,
+]:
+    """Parses the output string from the prompt optimizer endpoint."""
+    parsed_out = _clean_and_parse_optimized_prompt(output_str)
+    try:
+        return types.prompts.ParsedResponse(**parsed_out)
+    except ValidationError:
+        return types.prompts.ParsedResponseFewShot(**parsed_out)

agentplatform/_genai/_skills_utils.py

@@ -0,0 +1,69 @@
+# 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 Skills."""
+
+import base64
+import io
+import os
+import pathlib
+import zipfile
+
+
+def zip_directory(directory_path: pathlib.Path | str) -> bytes:
+    """Zips a directory into memory and returns the bytes.
+
+    Args:
+        directory_path (pathlib.Path | str): Required. The local path to the
+          directory.
+
+    Returns:
+        bytes: The zipped directory content.
+    """
+    directory_str = os.fspath(directory_path)
+    if not os.path.isdir(directory_str):
+        raise ValueError(f"Path is not a directory: {directory_str}")
+
+    zip_buffer = io.BytesIO()
+    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+        for root, _, files in os.walk(directory_str):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, directory_str)
+
+                # Read actual file data
+                with open(file_path, "rb") as f:
+                    file_data = f.read()
+
+                # Use deterministic ZipInfo (mtime: 1980-01-01 00:00:00)
+                zinfo = zipfile.ZipInfo(arcname, date_time=(1980, 1, 1, 0, 0, 0))
+                zinfo.compress_type = zipfile.ZIP_DEFLATED
+                zinfo.external_attr = 0o644 << 16  # Constant file permissions
+
+                zip_file.writestr(zinfo, file_data)
+    return zip_buffer.getvalue()
+
+
+def get_zipped_filesystem_payload(directory_path: pathlib.Path | str) -> str:
+    """Zips a directory and base64-encodes the result to a UTF-8 string.
+
+    Args:
+        directory_path (pathlib.Path | str): Required. The local path to the
+          directory.
+
+    Returns:
+        str: The base64-encoded zipped directory.
+    """
+    zip_bytes = zip_directory(directory_path)
+    return base64.b64encode(zip_bytes).decode("utf-8")