blaxel 0.64.0__py3-none-any.whl

blaxel/models/websocket_channel.py

@@ -0,0 +1,88 @@
+from typing import Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="WebsocketChannel")
+
+
+@_attrs_define
+class WebsocketChannel:
+    """WebSocket connection details
+
+    Attributes:
+        created_at (Union[Unset, str]): The date and time when the resource was created
+        updated_at (Union[Unset, str]): The date and time when the resource was updated
+        connection_id (Union[Unset, str]): Unique connection ID
+        workspace (Union[Unset, str]): Workspace the connection belongs to
+    """
+
+    created_at: Union[Unset, str] = UNSET
+    updated_at: Union[Unset, str] = UNSET
+    connection_id: Union[Unset, str] = UNSET
+    workspace: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        created_at = self.created_at
+
+        updated_at = self.updated_at
+
+        connection_id = self.connection_id
+
+        workspace = self.workspace
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if created_at is not UNSET:
+            field_dict["createdAt"] = created_at
+        if updated_at is not UNSET:
+            field_dict["updatedAt"] = updated_at
+        if connection_id is not UNSET:
+            field_dict["connection_id"] = connection_id
+        if workspace is not UNSET:
+            field_dict["workspace"] = workspace
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        created_at = d.pop("createdAt", UNSET)
+
+        updated_at = d.pop("updatedAt", UNSET)
+
+        connection_id = d.pop("connection_id", UNSET)
+
+        workspace = d.pop("workspace", UNSET)
+
+        websocket_channel = cls(
+            created_at=created_at,
+            updated_at=updated_at,
+            connection_id=connection_id,
+            workspace=workspace,
+        )
+
+        websocket_channel.additional_properties = d
+        return websocket_channel
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/models/workspace.py

@@ -0,0 +1,148 @@
+from typing import TYPE_CHECKING, Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+if TYPE_CHECKING:
+    from ..models.workspace_labels import WorkspaceLabels
+
+
+T = TypeVar("T", bound="Workspace")
+
+
+@_attrs_define
+class Workspace:
+    """Workspace
+
+    Attributes:
+        created_at (Union[Unset, str]): The date and time when the resource was created
+        updated_at (Union[Unset, str]): The date and time when the resource was updated
+        created_by (Union[Unset, str]): The user or service account who created the resource
+        updated_by (Union[Unset, str]): The user or service account who updated the resource
+        account_id (Union[Unset, str]): Workspace account id
+        display_name (Union[Unset, str]): Workspace display name
+        labels (Union[Unset, WorkspaceLabels]): Workspace labels
+        name (Union[Unset, str]): Workspace name
+        region (Union[Unset, str]): Workspace write region
+    """
+
+    created_at: Union[Unset, str] = UNSET
+    updated_at: Union[Unset, str] = UNSET
+    created_by: Union[Unset, str] = UNSET
+    updated_by: Union[Unset, str] = UNSET
+    account_id: Union[Unset, str] = UNSET
+    display_name: Union[Unset, str] = UNSET
+    labels: Union[Unset, "WorkspaceLabels"] = UNSET
+    name: Union[Unset, str] = UNSET
+    region: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        created_at = self.created_at
+
+        updated_at = self.updated_at
+
+        created_by = self.created_by
+
+        updated_by = self.updated_by
+
+        account_id = self.account_id
+
+        display_name = self.display_name
+
+        labels: Union[Unset, dict[str, Any]] = UNSET
+        if self.labels and not isinstance(self.labels, Unset) and not isinstance(self.labels, dict):
+            labels = self.labels.to_dict()
+        elif self.labels and isinstance(self.labels, dict):
+            labels = self.labels
+
+        name = self.name
+
+        region = self.region
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if created_at is not UNSET:
+            field_dict["createdAt"] = created_at
+        if updated_at is not UNSET:
+            field_dict["updatedAt"] = updated_at
+        if created_by is not UNSET:
+            field_dict["createdBy"] = created_by
+        if updated_by is not UNSET:
+            field_dict["updatedBy"] = updated_by
+        if account_id is not UNSET:
+            field_dict["accountId"] = account_id
+        if display_name is not UNSET:
+            field_dict["displayName"] = display_name
+        if labels is not UNSET:
+            field_dict["labels"] = labels
+        if name is not UNSET:
+            field_dict["name"] = name
+        if region is not UNSET:
+            field_dict["region"] = region
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        from ..models.workspace_labels import WorkspaceLabels
+
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        created_at = d.pop("createdAt", UNSET)
+
+        updated_at = d.pop("updatedAt", UNSET)
+
+        created_by = d.pop("createdBy", UNSET)
+
+        updated_by = d.pop("updatedBy", UNSET)
+
+        account_id = d.pop("accountId", UNSET)
+
+        display_name = d.pop("displayName", UNSET)
+
+        _labels = d.pop("labels", UNSET)
+        labels: Union[Unset, WorkspaceLabels]
+        if isinstance(_labels, Unset):
+            labels = UNSET
+        else:
+            labels = WorkspaceLabels.from_dict(_labels)
+
+        name = d.pop("name", UNSET)
+
+        region = d.pop("region", UNSET)
+
+        workspace = cls(
+            created_at=created_at,
+            updated_at=updated_at,
+            created_by=created_by,
+            updated_by=updated_by,
+            account_id=account_id,
+            display_name=display_name,
+            labels=labels,
+            name=name,
+            region=region,
+        )
+
+        workspace.additional_properties = d
+        return workspace
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/models/workspace_labels.py

