qwak-core 0.6.7__py3-none-any.whl → 0.7.0__py3-none-any.whl

frogml_storage/utils/_url_utils.py

@@ -0,0 +1,27 @@
+from typing import Optional, List
+from urllib.parse import urlparse
+
+
+def join_url(base_uri: str, *parts: str) -> str:
+    base_uri = base_uri.rstrip("/")
+
+    cleaned_parts: List[str] = [
+        part.strip("/") for part in parts if part is not None and part.strip("/")
+    ]
+    uri_parts: List[str] = [base_uri, *cleaned_parts]
+
+    return "/".join(uri_parts)
+
+
+def assemble_artifact_url(uri: Optional[str]) -> str:
+    if uri is None:
+        raise Exception("Artifactory URI is required")
+
+    parsed_url = urlparse(uri)
+    if parsed_url.scheme not in ["http", "https"]:
+        raise Exception(
+            f"Not a valid Artifactory URI: {uri}. "
+            f"Artifactory URI example: `https://frogger.jfrog.io/artifactory/ml-local`"
+        )
+
+    return f"{parsed_url.scheme}://{parsed_url.netloc}/artifactory"

qwak/__init__.py

@@ -1,7 +1,7 @@
 """Top-level package for qwak-core."""
 
 __author__ = "Qwak.ai"
-__version__ = "0.6.7"
+__version__ = "0.7.0"
 
 from qwak.inner.di_configuration import wire_dependencies
 from qwak.model.experiment_tracking import log_metric, log_param

qwak/clients/model_management/client.py

@@ -12,6 +12,7 @@ from _qwak_proto.qwak.models.models_pb2 import (
     ModelSpec,
 )
 from _qwak_proto.qwak.models.models_pb2_grpc import ModelsManagementServiceStub
+from _qwak_proto.qwak.projects.jfrog_project_spec_pb2 import ModelRepositoryJFrogSpec
 from dependency_injector.wiring import Provide
 from qwak.exceptions import QwakException
 from qwak.inner.di_configuration import QwakContainer
@@ -58,6 +59,7 @@ class ModelsManagementClient:
         project_id,
         model_name,
         model_description,
+        jfrog_project_key: Optional[str] = None,
     ):
         try:
             return self._models_management_service.CreateModel(
@@ -68,6 +70,9 @@ class ModelsManagementClient:
                         project_id=project_id,
                         model_description=model_description,
                     ),
+                    jfrog_project_spec=ModelRepositoryJFrogSpec(
+                        jfrog_project_key=jfrog_project_key,
+                    ),
                 )
             )
 

qwak/clients/project/client.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import grpc
 from _qwak_proto.qwak.projects.projects_pb2 import (
     CreateProjectRequest,
@@ -6,6 +8,7 @@ from _qwak_proto.qwak.projects.projects_pb2 import (
     ListProjectsRequest,
 )
 from _qwak_proto.qwak.projects.projects_pb2_grpc import ProjectsManagementServiceStub
+from _qwak_proto.qwak.projects.jfrog_project_spec_pb2 import ModelRepositoryJFrogSpec
 from dependency_injector.wiring import Provide, inject
 from qwak.exceptions import QwakException
 from qwak.inner.di_configuration import QwakContainer
@@ -31,12 +34,16 @@ class ProjectsManagementClient:
         self,
         project_name,
         project_description,
+        jfrog_project_key: Optional[str] = None,
     ):
         try:
             return self._projects_management_service.CreateProject(
                 CreateProjectRequest(
                     project_name=project_name,
                     project_description=project_description,
+                    jfrog_spec=ModelRepositoryJFrogSpec(
+                        jfrog_project_key=jfrog_project_key
+                    ),
                 )
             )
 

qwak/inner/const.py

@@ -35,6 +35,14 @@ class QwakConstants:
 
     TOKEN_AUDIENCE: str = "https://auth-token.qwak.ai/"  # nosec B105
 
