Extend permissions service support to Tuned Models (#219)

* Add metaclass to inject Permission methods

Set PermissionAdapterMeta as a metaclass for Corpus and TunedModel classes to inject permission methods dynamically.

* Add method to validate permission name

* Allow accepting partial name for get_permission methods

* Fix consistency issues

* Update tests

* Remove metaclasses and use inheritance instead

* format

* Update tests and fix pytpye error

* Update tests

* disable attribute error for PermissionsAdapter - pytpye for `self.name` field

* Format

* Handle partial names in get_permissions methods

* Update tests

* Update docsting for get_permission method

* format

* Add guard against invalid full name formats

* Update tests

* format

* Update Valid Permission Name regex

* Separate resource name validation for get_permission methods

* Make 	o_proto method internal for Permission class

	o_proto to _to_proto and update invlaid error message for permission_id.

* Update error messages

* started editing the permissions class

Change-Id: Idf1b9e8b0fff498a1c4fd4785c09eb49963c1d47

* Fix tests

Change-Id: Id8cc1cf6e59009f06e8db3ff59edd7ef6ac76943

* Add type hints.

Change-Id: I91f759fe3ebe4e3011440a3983a4cbfc317856d0

* Fix transfer_ownership for Corpora

Change-Id: I791514660881ce967c405e4c76f44739968362f9

* format

Change-Id: I0f98fc80fe81a8fc69aea9f00386aa2c519cb7e0

* Simplify error msgs

* Simplify name checks

* Add deprecation stubs on Corpus

* Fix decorator type checking

ast.Call.func can be of types ast.Attribute and ast.Name, previously only ast.Name was being considered.

* fix types

* fix types (#2)

* Update retriever_types.py

* Update retriever_types.py

---------

Co-authored-by: Mark Daoust <[email protected]>

Author: mayureshagashe2105

google/generativeai/permission.py

@@ -14,31 +14,156 @@
 # limitations under the License.
 from __future__ import annotations
 
+from typing import Callable
+
 import google.ai.generativelanguage as glm
 
 from google.generativeai.types import permission_types
+from google.generativeai.types import retriever_types
+from google.generativeai.types import model_types
+
+
+_RESOURCE_TYPE: dict[str, str] = {
+    "corpus": "corpora",
+    "corpora": "corpora",
+    "tunedmodel": "tunedModels",
+    "tunedmodels": "tunedModels",
+}
+
+
+def _to_resource_type(x: str) -> str:
+    if isinstance(x, str):
+        x = x.lower()
+    resource_type = _RESOURCE_TYPE.get(x, None)
+    if not resource_type:
+        raise ValueError(f"Unsupported resource type. Got: `{x}` instead.")
+
+    return resource_type
+
+
+def _validate_resource_name(x: str, resource_type: str) -> None:
+    if resource_type == "corpora":
+        if not retriever_types.valid_name(x):
+            raise ValueError(retriever_types.NAME_ERROR_MSG.format(length=len(x), name=x))
+
+    elif resource_type == "tunedModels":
+        if not model_types.valid_tuned_model_name(x):
+            raise ValueError(model_types.TUNED_MODEL_NAME_ERROR_MSG.format(length=len(x), name=x))
+
+    else:
+        raise ValueError(f"Unsupported resource type: {resource_type}")
+
+
+def _validate_permission_id(x: str) -> None:
+    if not permission_types.valid_id(x):
+        raise ValueError(permission_types.INVALID_PERMISSION_ID_MSG.format(permission_id=x))
+
+
+def _get_valid_name_components(name: str) -> str:
+    # name is of the format: resource_type/resource_name/permissions/permission_id
+    name_path_components = name.split("/")
+    if len(name_path_components) != 4:
+        raise ValueError(
+            f"Invalid name format. Expected format: \
+                `resource_type/<resource_name>/permissions/<permission_id>`. Got: `{name}` instead."
+        )
+
+    resource_type, resource_name, permission_placeholder, permission_id = name_path_components
+    resource_type = _to_resource_type(resource_type)
+
+    permission_id = "/".join([permission_placeholder, permission_id])
+
+    _validate_resource_name(resource_name, resource_type)
+    _validate_permission_id(permission_id)
+
+    return "/".join([resource_type, resource_name, permission_id])
+
+
+def _construct_name(
+    name: str | None = None,
+    resource_name: str | None = None,
+    permission_id: str | int | None = None,
+    resource_type: str | None = None,
+) -> str:
+    # resource_name is the name of the supported resource (corpus or tunedModel as of now) for which the permission is being created.
+    if not name:
+        # if name is not provided, then try to construct name via provided resource_name and permission_id.
+        if not (resource_name and permission_id):
+            raise ValueError(
+                "Either `name` or (`resource_name` and `permission_id`) must be provided."
+            )
+
+        if resource_type:
+            resource_type = _to_resource_type(resource_type)
+        else:
+            # if resource_type is not provided, then try to infer it from resource_name.
+            resource_path_components = resource_name.split("/")
+            if len(resource_path_components) != 2:
+                raise ValueError(
+                    f"Invalid `resource_name` format. Expected format: \
+                        `resource_type/resource_name`. Got: `{resource_name}` instead."
+                )
+            resource_type = _to_resource_type(resource_path_components[0])
+
+        if f"{resource_type}/" in resource_name:
+            name = f"{resource_name}/"
+        else:
+            name = f"{resource_type}/{resource_name}/"
+
+        if isinstance(permission_id, int) or "permissions/" not in permission_id:
+            name += f"permissions/{permission_id}"
+        else:
+            name += permission_id
+
+    # if name is provided, override resource_name and permission_id
+    name = _get_valid_name_components(name)
+    return name
 
 
 def get_permission(
-    name: str,
+    name: str | None = None,
+    *,
     client: glm.PermissionServiceClient | None = None,
+    resource_name: str | None = None,
+    permission_id: str | int | None = None,
+    resource_type: str | None = None,
 ) -> permission_types.Permission:
-    """Get a permission by name.
+    """Get information about a permission by name.
 
     Args:
         name: The name of the permission.
+        resource_name: The name of the supported resource for which the permission details are needed.
+        permission_id: The name of the permission.
+        resource_type: The type of the resource (corpus or tunedModel as of now) for which the permission details are needed.
+                        If not provided, it will be inferred from `resource_name`.
 
     Returns:
         The permission as an instance of `permission_types.Permission`.
     """
+    name = _construct_name(
+        name=name,
+        resource_name=resource_name,
+        permission_id=permission_id,
+        resource_type=resource_type,
+    )
     return permission_types.Permission.get(name=name, client=client)
 
 
 async def get_permission_async(
-    name: str,
+    name: str | None = None,
+    *,
     client: glm.PermissionServiceAsyncClient | None = None,
+    resource_name: str | None = None,
+    permission_id: str | int | None = None,
+    resource_type: str | None = None,
 ) -> permission_types.Permission:
     """
     This is the async version of `permission.get_permission`.
     """
+    name = _construct_name(
+        name=name,
+        resource_name=resource_name,
+        permission_id=permission_id,
+        resource_type=resource_type,
+    )
     return await permission_types.Permission.get_async(name=name, client=client)

google/generativeai/types/model_types.py

@@ -29,6 +29,7 @@
 from typing_extensions import TypedDict
 
 import google.ai.generativelanguage as glm
+from google.generativeai.types import permission_types
 from google.generativeai import string_utils
 
 
@@ -47,6 +48,17 @@
 
 TunedModelStateOptions = Union[None, str, int, TunedModelState]
 
+_TUNED_MODEL_VALID_NAME = r"[a-z](([a-z0-9-]{0,61}[a-z0-9])?)$"
+TUNED_MODEL_NAME_ERROR_MSG = """The `name` must consist of alphanumeric characters (or -) and be at most 63 characters; The name you entered:
+\tlen(name)== {length}
+\tname={name}
+"""
+
+
+def valid_tuned_model_name(name: str) -> bool:
+    return re.match(_TUNED_MODEL_VALID_NAME, name) is not None
+
+
 # fmt: off
 _TUNED_MODEL_STATES: dict[TunedModelStateOptions, TunedModelState] = {
     TunedModelState.ACTIVE: TunedModelState.ACTIVE,
@@ -183,6 +195,10 @@ class TunedModel:
     update_time: datetime.datetime | None = None
     tuning_task: TuningTask | None = None
 
+    @property
+    def permissions(self) -> permission_types.Permissions:
+        return permission_types.Permissions(self)
+
 
 @string_utils.prettyprint
 @dataclasses.dataclass