@@ -0,0 +1,45 @@
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="WorkspaceLabels")
+
+
+@_attrs_define
+class WorkspaceLabels:
+    """Workspace labels"""
+
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        workspace_labels = cls()
+
+        workspace_labels.additional_properties = d
+        return workspace_labels
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/models/workspace_user.py

@@ -0,0 +1,115 @@
+from typing import Any, TypeVar, Union
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="WorkspaceUser")
+
+
+@_attrs_define
+class WorkspaceUser:
+    """Workspace user
+
+    Attributes:
+        accepted (Union[Unset, bool]): Whether the user has accepted the workspace invitation
+        email (Union[Unset, str]): Workspace user email
+        email_verified (Union[Unset, bool]): Whether the user's email has been verified
+        family_name (Union[Unset, str]): Workspace user family name
+        given_name (Union[Unset, str]): Workspace user given name
+        role (Union[Unset, str]): Workspace user role
+        sub (Union[Unset, str]): Workspace user identifier
+    """
+
+    accepted: Union[Unset, bool] = UNSET
+    email: Union[Unset, str] = UNSET
+    email_verified: Union[Unset, bool] = UNSET
+    family_name: Union[Unset, str] = UNSET
+    given_name: Union[Unset, str] = UNSET
+    role: Union[Unset, str] = UNSET
+    sub: Union[Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        accepted = self.accepted
+
+        email = self.email
+
+        email_verified = self.email_verified
+
+        family_name = self.family_name
+
+        given_name = self.given_name
+
+        role = self.role
+
+        sub = self.sub
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if accepted is not UNSET:
+            field_dict["accepted"] = accepted
+        if email is not UNSET:
+            field_dict["email"] = email
+        if email_verified is not UNSET:
+            field_dict["email_verified"] = email_verified
+        if family_name is not UNSET:
+            field_dict["family_name"] = family_name
+        if given_name is not UNSET:
+            field_dict["given_name"] = given_name
+        if role is not UNSET:
+            field_dict["role"] = role
+        if sub is not UNSET:
+            field_dict["sub"] = sub
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: dict[str, Any]) -> T:
+        if not src_dict:
+            return None
+        d = src_dict.copy()
+        accepted = d.pop("accepted", UNSET)
+
+        email = d.pop("email", UNSET)
+
+        email_verified = d.pop("email_verified", UNSET)
+
+        family_name = d.pop("family_name", UNSET)
+
+        given_name = d.pop("given_name", UNSET)
+
+        role = d.pop("role", UNSET)
+
+        sub = d.pop("sub", UNSET)
+
+        workspace_user = cls(
+            accepted=accepted,
+            email=email,
+            email_verified=email_verified,
+            family_name=family_name,
+            given_name=given_name,
+            role=role,
+            sub=sub,
+        )
+
+        workspace_user.additional_properties = d
+        return workspace_user
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

blaxel/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561

blaxel/run.py