+    QWAK_AUTHENTICATION_URL = "https://grpc.qwak.ai/api/v1/authentication/qwak-api-key"
+
+    QWAK_AUTHENTICATED_USER_ENDPOINT: str = (
+        "https://grpc.qwak.ai/api/v0/runtime/get-authenticated-user-context"
+    )
+
+    JFROG_TENANT_HEADER_KEY = "X-JFrog-Tenant-Id"
+
     QWAK_APP_URL: str = "https://app.qwak.ai"
 
     CONTROL_PLANE_GRPC_ADDRESS_ENVAR_NAME: str = "CONTROL_PLANE_GRPC_ADDRESS"

qwak/inner/di_configuration/account.py

@@ -2,12 +2,13 @@ import configparser
 import errno
 import os
 from dataclasses import dataclass
-from typing import Optional, Type
+from typing import Optional, Type, Union
 
 from qwak.exceptions import QwakLoginException
 from qwak.inner.const import QwakConstants
 from qwak.inner.di_configuration.session import Session
-from qwak.inner.tool.auth import Auth0ClientBase
+from qwak.inner.tool.auth import Auth0ClientBase, FrogMLAuthClient
+from frogml_storage.authentication.login import frogml_login
 
 
 @dataclass
@@ -22,6 +23,18 @@ class UserAccount:
     # Assigned username
     username: Optional[str] = None
 
+    # Assigned password
+    password: Optional[str] = None
+
+    # Assigned URL
+    url: Optional[str] = None
+
+    # Anonymous login
+    anonymous: bool = False
+
+    # Interactive login
+    is_interactive: bool = False
+
 
 class UserAccountConfiguration:
     USER_FIELD = "user"
@@ -30,9 +43,11 @@ class UserAccountConfiguration:
 
     def __init__(
         self,
-        config_file: str = QwakConstants.QWAK_CONFIG_FILE,
-        auth_file: str = QwakConstants.QWAK_AUTHORIZATION_FILE,
-        auth_client: Type[Auth0ClientBase] = Auth0ClientBase,
+        config_file=QwakConstants.QWAK_CONFIG_FILE,
+        auth_file=QwakConstants.QWAK_AUTHORIZATION_FILE,
+        auth_client: Optional[
+            Union[Type[Auth0ClientBase], Type[FrogMLAuthClient]]
+        ] = None,
     ):
         self._config_file = config_file
         self._auth_file = auth_file
@@ -40,12 +55,51 @@ class UserAccountConfiguration:
         self._auth = configparser.ConfigParser()
         self._environment = Session().get_environment()
         self._auth_client = auth_client
+        self._force_qwak_auth = os.getenv("FORCE_QWAK_AUTH", "False") == "True"
+
+        if not self._auth_client:
+            # Determine auth client based on FrogML configuration
+            try:
+                from frogml_storage.authentication.utils import (
+                    get_frogml_configuration,
+                )
+
+                if (
+                    get_frogml_configuration() or os.getenv("JF_URL")
+                ) and not self._force_qwak_auth:
+                    self._auth_client = FrogMLAuthClient
+                else:
+                    self._auth_client = Auth0ClientBase
+            except Exception:
+                self._auth_client = Auth0ClientBase
 
     def configure_user(self, user_account: UserAccount):
         """
         Configure user authentication based on the authentication client type
         """
-        self.__qwak_login(user_account)
+        if issubclass(self._auth_client, Auth0ClientBase):
+            # Existing Qwak authentication flow
+            self.__qwak_login(user_account)
+
+        elif issubclass(self._auth_client, FrogMLAuthClient):
+            # Use FrogML's login flow
+            success = frogml_login(
+                url=user_account.url,
+                username=user_account.username,
+                password=user_account.password,
+                token=user_account.api_key,
+                anonymous=user_account.anonymous,
+                is_interactive=user_account.is_interactive,
+            )
+
+            if not success:
+                raise QwakLoginException("Failed to authenticate with JFrog")
+            # Validate access token
+            token = self._auth_client().get_token()
+            if not token or len(token) <= 64:
+                raise QwakLoginException(
+                    "Authentication with JFrog failed: Only Access Tokens are supported. Please ensure you are using a valid Access Token."
+                )
 
     def __qwak_login(self, user_account: UserAccount):
         self._auth.read(self._auth_file)
