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

agentplatform/_genai/_gcs_utils.py

@@ -0,0 +1,262 @@
+# 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.
+#
+
+from importlib.metadata import version as get_version
+import io
+import json
+import logging
+from typing import Any, Union
+import uuid
+
+from google.cloud import storage  # type: ignore[attr-defined]
+from google.genai._api_client import BaseApiClient
+from packaging.version import Version
+import pandas as pd
+
+
+logger = logging.getLogger(__name__)
+
+
+GCS_PREFIX = "gs://"
+
+
+# Detect google-cloud-storage version once at module load
+try:
+    _GCS_VERSION = Version(get_version("google-cloud-storage"))
+except Exception:
+    # Fallback if version detection fails (should not happen in normal use)
+    _GCS_VERSION = Version("3.0.0")
+
+_USE_FROM_URI = _GCS_VERSION >= Version("3.0.0")
+
+
+def blob_from_uri(uri: str, client: storage.Client) -> storage.Blob:
+    """Create a Blob from a GCS URI, compatible with v2 and v3.
+
+    This function provides compatibility across google-cloud-storage versions:
+    - v3.x: Uses Blob.from_uri()
+    - v2.x: Uses Blob.from_string() (deprecated in v3)
+
+    Args:
+        uri: GCS URI (e.g., 'gs://bucket/path/to/blob')
+        client: Storage client instance
+
+    Returns:
+        storage.Blob: Blob instance
+    """
+    if _USE_FROM_URI:
+        return storage.Blob.from_uri(uri, client=client)
+    else:
+        return storage.Blob.from_string(uri, client=client)
+
+
+class GcsUtils:
+    """Handles File I/O operations with Google Cloud Storage (GCS)"""
+
+    def __init__(self, api_client: BaseApiClient):
+        self.api_client = api_client
+        self.storage_client = storage.Client(
+            project=self.api_client.project,
+            credentials=self.api_client._credentials,
+        )
+
+    def parse_gcs_path(self, gcs_path: str) -> tuple[str, str]:
+        """Helper to parse gs://bucket/path into (bucket_name, blob_path)."""
+        if not gcs_path.startswith(GCS_PREFIX):
+            raise ValueError(
+                f"Invalid GCS path: '{gcs_path}'. It must start with '{GCS_PREFIX}'."
+            )
+        path_without_prefix = gcs_path[len(GCS_PREFIX) :]
+        if "/" not in path_without_prefix:
+            return path_without_prefix, ""
+        bucket_name, blob_path = path_without_prefix.split("/", 1)
+        return bucket_name, blob_path
+
+    def upload_file_to_gcs(self, upload_gcs_path: str, filename: str) -> None:
+        """Uploads the provided file to a Google Cloud Storage location."""
+
+        blob_from_uri(
+            uri=upload_gcs_path, client=self.storage_client
+        ).upload_from_filename(filename)
+
+    def upload_dataframe(
+        self,
+        df: "pd.DataFrame",
+        gcs_destination_blob_path: str,
+        file_type: str = "jsonl",
+    ) -> None:
+        """Uploads a Pandas DataFrame to a Google Cloud Storage location.
+
+        Args:
+          df: The Pandas DataFrame to upload.
+          gcs_destination_blob_path: The full GCS path for the destination blob
+            (e.g., 'gs://bucket/data/my_dataframe.jsonl').
+          file_type: The format to save the DataFrame ('jsonl' or 'csv'). Defaults
+            to 'jsonl'.
+        """
+        bucket_name, blob_name = self.parse_gcs_path(gcs_destination_blob_path)
+        if not blob_name:
+            raise ValueError(
+                f"Invalid GCS path for blob: '{gcs_destination_blob_path}'. "
+                "It must include the object name (e.g., gs://bucket/file.csv)."
+            )
+        bucket = self.storage_client.bucket(bucket_name)
+        blob = bucket.blob(blob_name)
+
+        buffer = io.StringIO()
+        if file_type == "csv":
+            df.to_csv(buffer, index=False)
+            content_type = "text/csv"
+        elif file_type == "jsonl":
+            df.to_json(buffer, orient="records", lines=True)
+            content_type = "application/jsonl"
+        else:
+            raise ValueError(
+                f"Unsupported file type: '{file_type}'. "
+                "Please provide 'jsonl' or 'csv'."
+            )
+        blob.upload_from_string(buffer.getvalue(), content_type=content_type)
+
+        logger.info(
+            f"DataFrame successfully uploaded to: gs://{bucket.name}/{blob.name}"
+        )
+
+    def upload_json(self, data: dict[str, Any], gcs_destination_blob_path: str) -> None:
+        """Uploads a dictionary as a JSON file to Google Cloud Storage."""
+        bucket_name, blob_name = self.parse_gcs_path(gcs_destination_blob_path)
+        if not blob_name:
+            raise ValueError(
+                f"Invalid GCS path for blob: '{gcs_destination_blob_path}'. "
+                "It must include the object name (e.g., gs://bucket/file.json)."
+            )
+        bucket = self.storage_client.bucket(bucket_name)
+        blob = bucket.blob(blob_name)
+
+        json_data = json.dumps(data, indent=2)
+        blob.upload_from_string(json_data, content_type="application/json")
+
+        logger.info(
+            f"JSON data successfully uploaded to: gs://{bucket_name}/{blob_name}"
+        )
+
+    def upload_json_to_prefix(
+        self,
+        data: dict[str, Any],
+        gcs_dest_prefix: str,
+        filename_prefix: str = "data",
+    ) -> str:
+        """Uploads a dictionary to a GCS prefix with a UUID JSON filename.
+
+        Args:
+          data: The dictionary to upload.
+          gcs_dest_prefix: The GCS prefix (e.g., 'gs://bucket/path/prefix/').
+          filename_prefix: Prefix for the generated filename. Defaults to 'data'.
+
+        Returns:
+          The full GCS path where the file was uploaded.
+
+        Raises:
+          ValueError: If the gcs_dest_prefix is not a valid GCS path.
+        """
+        if not gcs_dest_prefix.startswith(GCS_PREFIX):
+            raise ValueError(
+                f"Invalid GCS destination prefix: '{gcs_dest_prefix}'. Must start"
+                f" with '{GCS_PREFIX}'."
+            )
+
+        gcs_path_without_scheme = gcs_dest_prefix[len(GCS_PREFIX) :]
+        bucket_name, *path_parts = gcs_path_without_scheme.split("/")
+
+        user_prefix_path = "/".join(path_parts)
+        if user_prefix_path and not user_prefix_path.endswith("/"):
+            user_prefix_path += "/"
+
+        filename = f"{filename_prefix}_{uuid.uuid4()}.json"
+
+        blob_name = f"{user_prefix_path}{filename}"
+
+        full_gcs_path = f"{GCS_PREFIX}{bucket_name}/{blob_name}"
+
+        self.upload_json(data, full_gcs_path)
+        return full_gcs_path
+
+    def read_file_contents(self, gcs_filepath: str) -> Union[str, Any]:
+        """Reads the contents of a file from Google Cloud Storage."""
+
+        bucket_name, blob_path = self.parse_gcs_path(gcs_filepath)
+        if not blob_path:
+            raise ValueError(
+                f"Invalid GCS file path: '{gcs_filepath}'. Path must point to a file,"
+                " not just a bucket."
+            )
+        bucket = self.storage_client.bucket(bucket_name)
+        blob = bucket.blob(blob_path)
+        content = blob.download_as_bytes().decode("utf-8")
+        logger.info(f"Successfully read content from '{gcs_filepath}'")
+        return content
+
+    def read_gcs_file_to_dataframe(
+        self, gcs_filepath: str, file_type: str
+    ) -> "pd.DataFrame":
+        """Reads a file from Google Cloud Storage into a Pandas DataFrame."""
+        file_contents = self.read_file_contents(gcs_filepath)
+        if file_type == "csv":
+            return pd.read_csv(io.StringIO(file_contents), encoding="utf-8")
+        elif file_type == "jsonl":
+            return pd.read_json(io.StringIO(file_contents), lines=True)
+        else:
+            raise ValueError(
+                f"Unsupported file type: '{file_type}'. Please provide 'jsonl' or"
+                " 'csv'."
+            )
+
+    def _verify_bucket_ownership(
+        self,
+        bucket_name: str,
+        expected_project: str,
+    ) -> bool:
+        """Verifies that a GCS bucket belongs to the expected project.
+
+        This check mitigates bucket squatting attacks.
+
+        Args:
+            bucket_name: The GCS bucket to verify.
+            expected_project: The project ID or number that should own the bucket.
+
+        Returns:
+            True if the bucket belongs to the expected project, False otherwise.
+        """
+        try:
+            bucket = self.storage_client.bucket(bucket_name=bucket_name)
+            bucket.reload(client=self.storage_client)
+            bucket_project_number = str(bucket.project_number)
+
+            if expected_project.isdigit():
+                expected_project_number = expected_project
+            else:
+                from google.cloud import resourcemanager_v3
+
+                projects_client = resourcemanager_v3.ProjectsClient(
+                    credentials=self.storage_client._credentials
+                )
+                project = projects_client.get_project(
+                    name=f"projects/{expected_project}"
+                )
+
+                expected_project_number = project.name.split("/")[-1]
+
+            return bucket_project_number == expected_project_number
+        except Exception:
+            return False

