kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/client/pquery.py

@@ -0,0 +1,207 @@
+import copy
+from dataclasses import fields, make_dataclass
+from http import HTTPStatus
+from typing import List, Optional, Tuple, Union, get_args, get_origin
+
+from kumoapi.common import ValidationResponse
+from kumoapi.json_serde import to_json_dict
+from kumoapi.model_plan import (
+    MissingType,
+    ModelPlan,
+    ModelPlanInfo,
+    PlanMixin,
+    SuggestModelPlanRequest,
+    TrainingTableGenerationPlan,
+)
+from kumoapi.pquery import PQueryResource, QueryType
+from kumoapi.task import TaskType
+from kumoapi.typing import WITH_PYDANTIC_V2
+from pydantic.dataclasses import dataclass
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import PQueryEndpoints
+from kumoai.client.utils import (
+    Returns,
+    parse_id_response,
+    parse_response,
+    raise_on_error,
+)
+
+GraphID = str
+PQueryID = str
+
+
+class PQueryAPI:
+    r"""Typed API definition for Kumo Predictive Query definition."""
+    def __init__(self, client: KumoClient) -> None:
+        self._client = client
+
+    def create(self, pquery: PQueryResource) -> PQueryID:
+        r"""Creates a Predictive Query resource object in Kumo."""
+        resp = self._client._request(PQueryEndpoints.create,
+                                     json=to_json_dict(pquery))
+        raise_on_error(resp)
+        return parse_id_response(resp)
+
+    def get_if_exists(self, name_or_id: str) -> Optional[PQueryResource]:
+        r"""Fetches a Predictive Query given its name(id)."""
+        resp = self._client._request(PQueryEndpoints.get.with_id(name_or_id))
+        if resp.status_code == HTTPStatus.NOT_FOUND:
+            return None
+
+        raise_on_error(resp)
+        return parse_response(PQueryResource, resp)
+
+    def list(self, *,
+             name_pattern: Optional[str] = None) -> List[PQueryResource]:
+        params = {}
+        if name_pattern:
+            params['name_pattern'] = name_pattern
+        resp = self._client._request(PQueryEndpoints.list, params=params)
+        raise_on_error(resp)
+        return parse_response(List[PQueryResource], resp)
+
+    def delete(self, name: str) -> None:
+        resp = self._client._request(PQueryEndpoints.delete.with_id(name))
+        raise_on_error(resp)
+
+    def infer_task_type(
+        self,
+        pquery_string: str,
+        graph_id: GraphID,
+    ) -> Tuple[TaskType, str]:
+        resp = self._client._request(
+            PQueryEndpoints.infer_task_type,
+            params={
+                'pquery_string': pquery_string,
+                'graph_id': graph_id
+            },
+        )
+        raise_on_error(resp)
+        return parse_response(Returns[Tuple[TaskType, str]], resp)
+
+    def validate(self, pquery: PQueryResource) -> ValidationResponse:
+        r"""Validates a Predictive Query resource object in Kumo."""
+        resp = self._client._request(PQueryEndpoints.validate,
+                                     json=to_json_dict(pquery))
+        raise_on_error(resp)
+        return parse_response(Returns[ValidationResponse], resp)
+
+    def suggest_training_table_plan(
+        self,
+        req: SuggestModelPlanRequest,
+    ) -> TrainingTableGenerationPlan:
+        resp = self._client._request(
+            PQueryEndpoints.suggest_training_table_plan,
+            json=to_json_dict(req),
+        )
+        raise_on_error(resp)
+        return parse_response(TrainingTableGenerationPlan, resp)
+
+    def suggest_model_plan(
+        self,
+        req: SuggestModelPlanRequest,
+    ) -> ModelPlan:
+        resp = self._client._request(
+            PQueryEndpoints.suggest_model_plan,
+            json=to_json_dict(req),
+        )
+        raise_on_error(resp)
+        plan_info = parse_response(ModelPlanInfo, resp)
+
+        return filter_model_plan(
+            plan_info.model_plan,
+            plan_info.task_type,
+            plan_info.query_type,
+            plan_info.has_train_table_weight_col,
+        )
+
+
+def filter_model_plan(
+    plan: ModelPlan,
+    task_type: TaskType,
+    query_type: QueryType,
+    has_train_table_weight_col: bool,
+) -> ModelPlan:
+    r"""Filters the model plan ``plan`` down to its valid fields only, given
+    ``task_type`` and ``query_type``.
+    """
+    new_section_fields = []
+    new_sections = []
+    for section_field in fields(plan):
+        section = getattr(plan, section_field.name)
+
+        new_opt_fields = []
+        new_opts = []
+        for field in fields(section):
+            if WITH_PYDANTIC_V2:
+                metadata = field.default.json_schema_extra  # type: ignore
+            else:
+                metadata = field.default.extra['metadata']  # type: ignore
+            if not section.is_valid_option(
+                    name=field.name,
+                    metadata=metadata,
+                    task_type=task_type,
+                    query_type=query_type,
+                    has_train_table_weight_col=has_train_table_weight_col,
+            ):
+                continue
+
+            # Remove `MissingType` from type annotation:
+            if WITH_PYDANTIC_V2:
+                hidden = metadata["hidden"]
+            else:
+                hidden = metadata.hidden
+            if not hidden and get_origin(field.type) is Union:
+                _types = tuple(_type for _type in get_args(field.type)
+                               if _type is not MissingType)
+                assert len(_types) > 0
+                _type = Union[_types] if len(_types) > 1 else _types[0]
+            else:
+                _type = field.type
+
+            default = copy.copy(field.default)
+            if not WITH_PYDANTIC_V2:
+                # Pydantic v1
+                from pydantic.fields import Undefined
+                default.default = Undefined  # type: ignore
+            else:
+                try:
+                    # Pydantic v2 - Undefined moved to pydantic_core
+                    from pydantic_core import PydanticUndefined
+                    default.default = PydanticUndefined  # type: ignore
+                except ImportError:
+                    # Fallback: Neither v1 nor v2 Undefined available
+                    # In Pydantic v2, not setting a default is equivalent to
+                    # Undefined
+                    pass
+
+            # Forward compatibility - Remove any newly introduced arguments not
+            # returned yet by the backend:
+            value = getattr(section, field.name)
+            if value != MissingType.VALUE:
+                new_opt_fields.append((field.name, _type, default))
+                new_opts.append(value)
+
+        Section = dataclass(
+            config=dict(validate_assignment=True),
+            repr=False,
+        )(make_dataclass(
+            type(section).__name__,
+            new_opt_fields,
+            bases=(PlanMixin, ),
+            repr=False,
+        ))
+        new_section_fields.append((section_field.name, Section))
+        new_sections.append(Section(*new_opts))
+
+    Plan = dataclass(
+        config=dict(validate_assignment=True),
+        repr=False,
+    )(make_dataclass(
+        'ModelPlan',
+        new_section_fields,
+        bases=(ModelPlan, ),
+        repr=False,
+    ))
+    return Plan(*new_sections)