@@ -86,6 +140,8 @@ class UserAccountConfiguration:
         :return:
         """
         try:
+            if issubclass(self._auth_client, FrogMLAuthClient):
+                return UserAccount()
             username = os.environ.get("QWAK_USERNAME")
             api_key = os.environ.get("QWAK_API_KEY")
             if not api_key and (
@@ -124,6 +180,8 @@ class UserAccountConfiguration:
         :return:
         """
         try:
+            if issubclass(self._auth_client, FrogMLAuthClient):
+                return ""
             api_key = os.environ.get("QWAK_API_KEY")
             if api_key:
                 Session().set_environment(api_key)
@@ -156,4 +214,7 @@ class UserAccountConfiguration:
         auth_client_instance = self._auth_client()
         base_url = auth_client_instance.get_base_url()
 
+        if issubclass(self._auth_client, FrogMLAuthClient):
+            return f"{base_url}/ui/ml"
+
         return base_url

qwak/inner/di_configuration/dependency_wiring.py

@@ -50,7 +50,6 @@ def wire_dependencies():
         prompt_manager,
         system_secret,
         user_application_instance,
-        vector_store,
         workspace_manager,
     )
 
@@ -76,7 +75,6 @@ def wire_dependencies():
             user_application_instance,
             alerts_registry,
             workspace_manager,
-            vector_store,
             integration_management,
             system_secret,
             prompt_manager,

qwak/inner/tool/auth.py

@@ -1,9 +1,12 @@
 import warnings
 from filelock import FileLock
+from typing_extensions import Self
 
 from qwak.inner.di_configuration.session import Session
 from abc import ABC, abstractmethod
 from typing import Optional
+from frogml_storage.authentication.utils import get_credentials
+from frogml_storage.authentication.models import AuthConfig
 
 warnings.filterwarnings(action="ignore", module=".*jose.*")
 
@@ -137,3 +140,83 @@ class Auth0ClientBase(BaseAuthClient):
             raise e
         except Exception:
             raise QwakLoginException()
+
+
+class FrogMLAuthClient(BaseAuthClient):
+    __MIN_TOKEN_LENGTH: int = 64
+    __FAIL_TO_AUTH_ERROR_MESSAGE = (
+        "Failed to authenticate with JFrog. Please check your artifactory configuration"
+    )
+
+    def __init__(self, auth_config: Optional[AuthConfig] = None):
+        self.auth_config = auth_config
+        self._token = None
+        self._tenant_id = None
+
+    def get_token(self) -> Optional[str]:
+        if not self._token:
+            self.login()
+        return self._token
+
+    def get_tenant_id(self) -> Optional[str]:
+        if not self._tenant_id:
+            self.login()
+        return self._tenant_id
+
+    def get_base_url(self) -> str:
+        artifactory_url, _ = get_credentials(self.auth_config)
+        return self.__format_artifactory_url(artifactory_url)
+
+    def login(self) -> None:
+        artifactory_url, auth = get_credentials(self.auth_config)
+        # For now, we only support Bearer token authentication
+        if not hasattr(auth, "token"):
+            return
+
+        # noinspection PyUnresolvedReferences
+        self._token = auth.token
+        self.__validate_token()
+
+        base_url = self.__format_artifactory_url(artifactory_url)
+        try:
+            response = requests.get(
+                f"{base_url}/ui/api/v1/system/auth/screen/footer",
+                headers={"Authorization": f"Bearer {self._token}"},
+                timeout=60,
+            )
+            response.raise_for_status()  # Raises an HTTPError for bad responses
+            response_data = response.json()
+            if "serverId" not in response_data:
+                response = requests.get(
+                    f"{base_url}/jfconnect/api/v1/system/jpd_id",
+                    headers={"Authorization": f"Bearer {self._token}"},
+                    timeout=60,
+                )
+                if response.status_code == 200:
+                    self._tenant_id = response.text
+                elif response.status_code == 401:
+                    raise QwakLoginException(
+                        "Failed to authenticate with JFrog. Please check your credentials"
+                    )
+                else:
+                    raise QwakLoginException(self.__FAIL_TO_AUTH_ERROR_MESSAGE)
+            else:
+                self._tenant_id = response_data["serverId"]
+        except requests.exceptions.RequestException:
+            raise QwakLoginException(self.__FAIL_TO_AUTH_ERROR_MESSAGE)
+        except ValueError:  # This catches JSON decode errors
+            raise QwakLoginException(self.__FAIL_TO_AUTH_ERROR_MESSAGE)
+
+    def __validate_token(self: Self):
+        if self._token is None or len(self._token) <= self.__MIN_TOKEN_LENGTH:
+            raise QwakLoginException(
+                "Authentication with JFrog failed: Only JWT Access Tokens are supported. "
+                "Please ensure you are using a valid JWT Access Token."
+            )
+
+    @staticmethod
+    def __format_artifactory_url(artifactory_url: str) -> str:
+        # Remove '/artifactory' from the URL
+        base_url: str = artifactory_url.replace("/artifactory", "")
+        # Remove trailing slash if exists
+        return base_url.rstrip("/")