agentplatform/_genai/_logging_utils.py

@@ -0,0 +1,47 @@
+# 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.
+#
+
+import functools
+from typing import Any, Callable
+from google.genai import _common
+import warnings
+
+
+def show_deprecation_warning_once(
+    message: str,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """Decorator to show a deprecation warning once for a function."""
+
+    def decorator(func: Any) -> Any:
+        warning_done = False
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            nonlocal warning_done
+            if not warning_done:
+                warning_done = True
+                warnings.warn(message, DeprecationWarning, stacklevel=2)
+
+            # Suppress ExperimentalWarning while executing the deprecated wrapper
+            with warnings.catch_warnings():
+                # We ignore ExperimentalWarning because the user will see it
+                # when they migrate to the new prompts module
+                warnings.simplefilter("ignore", category=_common.ExperimentalWarning)
+                return func(*args, **kwargs)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

agentplatform/_genai/_memory_bank_utils.py

@@ -0,0 +1,206 @@
+# 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 memory banks."""
+
+import asyncio
+import json
+import re
+import time
+from typing import (
+    Any,
+    Protocol,
+    Union,
+)
+
+from . import types as genai_types
+
+
+MemoryBankOperation = Union[
+    genai_types.MemoryBankOperation,
+    genai_types.MemoryOperation,
+    genai_types.GenerateMemoriesOperation,
+]
+
+
+class GetOperationFunction(Protocol):
+    def __call__(self, *, operation_name: str, **kwargs: Any) -> MemoryBankOperation:
+        pass
+
+
+class GetAsyncOperationFunction(Protocol):
+    async def __call__(
+        self, *, operation_name: str, **kwargs: Any
+    ) -> MemoryBankOperation:
+        pass
+
+
+def _get_memory_bank_id(operation_name: str = "", resource_name: str = "") -> str:
+    """Returns Memory Bank ID from operation name or resource name."""
+    if not resource_name and not operation_name:
+        raise ValueError("Resource name or operation name cannot be empty.")
+
+    if resource_name:
+        match = re.match(
+            r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)$",
+            resource_name,
+        )
+        if match:
+            return match.group(1)
+        match = re.match(
+            r"^projects/[^/]+/locations/[^/]+/memoryBanks/([^/]+)$",
+            resource_name,
+        )
+        if match:
+            return match.group(1)
+        raise ValueError(
+            "Failed to parse Memory Bank ID from resource name: " f"`{resource_name}`"
+        )
+
+    if not operation_name:
+        raise ValueError("Operation name cannot be empty.")
+
+    match = re.match(
+        r"^projects/[^/]+/locations/[^/]+/reasoningEngines/([^/]+)/operations/[^/]+$",
+        operation_name,
+    )
+    if match:
+        return match.group(1)
+
+    match = re.match(
+        r"^projects/[^/]+/locations/[^/]+/memoryBanks/([^/]+)/operations/[^/]+$",
+        operation_name,
+    )
+    raise ValueError(
+        "Failed to parse Memory Bank ID from operation name: " f"`{operation_name}`"
+    )
+
+
+def _await_operation(
+    *,
+    operation_name: str,
+    get_operation_fn: GetOperationFunction,
+    poll_interval_seconds: float = 1,
+) -> Any:
+    """Waits for the operation to complete.
+
+    Args:
+        operation_name (str):
+            Required. The name of the operation.
+        poll_interval_seconds (float):
+            The number of seconds to wait between each poll.
+        get_operation_fn (Callable[[str], Any]):
+            Optional. The function to use for getting the operation. If not
+            provided, `self._get_memory_bank_operation` will be used.
+
+    Returns:
+        The operation that has completed (i.e. `operation.done==True`).
+    """
+    operation = get_operation_fn(operation_name=operation_name)
+    while not operation.done:
+        time.sleep(poll_interval_seconds)
+        operation = get_operation_fn(operation_name=operation.name)
+
+    return operation
+
+
+async def _await_async_operation(
+    *,
+    operation_name: str,
+    get_operation_fn: GetAsyncOperationFunction,
+    poll_interval_seconds: float = 1,
+) -> Any:
+    """Waits for the operation to complete.
+
+    Args:
+        operation_name (str):
+            Required. The name of the operation.
+        poll_interval_seconds (float):
+            The number of seconds to wait between each poll.
+        get_operation_fn (Callable[[str], Awaitable[Any]]):
+            Optional. The async function to use for getting the operation. If not
+            provided, `self._get_memory_bank_operation` will be used.
+
+    Returns:
+        The operation that has completed (i.e. `operation.done==True`).
+    """
+    operation = await get_operation_fn(operation_name=operation_name)
+    while not operation.done:
+        await asyncio.sleep(poll_interval_seconds)
+        operation = await get_operation_fn(operation_name=operation.name)
+
+    return operation
+
+
+def _managed_semantic_memory_config_to_memory_bank_config(
+    semantic_memory_config: genai_types.ManagedSemanticMemoryConfigOrDict,
+) -> genai_types.ReasoningEngineContextSpecMemoryBankConfigDict:
+    """Converts ManagedSemanticMemoryConfig to MemoryBankConfig."""
+    if semantic_memory_config is None:
+        semantic_memory_config = {}
+    if isinstance(semantic_memory_config, dict):
+        semantic_memory_config = genai_types.ManagedSemanticMemoryConfig.model_validate(
+            semantic_memory_config
+        )
+    elif not isinstance(
+        semantic_memory_config, genai_types.ManagedSemanticMemoryConfig
+    ):
+        raise TypeError(
+            "managed_semantic_memory_config must be a dict or "
+            "ManagedSemanticMemoryConfig, "
+            f"but got {type(semantic_memory_config)}."
+        )
+
+    memory_bank_config = json.loads(semantic_memory_config.model_dump_json())
+    if "unstructured_memory_configs" in memory_bank_config:
+        memory_bank_config["customization_configs"] = memory_bank_config.pop(
+            "unstructured_memory_configs"
+        )
+    return memory_bank_config
+
+
+def _memory_bank_config_to_managed_semantic_memories_config(
+    memory_bank_config: genai_types.ReasoningEngineContextSpecMemoryBankConfig,
+) -> genai_types.ManagedSemanticMemoryConfigDict:
+    """Converts MemoryBankConfig to ManagedSemanticMemoriesConfig."""
+    memory_bank_config = json.loads(memory_bank_config.model_dump_json())
+    if "customization_configs" in memory_bank_config:
+        memory_bank_config["unstructured_memory_configs"] = memory_bank_config.pop(
+            "customization_configs"
+        )
+    return memory_bank_config
+
+
+def _reasoning_engine_to_memory_bank(
+    reasoning_engine: genai_types.ReasoningEngine,
+) -> genai_types.MemoryBank:
+    """Converts ReasoningEngine to MemoryBank."""
+    if reasoning_engine.context_spec is not None:
+        semantic_memory_config = (
+            _memory_bank_config_to_managed_semantic_memories_config(
+                reasoning_engine.context_spec.memory_bank_config
+            )
+        )
+    else:
+        semantic_memory_config = {}
+    memory_bank = genai_types.MemoryBank(
+        name=reasoning_engine.name,
+        create_time=reasoning_engine.create_time,
+        update_time=reasoning_engine.update_time,
+        display_name=reasoning_engine.display_name,
+        description=reasoning_engine.description,
+        encryption_spec=reasoning_engine.encryption_spec,
+        managed_semantic_memory_config=semantic_memory_config,
+    )
+    return memory_bank

