qwak-core 0.4.278__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/inner/tool/grpc/grpc_auth.py

@@ -1,7 +1,6 @@
 from typing import Callable, Tuple
 
 import grpc
-from qwak.inner.const import QwakConstants
 
 _SIGNATURE_HEADER_KEY = "authorization"
 
@@ -35,34 +34,3 @@ 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.
-        """
-        # Get token from FrogML client
-        if not self._auth_client:
-            from qwak.inner.tool.auth import FrogMLAuthClient
-
-            self._auth_client = FrogMLAuthClient()
-        token = self._auth_client.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)

qwak/inner/tool/grpc/grpc_tools.py

@@ -1,17 +1,18 @@
 import logging
+import re
 import time
 from abc import ABC, abstractmethod
 from random import randint
 from typing import Callable, Optional, Tuple
+from urllib.parse import urlparse, ParseResult
 
 import grpc
-from qwak.exceptions import QwakException
-from qwak.inner.di_configuration.account import UserAccountConfiguration
-from qwak.inner.tool.auth import Auth0ClientBase
 
-from .grpc_auth import Auth0Client, FrogMLGrpcClient
+from qwak.exceptions import QwakException, QwakGrpcAddressException
+from .grpc_auth import Auth0Client
 
 logger = logging.getLogger()
+HOSTNAME_REGEX: str = r"^(?!-)(?:[A-Za-z0-9-]{1,63}\.)*[A-Za-z0-9-]{1,63}(?<!-)$"
 
 
 def create_grpc_channel(
@@ -21,7 +22,7 @@ def create_grpc_channel(
     auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
     timeout: int = 100,
     options=None,
-    backoff_options={},
+    backoff_options=None,
     max_attempts=4,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
@@ -42,6 +43,9 @@ def create_grpc_channel(
         status_for_retry: grpc statuses to retry upon
     Returns: Returns a grpc.Channel
     """
+    if backoff_options is None:
+        backoff_options = {}
+
     if not url:
         raise QwakException("Unable to create gRPC channel. URL has not been defined.")
 