qwak/inner/tool/grpc/grpc_auth.py

@@ -1,6 +1,7 @@
 from typing import Callable, Tuple
 
 import grpc
+from qwak.inner.const import QwakConstants
 
 _SIGNATURE_HEADER_KEY = "authorization"
 
@@ -34,3 +35,37 @@ class Auth0Client(grpc.AuthMetadataPlugin):
 
             self._auth_client = Auth0ClientBase()
         return self._auth_client.get_token()
+
+
+class FrogMLGrpcClient(grpc.AuthMetadataPlugin):
+    def __init__(self):
+        self._auth_client = None
+
+    def __call__(
+        self,
+        context: grpc.AuthMetadataContext,
+        callback: Callable[[Tuple[Tuple[str, str]], None], None],
+    ):
+        """Implements authentication by passing metadata to a callback.
+
+        Args:
+            context: An AuthMetadataContext providing information on the RPC that
+                the plugin is being called to authenticate.
+            callback: A callback that accepts a tuple of metadata key/value pairs and a None
+                parameter.
+        """
+        token = self.get_token()
+        jfrog_tenant_id = self._auth_client.get_tenant_id()
+        metadata = (
+            (_SIGNATURE_HEADER_KEY, f"Bearer {token}"),
+            (QwakConstants.JFROG_TENANT_HEADER_KEY.lower(), jfrog_tenant_id),
+        )
+        callback(metadata, None)
+
+    def get_token(self) -> str:
+        """Get the authentication token."""
+        if not self._auth_client:
+            from qwak.inner.tool.auth import FrogMLAuthClient
+
+            self._auth_client = FrogMLAuthClient()
+        return self._auth_client.get_token()

qwak/inner/tool/grpc/grpc_tools.py

@@ -7,9 +7,12 @@ from typing import Callable, Optional, Tuple
 from urllib.parse import urlparse, ParseResult
 
 import grpc
+from grpc import AuthMetadataPlugin, Channel
 
 from qwak.exceptions import QwakException, QwakGrpcAddressException
-from .grpc_auth import Auth0Client
+from qwak.inner.di_configuration.account import UserAccountConfiguration
+from qwak.inner.tool.auth import Auth0ClientBase
+from .grpc_auth import Auth0Client, FrogMLGrpcClient
 
 logger = logging.getLogger()
 HOSTNAME_REGEX: str = r"^(?!-)(?:[A-Za-z0-9-]{1,63}\.)*[A-Za-z0-9-]{1,63}(?<!-)$"