@@ -0,0 +1,108 @@
+"""
+This module provides functionality for executing HTTP requests against Blaxel resources.
+"""
+import urllib.parse
+from typing import Any
+
+import requests
+
+from blaxel.client import AuthenticatedClient
+from blaxel.common import HTTPError, get_settings
+
+
+class RunClient:
+    """Provides functionality for executing HTTP requests against Blaxel resources.
+
+    This module contains the RunClient class which handles authenticated HTTP requests to Blaxel
+    resources. It allows users to interact with different resource types (like functions or services), supporting various HTTP methods and request parameters.
+
+    Example:
+        ```python
+        client = new_client()
+        run_client = RunClient(client)
+        response = run_client.run(
+            resource_type="function",
+            resource_name="my-function",
+            method="POST",
+            json={"key": "value"}
+        )
+        ```
+
+    Args:
+        client (AuthenticatedClient): An authenticated client instance for making HTTP requests.
+    """
+
+    def __init__(self, client: AuthenticatedClient):
+        self.client = client
+
+    def run(
+        self,
+        resource_type: str,
+        resource_name: str,
+        method: str,
+        path: str = "",
+        headers: dict[str, str] | None = None,
+        json: dict[str, Any] | None = None,
+        data: str | None = None,
+        params: dict[str, str] | None = None,
+        cloud: bool = False,
+        service_name: str | None = None,
+    ) -> requests.Response:
+        """Execute an HTTP request against a Blaxel resource.
+
+        Args:
+            resource_type (str): The type of resource to interact with (e.g., 'function', 'service').
+            resource_name (str): The name of the specific resource.
+            method (str): The HTTP method to use (e.g., 'GET', 'POST', 'PUT', 'DELETE').
+            path (str, optional): Additional path segments to append to the resource URL. Defaults to "".
+            headers (dict[str, str] | None, optional): HTTP headers to include in the request. Defaults to None.
+            json (dict[str, Any] | None, optional): JSON payload to send with the request. Defaults to None.
+            data (str | None, optional): Raw data to send with the request. Defaults to None.
+            params (dict[str, str] | None, optional): Query parameters to include in the URL. Defaults to None.
+            cloud (bool, optional): Whether to use the cloud endpoint. Defaults to False.
+            service_name (str | None, optional): The name of the service to use. Defaults to None.
+
+        Returns:
+            requests.Response: The HTTP response from the server.
+
+        Raises:
+            HTTPError: If the server responds with a status code >= 400.
+        """
+        settings = get_settings()
+        headers = headers or {}
+        params = params or {}
+
+        if cloud and path and service_name:
+            url = f"https://{service_name}.{settings.run_internal_hostname}/{path}"
+
+        if cloud and not path and service_name:
+            url = f"https://{service_name}.{settings.run_internal_hostname}"
+
+        if not cloud and path:
+            url = urllib.parse.urljoin(settings.run_url, f"{settings.workspace}/{resource_type}s/{resource_name}/{path}")
+
+        if not cloud and not path:
+            url = urllib.parse.urljoin(settings.run_url, f"{settings.workspace}/{resource_type}s/{resource_name}")
+
+        kwargs = {
+            "headers": headers,
+            "params": {**params},
+        }
+        if data:
+            kwargs["data"] = data
+        if json:
+            kwargs["json"] = json
+
+        response = self.client.get_httpx_client().request(method, url, **kwargs)
+        if response.status_code >= 400 and not cloud:
+            raise HTTPError(response.status_code, response.text)
+        if response.status_code >= 400 and cloud: # Redirect to the public endpoint if the resource is in the cloud and the request fails
+            if path:
+                url = urllib.parse.urljoin(settings.run_url, f"{settings.workspace}/{resource_type}s/{resource_name}/{path}")
+            else:
+                url = urllib.parse.urljoin(settings.run_url, f"{settings.workspace}/{resource_type}s/{resource_name}")
+            response = self.client.get_httpx_client().request(method, url, **kwargs)
+            if response.status_code >= 400:
+                raise HTTPError(response.status_code, response.text)
+            return response
+        return response

blaxel/serve/app.py