@@ -49,11 +53,7 @@ def create_grpc_channel(
         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()
+                auth_metadata_plugin = Auth0Client()
             credentials = grpc.composite_channel_credentials(
                 credentials, grpc.metadata_call_credentials(auth_metadata_plugin)
             )
@@ -107,11 +107,14 @@ def create_grpc_channel_or_none(
     auth_metadata_plugin: grpc.AuthMetadataPlugin = None,
     timeout: int = 30,
     options=None,
-    backoff_options={},
+    backoff_options=None,
     max_attempts=2,
     status_for_retry=(grpc.StatusCode.UNAVAILABLE,),
     attempt=0,
 ) -> Callable[[Optional[str], Optional[bool]], Optional[grpc.Channel]]:
+    if backoff_options is None:
+        backoff_options = {}
+
     def deferred_channel(
         url_overwrite: Optional[str] = None, ssl_overwrite: Optional[bool] = None
     ):
@@ -135,6 +138,117 @@ def create_grpc_channel_or_none(
     return deferred_channel
 
 
+def validate_grpc_address(
+    grpc_address: str,
+    is_port_specification_allowed: bool = False,
+    is_url_scheme_allowed: bool = False,
+):
+    """
+    Validate gRPC address format
+    Args:
+        grpc_address (str): gRPC address to validate
+        is_port_specification_allowed (bool): Whether to allow port specification in the address
+        is_url_scheme_allowed (bool): Whether to allow URL scheme in the address
+    Raises:
+        QwakGrpcAddressException: If the gRPC address is invalid
+    """
+    parsed_grpc_address: ParseResult = parse_address(grpc_address)
+    hostname: str = get_hostname_from_address(parsed_grpc_address)
+    validate_paths_are_not_included_in_address(parsed_grpc_address)
+
+    if not is_url_scheme_allowed:
+        __validate_url_scheme_not_included_in_address(parsed_grpc_address)
+
+    if not is_port_specification_allowed:
+        __validate_port_not_included_in_address(parsed_grpc_address)
+
+    if not is_valid_hostname(hostname):
+        raise QwakGrpcAddressException(
+            "gRPC address must be a simple hostname or fully qualified domain name.",
+            parsed_grpc_address,
+        )
+
+
+def validate_paths_are_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    has_invalid_path: bool = (
+        parsed_grpc_address.path not in {"", "/"}
+        or parsed_grpc_address.query
+        or parsed_grpc_address.fragment
+    )
+
+    if has_invalid_path:
+        raise QwakGrpcAddressException(
+            "gRPC address must not contain paths, queries, or fragments.",
+            parsed_grpc_address,
+        )
+
+
+def get_hostname_from_address(parsed_grpc_address: ParseResult) -> str:
+    hostname: Optional[str] = parsed_grpc_address.hostname
+    if not hostname:
+        raise QwakGrpcAddressException(
+            "gRPC address must contain a valid hostname.", parsed_grpc_address
+        )
+
+    return hostname
+
+
+def __validate_url_scheme_not_included_in_address(
+    parsed_grpc_address: ParseResult,
+) -> None:
+    if parsed_grpc_address.scheme:
+        raise QwakGrpcAddressException(
+            "URL scheme is not allowed in the gRPC address.", parsed_grpc_address
+        )
+
+
+def __validate_port_not_included_in_address(parsed_grpc_address: ParseResult):
+    try:
+        port: Optional[int] = parsed_grpc_address.port
+    except ValueError as exc:
+        raise QwakGrpcAddressException(
+            "Invalid port specification in the gRPC address.", parsed_grpc_address
+        ) from exc
+
+    if port:
+        raise QwakGrpcAddressException(
+            "Port specification is not allowed in the gRPC address.",
+            parsed_grpc_address,
+        )
+
+
+def parse_address(grpc_address: str) -> ParseResult:
+    if not grpc_address or not grpc_address.strip():
+        raise QwakGrpcAddressException(
+            "gRPC address must not be empty or whitespace.", grpc_address
+        )
+
+    trimmed_address: str = grpc_address.strip()
+    parsed_address: ParseResult = urlparse(
+        trimmed_address if "://" in trimmed_address else f"//{trimmed_address}"
+    )
+
+    return parsed_address
+
+
+def is_valid_hostname(hostname: str) -> bool:
+    """
+    Validate that the supplied hostname conforms to RFC-style label rules:
+    anchored pattern enforces full-string validation, negative lookahead/lookbehind block
+    leading or trailing hyphens per label, and each dot-separated label must be 1-63
+    alphanumeric/hyphen characters.
+
+    Args:
+        hostname (str): The hostname to validate.
+    Returns:
+        bool: True if the hostname is valid, False otherwise.
+    """
+    hostname_pattern: re.Pattern = re.compile(HOSTNAME_REGEX)
+    return bool(hostname_pattern.fullmatch(hostname))
+
+
 class SleepingPolicy(ABC):
     @abstractmethod
     def sleep(self, try_i: int):

qwak/inner/tool/grpc/grpc_try_wrapping.py

@@ -1,15 +1,152 @@
+import logging
+import functools
+import inspect
+from inspect import BoundArguments, Signature
+from typing import Any, Callable, Optional
+
 import grpc
 from qwak.exceptions import QwakException
 
+logger = logging.getLogger()
+
+
+def grpc_try_catch_wrapper(
+    exception_message: str,
+    reraise_non_grpc_error_original_exception: Optional[bool] = False,
+) -> Callable:
+    """A decorator for handling exceptions in client methods.
+
+    This decorator wraps a function, catching any `grpc.RpcError` or other
+    `Exception`. it re-raises them as a `QwakException`. The error message can be a
+    static string or a template formatted with the decorated function's
+    runtime arguments.
+
+    Parameters
+    ----------
+    exception_message : str
+        The error message template. This can be a static string or a format
+        string using placeholders that match the decorated function's
+        parameter names (e.g., "Failed on item {item_id}").
+    reraise_non_grpc_error_original_exception : bool, optional
+        If `True`, any exception that is *not* a `grpc.RpcError` will be
+        re-raised in its original form. If `False` (the default), all
+        exceptions are wrapped in `QwakException`.
+
+    Returns
+    -------
+    Callable
+        A decorator that wraps the function with the error handling logic.
+
+    Raises
+    ------
+    QwakException
+        Raised when the decorated function encounters any exception, unless
+        `reraise_non_grpc_error_original_exception` is `True` and the error
+        is not a `grpc.RpcError`. The new exception's message is formatted
+        and includes details from the original error.
+    Exception
+        The original non-gRPC exception is re-raised if
+        `reraise_non_grpc_error_original_exception` is set to `True`.
+
+    Usage
+    -----
+    The decorator can be used with static strings and formatted strings
+
+    **1. Formatted Error Messages**
+
+    Use placeholders matching the function's parameter names to create
+    context-specific error messages.
+
+        @grpc_try_catch_wrapper("Failed to delete version {version_number} for FeatureSet {featureset_name}")
+        def delete_featureset_version(self, featureset_name: str, version_number: int):
+            # ... gRPC call ...
+
+    If this function fails with a gRPC error, a `QwakException` is raised with a message like:
+    `"Failed to delete version 2 for FeatureSet customer_churn - <original gRPC error code> - <original gRPC error details>."`
+
+    **2. Static Error Messages**
+
+    When no dynamic context is needed, a simple string is sufficient.
 
-def grpc_try_catch_wrapper(exception_message):
-    def decorator(function):
+        @grpc_try_catch_wrapper("Failed to retrieve featuresets mapping")
+        def get_featuresets_mapping(self):
+            # ... gRPC call ...
+
+    If this fails, the `QwakException` message will be:
+    `"Failed to retrieve featuresets mapping - <original gRPC error code> - <original gRPC error details>."`
+    """
+
+    def decorator(function: Callable):
+        @functools.wraps(function)
         def _inner_wrapper(*args, **kwargs):
             try:
                 return function(*args, **kwargs)
             except grpc.RpcError as e:
-                raise QwakException(exception_message + f". Error is: {e.details()}.")
+                error_message: str = __get_error_msg(
+                    exception_message=exception_message, e=e
+                )
+
+                formatted_message = __get_formatted_error_message(
+                    exception_message=error_message,
+                    function=function,
+                    args=args,
+                    kwargs=kwargs,
+                )
+
+                raise QwakException(formatted_message) from e
+
+            except Exception as e:
+                if reraise_non_grpc_error_original_exception:
+                    raise e
+
+                formatted_message = __get_formatted_error_message(
+                    exception_message=exception_message,
+                    function=function,
+                    args=args,
+                    kwargs=kwargs,
+                )
+                raise QwakException(f"{formatted_message}. Error is: {e}.") from e
 
         return _inner_wrapper
 
     return decorator
+
+
+def __get_error_msg(exception_message: str, e: grpc.RpcError) -> str:
+    err_msg: str = (
+        f"{exception_message} - {e.code() if hasattr(e, 'code') else 'UNKNOWN'}"
+    )
+    if hasattr(e, "details") and e.details() is not None:
+        err_msg: str = f"{err_msg} - {e.details()}"
+    elif hasattr(e, "debug_error_string"):
+        err_msg: str = f"{err_msg} - {e.debug_error_string()}"
+    logger.debug(f"{err_msg}: {e}")
+    return err_msg
+
+
+def __get_formatted_error_message(
+    exception_message: str, function: Callable, args: tuple[Any], kwargs: dict[str, Any]
+) -> str:
+    """Formats an error message string with the runtime arguments of a function.
+
+    It inspects the function's signature and binds the passed args and kwargs
+    to their parameter names, then uses them to format the message string.
+
+    Args:
+        exception_message (str): The error message template with placeholders.
+        function (Callable): The decorated function.
+        args (tuple[Any]): The positional arguments passed to the function.
+        kwargs (dict[str, Any]): The keyword arguments passed to the function.
+
+    Returns:
+        str: The error message with placeholders filled in.
+    """
+    try:
+        sig: Signature = inspect.signature(function)
+        bound_args: BoundArguments = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+
+        return exception_message.format(**bound_args.arguments)
+    except (ValueError, KeyError, TypeError):
+        # Fallback if formatting fails (e.g., missing key in template)
+        return exception_message

qwak/llmops/generation/chat/openai/types/chat/chat_completion.py

@@ -16,8 +16,6 @@
 from dataclasses import dataclass
 from typing import List, Optional
 
-from typing_extensions import Literal
-
 from qwak.llmops.generation.base import ModelResponse
 from .chat_completion_message import ChatCompletionMessage
 from .chat_completion_token_logprob import ChatCompletionTokenLogprob
@@ -34,9 +32,7 @@ class ChoiceLogprobs:
 
 @dataclass
 class Choice:
-    finish_reason: Literal[
-        "stop", "length", "tool_calls", "content_filter", "function_call"
-    ]
+    finish_reason: str
     """The reason the model stopped generating tokens.
 
     This will be `stop` if the model hit a natural stop point or a provided stop
@@ -55,6 +51,21 @@ class Choice:
     logprobs: Optional[ChoiceLogprobs] = None
     """Log probability information for the choice."""
 
+    def __post_init__(self):
+        """Validates that finish_reason is one of the allowed values."""
+        allowed_reasons = {
+            "stop",
+            "length",
+            "tool_calls",
+            "content_filter",
+            "function_call",
+        }
+        if self.finish_reason not in allowed_reasons:
+            raise ValueError(
+                f"Invalid finish_reason: '{self.finish_reason}'. "
+                f"Must be one of {allowed_reasons}"
+            )
+
 
 @dataclass
 class ChatCompletion(ModelResponse):
@@ -73,7 +84,7 @@ class ChatCompletion(ModelResponse):
     model: str
     """The model used for the chat completion."""
 
-    object: Literal["chat.completion"]
+    object: str
     """The object type, which is always `chat.completion`."""
 
     system_fingerprint: Optional[str] = None
@@ -85,3 +96,10 @@ class ChatCompletion(ModelResponse):
 
     usage: Optional[CompletionUsage] = None
     """Usage statistics for the completion request."""
+
+    def __post_init__(self):
+        """Validates that the object type is correct."""
+        if self.object != "chat.completion":
+            raise ValueError(
+                f"Invalid object type: '{self.object}'. Must be 'chat.completion'"
+            )

qwak/llmops/generation/chat/openai/types/chat/chat_completion_chunk.py

@@ -16,8 +16,6 @@
 from dataclasses import dataclass
 from typing import List, Optional
 
-from typing_extensions import Literal
-
 from qwak.llmops.generation.base import ModelResponse
 from .chat_completion_token_logprob import ChatCompletionTokenLogprob
 
@@ -69,9 +67,16 @@ class ChoiceDeltaToolCall:
 
     function: Optional[ChoiceDeltaToolCallFunction] = None
 
-    type: Optional[Literal["function"]] = None
+    type: Optional[str] = None
     """The type of the tool. Currently, only `function` is supported."""
 
+    def __post_init__(self):
+        """Validates that type is 'function' if present."""
+        if self.type is not None and self.type != "function":
+            raise ValueError(
+                f"Invalid type: '{self.type}'. Must be 'function' or None."
+            )
+
 
 @dataclass
 class ChoiceDelta:
@@ -85,11 +90,21 @@ class ChoiceDelta:
     model.
     """
 
-    role: Optional[Literal["system", "user", "assistant", "tool"]] = None
+    role: Optional[str] = None
     """The role of the author of this message."""
 
     tool_calls: Optional[List[ChoiceDeltaToolCall]] = None
 
+    def __post_init__(self):
+        """Validates that role is one of the allowed values if present."""
+        if self.role is not None:
+            allowed_roles = {"system", "user", "assistant", "tool"}
+            if self.role not in allowed_roles:
+                raise ValueError(
+                    f"Invalid role: '{self.role}'. "
+                    f"Must be one of {allowed_roles} or None."
+                )
+
 
 @dataclass
 class ChoiceLogprobs:
@@ -105,9 +120,7 @@ class Choice:
     index: int
     """The index of the choice in the list of choices."""
 
-    finish_reason: Optional[
-        Literal["stop", "length", "tool_calls", "content_filter", "function_call"]
-    ] = None
+    finish_reason: Optional[str] = None
     """The reason the model stopped generating tokens.
 
     This will be `stop` if the model hit a natural stop point or a provided stop
@@ -120,6 +133,22 @@ class Choice:
     logprobs: Optional[ChoiceLogprobs] = None
     """Log probability information for the choice."""
 
+    def __post_init__(self):
+        """Validates that finish_reason is one of the allowed values if present."""
+        if self.finish_reason is not None:
+            allowed_reasons = {
+                "stop",
+                "length",
+                "tool_calls",
+                "content_filter",
+                "function_call",
+            }
+            if self.finish_reason not in allowed_reasons:
+                raise ValueError(
+                    f"Invalid finish_reason: '{self.finish_reason}'. "
+                    f"Must be one of {allowed_reasons} or None."
+                )
+
 
 @dataclass
 class ChatCompletionChunk(ModelResponse):
@@ -141,7 +170,7 @@ class ChatCompletionChunk(ModelResponse):
     model: str
     """The model to generate the completion."""
 
-    object: Literal["chat.completion.chunk"]
+    object: str
     """The object type, which is always `chat.completion.chunk`."""
 
     system_fingerprint: Optional[str] = None
@@ -150,3 +179,10 @@ class ChatCompletionChunk(ModelResponse):
     Can be used in conjunction with the `seed` request parameter to understand when
     backend changes have been made that might impact determinism.
     """
+
+    def __post_init__(self):
+        """Validates that the object type is correct."""
+        if self.object != "chat.completion.chunk":
+            raise ValueError(
+                f"Invalid object type: '{self.object}'. Must be 'chat.completion.chunk'"
+            )

qwak/llmops/generation/chat/openai/types/chat/chat_completion_message.py

@@ -16,8 +16,6 @@
 from dataclasses import dataclass
 from typing import List, Optional
 
-from typing_extensions import Literal
-
 from .chat_completion_message_tool_call import ChatCompletionMessageToolCall
 
 __all__ = ["ChatCompletionMessage", "FunctionCall"]
@@ -39,7 +37,7 @@ class FunctionCall:
 
 @dataclass
 class ChatCompletionMessage:
-    role: Literal["assistant"]
+    role: str
     """The role of the author of this message."""
 
     content: Optional[str] = None
@@ -54,3 +52,8 @@ class ChatCompletionMessage:
 
     tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
     """The tool calls generated by the model, such as function calls."""
+
+    def __post_init__(self):
+        """Validates that the role type is correct."""
+        if self.role != "assistant":
+            raise ValueError(f"Invalid object type: '{self.role}'. Must be 'assistant'")

qwak/qwak_client/client.py

@@ -8,7 +8,6 @@ from _qwak_proto.qwak.build.v1.build_pb2 import DESCRIPTOR
 from qwak.clients.analytics.client import AnalyticsEngineClient
 from qwak.clients.automation_management.client import AutomationsManagementClient
 from qwak.clients.batch_job_management.client import BatchJobManagerClient
-from qwak.clients.build_management.client import BuildsManagementClient
 from qwak.clients.build_orchestrator.client import BuildFilter, BuildOrchestratorClient
 from qwak.clients.data_versioning.client import DataVersioningManagementClient
 from qwak.clients.data_versioning.data_tag_filter import DataTagFilter
@@ -58,10 +57,6 @@ class QwakClient:
     def _get_model_management(self):
         return ModelsManagementClient()
 
-    @lru_cache(maxsize=1)
-    def _get_build_management(self):
-        return BuildsManagementClient()
-
     @lru_cache(maxsize=1)
     def _get_build_orchestrator(self):
         return BuildOrchestratorClient()
@@ -243,8 +238,8 @@ class QwakClient:
         model = self._get_model_management().get_model(model_id=model_id)
         model_uuid = model.uuid
 
-        metric_filters = list()
-        parameter_filters = list()
+        metric_filters = []
+        parameter_filters = []
 
         for build_filter in filters:
             if isinstance(build_filter, MetricFilter):
@@ -361,7 +356,6 @@ class QwakClient:
         self,
         project_name: str,
         project_description: str,
-        jfrog_project_key: Optional[str] = None,
     ) -> str:
         """
         Create project
@@ -369,7 +363,6 @@ 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
@@ -378,7 +371,6 @@ 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
@@ -424,7 +416,6 @@ class QwakClient:
         project_id: str,
         model_name: str,
         model_description: str,
-        jfrog_project_key: Optional[str] = None,
     ) -> str:
         """
         Create model
@@ -433,7 +424,6 @@ 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
@@ -443,7 +433,6 @@ 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/vector_store/rest_helpers.py

@@ -4,25 +4,16 @@ import socket
 from datetime import datetime
 
 import requests
-from qwak.inner.di_configuration.account import UserAccountConfiguration
-from qwak.inner.tool.auth import Auth0ClientBase, FrogMLAuthClient
-from qwak.inner.const import QwakConstants
+from qwak.inner.tool.auth import Auth0ClientBase
 
 
 def _get_authorization():
-    user_account_configuration = UserAccountConfiguration()
-    if issubclass(user_account_configuration._auth_client, Auth0ClientBase):
-        auth_client = Auth0ClientBase()
-        tenant_id = None
-    else:
-        auth_client = FrogMLAuthClient()
-        tenant_id = auth_client.get_tenant_id()
-
+    auth_client = Auth0ClientBase()
     token = auth_client.get_token()
     token_split = token.split(".")
     decoded_token = json.loads(_base64url_decode(token_split[1]).decode("utf-8"))
     token_expiration = datetime.fromtimestamp(decoded_token["exp"])
-    return f"Bearer {token}", token_expiration, tenant_id
+    return f"Bearer {token}", token_expiration
 
 
 def _base64url_decode(input):
@@ -77,8 +68,5 @@ class RestSession(requests.Session):
         return super().prepare_request(request)
 
     def prepare_request_token(self):
-        auth_token, self.jwt_expiration, tenant_id = _get_authorization()
+        auth_token, self.jwt_expiration = _get_authorization()
         self.headers["Authorization"] = auth_token
-
-        if tenant_id:
-            self.headers[QwakConstants.JFROG_TENANT_HEADER_KEY] = tenant_id