@@ -19,7 +22,7 @@ def create_grpc_channel(
     url: str,
     enable_ssl: bool = True,
     enable_auth: bool = True,
-    auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
+    auth_metadata_plugin: Optional[grpc.AuthMetadataPlugin] = None,
     timeout: int = 100,
     options=None,
     backoff_options=None,
@@ -49,18 +52,9 @@ def create_grpc_channel(
     if not url:
         raise QwakException("Unable to create gRPC channel. URL has not been defined.")
 
-    if enable_ssl or url.endswith(":443"):
-        credentials = grpc.ssl_channel_credentials()
-        if enable_auth:
-            if auth_metadata_plugin is None:
-                auth_metadata_plugin = Auth0Client()
-            credentials = grpc.composite_channel_credentials(
-                credentials, grpc.metadata_call_credentials(auth_metadata_plugin)
-            )
-
-        channel = grpc.secure_channel(url, credentials=credentials, options=options)
-    else:
-        channel = grpc.insecure_channel(url, options=options)
+    auth_metadata_plugin, channel = __get_channel_definition(
+        auth_metadata_plugin, enable_auth, enable_ssl, options, url
+    )
     try:
         interceptors = (
             RetryOnRpcErrorClientInterceptor(
@@ -100,6 +94,35 @@ def create_grpc_channel(
         ) from e
 
 
+def __get_channel_definition(
+    auth_metadata_plugin: Optional[AuthMetadataPlugin],
+    enable_auth: bool,
+    enable_ssl: bool,
+    options,
+    url: str,
+) -> tuple[Optional[AuthMetadataPlugin], Channel]:
+    if enable_ssl or url.endswith(":443"):
+        credentials = grpc.ssl_channel_credentials()
+        if enable_auth:
+            if auth_metadata_plugin is None:
+                user_config = UserAccountConfiguration()
+
+                if issubclass(user_config._auth_client, Auth0ClientBase):
+                    auth_metadata_plugin = Auth0Client()
+                else:
+                    auth_metadata_plugin = FrogMLGrpcClient()
+
+            credentials = grpc.composite_channel_credentials(
+                credentials, grpc.metadata_call_credentials(auth_metadata_plugin)
+            )
+
+        channel = grpc.secure_channel(url, credentials=credentials, options=options)
+    else:
+        channel = grpc.insecure_channel(url, options=options)
+
+    return auth_metadata_plugin, channel
+
+
 def create_grpc_channel_or_none(
     url: str,
     enable_ssl: bool = True,

qwak/inner/tool/grpc/grpc_try_wrapping.py

@@ -1,14 +1,12 @@
-import logging
 import functools
 import inspect
 from inspect import BoundArguments, Signature
 from typing import Any, Callable, Optional
 
 import grpc
+from frogml_storage.logging import logger
 from qwak.exceptions import QwakException
 
-logger = logging.getLogger()
-
 
 def grpc_try_catch_wrapper(
     exception_message: str,

qwak/qwak_client/client.py

@@ -356,6 +356,7 @@ class QwakClient:
         self,
         project_name: str,
         project_description: str,
+        jfrog_project_key: Optional[str] = None,
     ) -> str:
         """
         Create project
@@ -363,6 +364,7 @@ class QwakClient:
         Args:
             project_name (str): The requested name
             project_description (str): The requested description
+            jfrog_project_key (Optional[str]): The requested jfrog project key
 
         Returns:
              str: The project ID of the newly created project
@@ -371,6 +373,7 @@ class QwakClient:
         project = self._get_project_management().create_project(
             project_name=project_name,
             project_description=project_description,
+            jfrog_project_key=jfrog_project_key,
         )
 
         return project.project.project_id
@@ -416,6 +419,7 @@ class QwakClient:
         project_id: str,
         model_name: str,
         model_description: str,
+        jfrog_project_key: Optional[str] = None,
     ) -> str:
         """
         Create model
@@ -424,6 +428,7 @@ class QwakClient:
             project_id (str): The project ID to associate the model
             model_name (str): The requested name
             model_description (str): The requested description
+            jfrog_project_key (Optional[str]): The jfrog project key
 
         Returns:
              str: The model ID of the newly created project
@@ -433,6 +438,7 @@ class QwakClient:
             project_id=project_id,
             model_name=model_name,
             model_description=model_description,
+            jfrog_project_key=jfrog_project_key,
         )
 
         return model.model_id

{qwak_core-0.6.7.dist-info → qwak_core-0.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: qwak-core
-Version: 0.6.7
+Version: 0.7.0
 Summary: Qwak Core contains the necessary objects and communication tools for using the Qwak Platform
 License: Apache-2.0
 Keywords: mlops,ml,deployment,serving,model