@@ -0,0 +1,131 @@
+"""Module: app
+
+This module sets up and runs the Blaxel server using FastAPI.
+It configures middleware, handles server lifespan events, and defines endpoints.
+"""
+
+import asyncio
+import importlib
+import inspect
+import os
+import sys
+import traceback
+from contextlib import asynccontextmanager
+from logging import getLogger
+from uuid import uuid4
+
+from asgi_correlation_id import CorrelationIdMiddleware
+from fastapi import FastAPI, Request, Response, WebSocket
+from fastapi.responses import JSONResponse
+
+from blaxel.aimon import HTTPError, get_settings, init
+from blaxel.aimon.instrumentation import instrument_app, shutdown_instrumentation
+
+from .middlewares import AccessLogMiddleware, AddProcessTimeHeader
+
+sys.path.insert(0, os.getcwd())
+sys.path.insert(0, os.path.join(os.getcwd(), "src"))
+
+
+def import_module():
+    """Dynamically imports the main server module based on settings.
+
+    Returns:
+        Callable: The main function to run the server.
+    """
+    settings = get_settings()
+    main_module = importlib.import_module(".".join(settings.server.module.split(".")[0:-1]))
+    func = getattr(main_module, settings.server.module.split(".")[-1])
+    return func
+
+
+settings = init()
+logger = getLogger(__name__)
+logger.info(f"Importing server module: {settings.server.module}")
+func = import_module()
+logger.info(
+    f"Running server"
+    f" on {settings.server.host}:{settings.server.port}"
+)
+func_params = inspect.signature(func).parameters
+websocket_detected = False
+if "websocket" in func_params:
+    websocket_detected = True
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Manages the lifespan events of the FastAPI application.
+
+    Args:
+        app (FastAPI): The FastAPI application instance.
+    """
+    yield
+    shutdown_instrumentation()
+
+
+app = FastAPI(docs_url=None, redoc_url=None, lifespan=lifespan)
+app.add_middleware(
+    CorrelationIdMiddleware,
+    header_name="x-blaxel-request-id",
+    generator=lambda: str(uuid4()),
+)
+app.add_middleware(AddProcessTimeHeader)
+app.add_middleware(AccessLogMiddleware)
+instrument_app(app)
+
+@app.get("/health")
+async def health():
+    """Health check endpoint.
+
+    Returns:
+        dict: A simple status message indicating the server is running.
+    """
+    return {"status": "ok"}
+
+if websocket_detected:
+    @app.websocket("/ws")
+    async def websocket_endpoint(websocket: WebSocket):
+        await websocket.accept()
+
+        original_func = getattr(func, "__wrapped__", func)
+        if asyncio.iscoroutinefunction(func) or asyncio.iscoroutinefunction(original_func):
+            await func(websocket)
+        else:
+            func(websocket)
+
+else:
+    @app.post("/")
+    async def root(request: Request):
+        logger = getLogger(__name__)
+        try:
+            original_func = getattr(func, "__wrapped__", func)
+            if asyncio.iscoroutinefunction(func) or asyncio.iscoroutinefunction(original_func):
+                response = await func(request)
+            else:
+                response = func(request)
+
+            if isinstance(response, Response):
+                return response
+            if type(response) is str:
+                return Response(
+                    content=response,
+                    headers={"Content-Type": "text/plain"},
+                    media_type="text/plain",
+                    status_code=200,
+                )
+            return JSONResponse(status_code=200, content=response)
+        except ValueError as e:
+            content = {"error": str(e)}
+            content["traceback"] = str(traceback.format_exc())
+            logger.error(str(traceback.format_exc()))
+            return JSONResponse(status_code=400, content=content)
+        except HTTPError as e:
+            content = {"error": e.message, "status_code": e.status_code}
+            content["traceback"] = str(traceback.format_exc())
+            logger.error(f"{e.status_code} {str(traceback.format_exc())}")
+            return JSONResponse(status_code=e.status_code, content=content)
+        except Exception as e:
+            content = {"error": f"Internal server error, {e}"}
+            content["traceback"] = str(traceback.format_exc())
+            logger.error(str(traceback.format_exc()))
+            return JSONResponse(status_code=500, content=content)

blaxel/serve/middlewares/__init__.py

@@ -0,0 +1,10 @@
+"""Package: middlewares
+
+This package contains custom middleware classes for the Blaxel server,
+including access logging and process time header addition.
+"""
+
+from .accesslog import AccessLogMiddleware
+from .processtime import AddProcessTimeHeader
+
+__all__ = ["AccessLogMiddleware", "AddProcessTimeHeader"]

blaxel/serve/middlewares/accesslog.py

@@ -0,0 +1,32 @@
+"""Module: accesslog
+
+Defines the AccessLogMiddleware for logging HTTP requests and responses.
+"""
+
+from logging import getLogger
+
+from starlette.middleware.base import BaseHTTPMiddleware
+
+
+class AccessLogMiddleware(BaseHTTPMiddleware):
+    """Middleware for logging access information of each HTTP request."""
+
+    async def dispatch(self, request, call_next):
+        """Processes each request to log access details.
+
+        Args:
+            request (Request): The incoming HTTP request.
+            call_next (Callable): The next middleware or endpoint handler.
+
+        Returns:
+            Response: The HTTP response generated by the next handler.
+        """
+        logger = getLogger(__name__)
+        response = await call_next(request)
+        process_time = response.headers.get("X-Process-Time")
+        rid_header = response.headers.get("X-Request-Id")
+        request_id = rid_header or response.headers.get("X-Blaxel-Request-Id")
+        logger.info(
+            f"{request.method} {request.url.path} {response.status_code} {process_time}ms rid={request_id}"
+        )
+        return response