kumoai/client/rfm.py

@@ -0,0 +1,112 @@
+from typing import Any
+
+from kumoapi.json_serde import to_json_dict
+from kumoapi.rfm import (
+    RFMEvaluateResponse,
+    RFMExplanationResponse,
+    RFMParseQueryRequest,
+    RFMParseQueryResponse,
+    RFMPredictResponse,
+    RFMValidateQueryRequest,
+    RFMValidateQueryResponse,
+)
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import RFMEndpoints
+from kumoai.client.utils import parse_response, raise_on_error
+
+
+class RFMAPI:
+    r"""Typed API definition for Kumo RFM (Relational Foundation Model)."""
+    def __init__(self, client: KumoClient) -> None:
+        self._client = client
+
+    def predict(self, request: bytes) -> RFMPredictResponse:
+        """Make predictions using the RFM model.
+
+        Args:
+            request: The predict request as serialized protobuf.
+
+        Returns:
+            RFMPredictResponse containing the predictions
+        """
+        response = self._client._request(
+            RFMEndpoints.predict,
+            data=request,
+            headers={'Content-Type': 'application/x-protobuf'},
+        )
+        raise_on_error(response)
+        return parse_response(RFMPredictResponse, response)
+
+    def explain(
+        self,
+        request: bytes,
+        skip_summary: bool = False,
+    ) -> RFMExplanationResponse:
+        """Explain the RFM model on the given context.
+
+        Args:
+            request: The predict request as serialized protobuf.
+            skip_summary: Whether to skip generating a human-readable summary
+                of the explanation.
+
+        Returns:
+            RFMPredictResponse containing the explanations
+        """
+        params: dict[str, Any] = {'generate_summary': not skip_summary}
+        response = self._client._request(
+            RFMEndpoints.explain, data=request, params=params,
+            headers={'Content-Type': 'application/x-protobuf'})
+        raise_on_error(response)
+        return parse_response(RFMExplanationResponse, response)
+
+    def evaluate(self, request: bytes) -> RFMEvaluateResponse:
+        """Evaluate the RFM model on the given context.
+
+        Args:
+            request: The evaluate request as serialized protobuf.
+
+        Returns:
+            RFMEvaluateResponse containing the computed metrics
+        """
+        response = self._client._request(
+            RFMEndpoints.evaluate, data=request,
+            headers={'Content-Type': 'application/x-protobuf'})
+        raise_on_error(response)
+        return parse_response(RFMEvaluateResponse, response)
+
+    def validate_query(
+        self,
+        request: RFMValidateQueryRequest,
+    ) -> RFMValidateQueryResponse:
+        """Validate a predictive query against a graph.
+
+        Args:
+            request: The request object containing
+                the query and graph definition
+
+        Returns:
+            RFMValidateQueryResponse containing the QueryDefinition
+        """
+        response = self._client._request(RFMEndpoints.validate_query,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(RFMValidateQueryResponse, response)
+
+    def parse_query(
+        self,
+        request: RFMParseQueryRequest,
+    ) -> RFMParseQueryResponse:
+        """Validate a predictive query against a graph.
+
+        Args:
+            request: The request object containing
+                the query and graph definition
+
+        Returns:
+            RFMParseQueryResponse containing the QueryDefinition
+        """
+        response = self._client._request(RFMEndpoints.parse_query,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(RFMParseQueryResponse, response)

kumoai/client/source_table.py

@@ -0,0 +1,53 @@
+from typing import List
+
+from kumoapi.json_serde import to_json_dict
+from kumoapi.source_table import (
+    SourceTableConfigRequest,
+    SourceTableConfigResponse,
+    SourceTableDataRequest,
+    SourceTableDataResponse,
+    SourceTableListRequest,
+    SourceTableListResponse,
+    SourceTableValidateRequest,
+    SourceTableValidateResponse,
+)
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import SourceTableEndpoints
+from kumoai.client.utils import parse_response, raise_on_error
+
+
+class SourceTableAPI:
+    r"""Typed API definition for Kumo source tables."""
+    def __init__(self, client: KumoClient) -> None:
+        self._client = client
+
+    def validate_table(
+            self, request: SourceTableValidateRequest
+    ) -> SourceTableValidateResponse:
+        response = self._client._request(SourceTableEndpoints.validate_table,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(SourceTableValidateResponse, response)
+
+    def list_tables(
+            self, request: SourceTableListRequest) -> SourceTableListResponse:
+        response = self._client._request(SourceTableEndpoints.list_tables,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(SourceTableListResponse, response)
+
+    def get_table_data(
+            self, request: SourceTableDataRequest) -> SourceTableDataResponse:
+        response = self._client._request(SourceTableEndpoints.get_table_data,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(List[SourceTableDataResponse], response)
+
+    def get_table_config(
+            self,
+            request: SourceTableConfigRequest) -> SourceTableConfigResponse:
+        response = self._client._request(SourceTableEndpoints.get_table_config,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(SourceTableConfigResponse, response)

kumoai/client/table.py

@@ -0,0 +1,101 @@
+from http import HTTPStatus
+from typing import Any, Dict, Optional
+
+from kumoapi.common import ValidationResponse
+from kumoapi.data_snapshot import TableSnapshotID, TableSnapshotResource
+from kumoapi.json_serde import to_json_dict
+from kumoapi.table import (
+    TableDefinition,
+    TableMetadataRequest,
+    TableMetadataResponse,
+    TableResource,
+    TableValidationRequest,
+)
+
+from kumoai.client import KumoClient
+from kumoai.client.endpoints import TableEndpoints
+from kumoai.client.utils import (
+    parse_id_response,
+    parse_response,
+    raise_on_error,
+)
+
+TableID = str
+
+
+class TableAPI:
+    r"""Typed API definition for Kumo table definition."""
+    def __init__(self, client: KumoClient) -> None:
+        self._client = client
+
+    def create_table(
+        self,
+        table_def: TableDefinition,
+        *,
+        name_alias: Optional[str] = None,
+        force_rename: bool = False,
+    ) -> TableID:
+        r"""Creates a Table (metadata definition) resource object in Kumo."""
+        params: Dict[str, Any] = {'force_rename': force_rename}
+        if name_alias:
+            params['name_alias'] = name_alias
+        resp = self._client._request(
+            TableEndpoints.create,
+            params=params,
+            json=to_json_dict(table_def),
+        )
+        raise_on_error(resp)
+        return parse_id_response(resp)
+
+    def get_table_if_exists(
+        self,
+        table_id_or_name: str,
+    ) -> Optional[TableResource]:
+        r"""Fetches a connector given its ID."""
+        resp = self._client._request(
+            TableEndpoints.get.with_id(table_id_or_name))
+        if resp.status_code == HTTPStatus.NOT_FOUND:
+            return None
+
+        raise_on_error(resp)
+        return parse_response(TableResource, resp)
+
+    def create_snapshot(
+        self,
+        table_definition: TableDefinition,
+        *,
+        refresh_source: bool = False,
+    ) -> TableSnapshotID:
+        params: Dict[str, Any] = {'refresh_source': refresh_source}
+        resp = self._client._request(TableEndpoints.create_snapshot,
+                                     params=params,
+                                     json=to_json_dict(table_definition))
+        raise_on_error(resp)
+        return parse_id_response(resp)
+
+    def get_snapshot(
+        self,
+        snapshot_id: TableSnapshotID,
+    ) -> TableSnapshotResource:
+        resp = self._client._request(
+            TableEndpoints.get_snapshot.with_id(snapshot_id))
+        raise_on_error(resp)
+        return parse_response(TableSnapshotResource, resp)
+
+    def infer_metadata(
+        self,
+        request: TableMetadataRequest,
+    ) -> TableMetadataResponse:
+        response = self._client._request(TableEndpoints.infer_metadata,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(TableMetadataResponse, response)
+
+    def validate_table(
+        self,
+        request: TableValidationRequest,
+    ) -> ValidationResponse:
+        response = self._client._request(TableEndpoints.validate,
+                                         json=to_json_dict(request))
+        raise_on_error(response)
+        return parse_response(ValidationResponse, response)

kumoai/client/utils.py

@@ -0,0 +1,130 @@
+import logging
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Generic,
+    List,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+    overload,
+)
+
+import requests
+from kumoapi.json_serde import from_json
+
+from kumoai.exceptions import HTTPException
+
+logger = logging.getLogger(__name__)
+
+_DataclassT = TypeVar('_DataclassT')
+ResponseT = TypeVar('ResponseT')
+
+
+# See https://github.com/python/mypy/issues/9003
+class Returns(Generic[ResponseT]):
+    r"""Use this class if you'd like to use Python types as annotations for
+    mypy, where `Type` is expected. See `parse_response` as an example.
+    """
+    if not TYPE_CHECKING:
+
+        def __class_getitem__(cls, item: object) -> object:
+            """Called when Returns is used as an Annotation at runtime.
+            We just return the type we're given
+            """
+            return item
+
+
+def _parse_dataclass_list(
+    data_class: Type[_DataclassT],
+    response: requests.Response,
+) -> List[_DataclassT]:
+    json_response = response.json()
+    assert isinstance(json_response, list)
+
+    def _parse_elem(v: Any) -> _DataclassT:
+        if issubclass(data_class, (str, int)):
+            return cast(_DataclassT, data_class(v))
+        if data_class == bool:
+            return cast(_DataclassT, str(v).lower() == 'true')
+        return from_json(v, data_class)
+
+    return [_parse_elem(item) for item in json_response]
+
+
+@overload
+def parse_response(response_type: Type[Returns[ResponseT]],
+                   response: requests.Response) -> ResponseT:
+    ...
+
+
+@overload
+def parse_response(response_type: Type[ResponseT],
+                   response: requests.Response) -> ResponseT:
+    ...
+
+
+def parse_response(response_type: Union[Type[ResponseT],
+                                        Type[Returns[ResponseT]]],
+                   response: requests.Response) -> ResponseT:
+    """Parse the HTTP response body into a Pydantic data structure.
+
+    Args:
+        response_type: a class definition for a Pydantic dataclass, or a
+            :class:`list` of Pydantic dataclasses.
+        response: a requests.Response object containing
+            json data in the body
+
+    Returns:
+        The parsed response.
+    """
+    # Case 0: no parsing needed
+    if response_type == str:
+        return response.text  # type: ignore
+    if response_type == bool:
+        return response.text.lower() == 'true'  # type: ignore
+
+    # Case 1: dict[str, V] type
+    if response_type == dict:
+        k_type, v_type = get_args(response_type)
+        assert issubclass(k_type, str)
+        data = response.json()
+        assert isinstance(data, dict)
+        return {
+            k: from_json(v, v_type)
+            for k, v in data.items()
+        }  # type: ignore
+
+    # Case 2: response is a list of pydantic data classes
+    if get_origin(response_type) is list:
+        list_inner_type = get_args(response_type)[0]
+        result = _parse_dataclass_list(list_inner_type, response)
+        return cast(ResponseT, result)
+
+    # Case 3: response is pydantic data classes
+    try:
+        return from_json(response.text, response_type)
+    except Exception:
+        logger.error(f"Unable to parse response {response.text}")
+        raise
+
+
+def parse_id_response(response: requests.Response) -> str:
+    return parse_response(Dict[str, str], response)['id']
+
+
+def parse_patch_response(response: requests.Response) -> bool:
+    """Parse PATCH response to indicate whether the resource was modified."""
+    return parse_response(Dict[str, bool], response)['resource_updated']
+
+
+def raise_on_error(response: requests.Response) -> None:
+    r"""Raises an :class:`~kumoai.exceptions.HTTPException` if a response does
+    not return with an OK status code.
+    """
+    if not response.ok:
+        raise HTTPException(response.status_code, response.text)

kumoai/codegen/__init__.py

@@ -0,0 +1,19 @@
+"""Kumo SDK Code Generation Utility.
+
+Generates Python SDK code from Kumo UI entities.
+Supports both ID-based loading and JSON-based loading.
+"""
+
+from .generate import generate_code
+from .exceptions import (
+    CodegenError,
+    CyclicDependencyError,
+    UnsupportedEntityError,
+)
+
+__all__ = [
+    "generate_code",
+    "CodegenError",
+    "CyclicDependencyError",
+    "UnsupportedEntityError",
+]