agentplatform/_genai/_observability_data_converter.py

@@ -0,0 +1,186 @@
+# 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.
+#
+"""Dataset converter for Google Observability GenAI data."""
+
+import json
+import logging
+from typing import Any, Optional
+
+from google.genai import types as genai_types
+from typing_extensions import override
+
+from . import _evals_utils
+from . import types
+
+
+logger = logging.getLogger("agentplatform_genai._observability_data_converters")
+
+
+def _load_jsonl(data: Any, case_id: str) -> list[dict[Any, Any]]:
+    """Parses the raw JSONL data into a list of dict possible."""
+    if isinstance(data, str):
+        json_list = []
+        for line in data.splitlines():
+            loaded_json = json.loads(line)
+            if not isinstance(loaded_json, dict):
+                raise TypeError(
+                    f"Decoded JSON payload is not a dict for case "
+                    f"{case_id}. Type found: {type(loaded_json).__name__}"
+                )
+            json_list.append(loaded_json)
+        return json_list
+    else:
+        raise TypeError(
+            f"Payload is not a JSONL string for case {case_id}. Type "
+            f"found: {type(data).__name__}"
+        )
+
+
+class ObservabilityDataConverter(_evals_utils.EvalDataConverter):
+    """Converter for dataset in GCP Observability GenAI format."""
+
+    def _message_to_content(self, message: dict[str, Any]) -> genai_types.Content:
+        """Converts Observability GenAI Message format to Content."""
+        parts = []
+        message_parts = message.get("parts", [])
+        if isinstance(message_parts, list):
+            for message_part in message_parts:
+                part = None
+                part_type = message_part.get("type", "")
+                if part_type == "text":
+                    part = genai_types.Part(text=message_part.get("content", ""))
+                elif part_type == "blob":
+                    part = genai_types.Part(
+                        inline_data=genai_types.Blob(
+                            data=message_part.get("data", ""),
+                            mime_type=message_part.get("mime_type", ""),
+                        )
+                    )
+                elif part_type == "file_data":
+                    part = genai_types.Part(
+                        file_data=genai_types.FileData(
+                            file_uri=message_part.get("file_uri", ""),
+                            mime_type=message_part.get("mime_type", ""),
+                        )
+                    )
+                elif part_type == "tool_call":
+                    # O11y format requires use of id in place of name
+                    part = genai_types.Part(
+                        function_call=genai_types.FunctionCall(
+                            id=message_part.get("id", ""),
+                            name=message_part.get("id", ""),
+                            args=message_part.get("arguments", {}),
+                        )
+                    )
+                elif part_type == "tool_call_response":
+                    # O11y format requires use of id in place of name
+                    part = genai_types.Part(
+                        function_response=genai_types.FunctionResponse(
+                            id=message_part.get("id", ""),
+                            name=message_part.get("id", ""),
+                            response=message_part.get("result", {}),
+                        )
+                    )
+                else:
+                    logger.warning(
+                        "Skipping message part due to unrecognized message "
+                        "part type of '%s'",
+                        part_type,
+                    )
+
+                if part is not None:
+                    parts.append(part)
+
+        return genai_types.Content(parts=parts, role=message.get("role", ""))
+
+    def _parse_messages(
+        self,
+        eval_case_id: str,
+        request_msgs: list[Any],
+        response_msgs: list[Any],
+        system_instruction_msg: Optional[dict[str, Any]] = None,
+    ) -> types.EvalCase:
+        """Parses a set of Observability messages into an EvalCase."""
+        # System instruction message
+        system_instruction = None
+        if system_instruction_msg is not None:
+            system_instruction = self._message_to_content(system_instruction_msg)
+
+        # Request messages
+        prompt = None
+        conversation_history = []
+        if request_msgs:
+            # Extract latest message as prompt
+            prompt = self._message_to_content(request_msgs[-1])
+
+            # All previous messages are conversation history
+            if len(request_msgs) > 1:
+                for i, msg in enumerate(request_msgs[:-1]):
+                    conversation_history.append(
+                        types.evals.Message(
+                            turn_id=str(i),
+                            content=self._message_to_content(msg),
+                            author=msg.get("role", ""),
+                        )
+                    )
+
+        # Output messages
+        responses = []
+        for msg in response_msgs:
+            response = types.ResponseCandidate(response=self._message_to_content(msg))
+            responses.append(response)
+
+        return types.EvalCase(
+            eval_case_id=eval_case_id,
+            prompt=prompt,
+            responses=responses,
+            system_instruction=system_instruction,
+            conversation_history=conversation_history,
+            reference=None,
+        )
+
+    @override
+    def convert(self, raw_data: list[dict[str, Any]]) -> types.EvaluationDataset:
+        """Converts a list of GCP Observability GenAI cases into an EvaluationDataset."""
+        eval_cases = []
+
+        for i, case in enumerate(raw_data):
+            eval_case_id = f"observability_eval_case_{i}"
+
+            if "request" not in case or "response" not in case:
+                logger.warning(
+                    "Skipping case %s due to missing 'request' or 'response' key.",
+                    eval_case_id,
+                )
+                continue
+
+            request_data = case.get("request", [])
+            request_list = _load_jsonl(request_data, eval_case_id)
+
+            response_data = case.get("response", [])
+            response_list = _load_jsonl(response_data, eval_case_id)
+
+            system_dict = None
+            if "system_instruction" in case:
+                system_data = case.get("system_instruction", {})
+                system_list = _load_jsonl(system_data, eval_case_id)
+                system_dict = system_list[0] if system_list else {}
+
+            eval_case = self._parse_messages(
+                eval_case_id, request_list, response_list, system_dict
+            )
+            eval_cases.append(eval_case)
+
+        return types.EvaluationDataset(eval_cases=eval_cases)