blaxel 0.64.0__py3-none-any.whl

blaxel/client.py

@@ -0,0 +1,281 @@
+import ssl
+from typing import Any, Optional, Union
+
+import httpx
+from attrs import define, evolve, field
+
+
+@define
+class Client:
+    """A class for keeping track of data related to the API
+
+    The following are accepted as keyword arguments and will be used to construct httpx Clients internally:
+
+        ``base_url``: The base URL for the API, all requests are made to a relative path to this URL
+
+        ``cookies``: A dictionary of cookies to be sent with every request
+
+        ``headers``: A dictionary of headers to be sent with every request
+
+        ``provider``: An implementation of httpx.Auth to use for authentication
+
+        ``timeout``: The maximum amount of a time a request can take. API functions will raise
+        httpx.TimeoutException if this is exceeded.
+
+        ``verify_ssl``: Whether or not to verify the SSL certificate of the API server. This should be True in production,
+        but can be set to False for testing purposes.
+
+        ``follow_redirects``: Whether or not to follow redirects. Default value is False.
+
+        ``httpx_args``: A dictionary of additional arguments to be passed to the ``httpx.Client`` and ``httpx.AsyncClient`` constructor.
+
+
+    Attributes:
+        raise_on_unexpected_status: Whether or not to raise an errors.UnexpectedStatus if the API returns a
+            status code that was not documented in the source OpenAPI document. Can also be provided as a keyword
+            argument to the constructor.
+
+    """
+
+    raise_on_unexpected_status: bool = field(default=True, kw_only=True)
+    _base_url: str = field(alias="base_url", default="")
+    _cookies: dict[str, str] = field(factory=dict, kw_only=True, alias="cookies")
+    _headers: dict[str, str] = field(factory=dict, kw_only=True, alias="headers")
+    _provider: httpx.Auth = field(default=None, alias="provider")
+    _timeout: Optional[httpx.Timeout] = field(default=None, kw_only=True, alias="timeout")
+    _verify_ssl: Union[str, bool, ssl.SSLContext] = field(default=True, kw_only=True, alias="verify_ssl")
+    _follow_redirects: bool = field(default=False, kw_only=True, alias="follow_redirects")
+    _httpx_args: dict[str, Any] = field(factory=dict, kw_only=True, alias="httpx_args")
+    _client: Optional[httpx.Client] = field(default=None, init=False)
+    _async_client: Optional[httpx.AsyncClient] = field(default=None, init=False)
+
+    def __post_init__(self):
+        from .common.settings import get_settings
+
+        settings = get_settings()
+        self._base_url = settings.base_url
+
+    def with_headers(self, headers: dict[str, str]) -> "Client":
+        """Get a new client matching this one with additional headers"""
+        if self._client is not None:
+            self._client.headers.update(headers)
+        if self._async_client is not None:
+            self._async_client.headers.update(headers)
+        return evolve(self, headers={**self._headers, **headers})
+
+    def with_cookies(self, cookies: dict[str, str]) -> "Client":
+        """Get a new client matching this one with additional cookies"""
+        if self._client is not None:
+            self._client.cookies.update(cookies)
+        if self._async_client is not None:
+            self._async_client.cookies.update(cookies)
+        return evolve(self, cookies={**self._cookies, **cookies})
+
+    def with_timeout(self, timeout: httpx.Timeout) -> "Client":
+        """Get a new client matching this one with a new timeout (in seconds)"""
+        if self._client is not None:
+            self._client.timeout = timeout
+        if self._async_client is not None:
+            self._async_client.timeout = timeout
+        return evolve(self, timeout=timeout)
+
+    def set_httpx_client(self, client: httpx.Client) -> "Client":
+        """Manually set the underlying httpx.Client
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._client = client
+        return self
+
+    def get_httpx_client(self) -> httpx.Client:
+        """Get the underlying httpx.Client, constructing a new one if not previously set"""
+        if self._client is None:
+            self._client = httpx.Client(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                auth=self._provider,
+                **self._httpx_args,
+            )
+        return self._client
+
+    def __enter__(self) -> "Client":
+        """Enter a context manager for self.client—you cannot enter twice (see httpx docs)"""
+        self.get_httpx_client().__enter__()
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for internal httpx.Client (see httpx docs)"""
+        self.get_httpx_client().__exit__(*args, **kwargs)
+
+    def set_async_httpx_client(self, async_client: httpx.AsyncClient) -> "Client":
+        """Manually the underlying httpx.AsyncClient
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._async_client = async_client
+        return self
+
+    def get_async_httpx_client(self) -> httpx.AsyncClient:
+        """Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._async_client
+
+    async def __aenter__(self) -> "Client":
+        """Enter a context manager for underlying httpx.AsyncClient—you cannot enter twice (see httpx docs)"""
+        await self.get_async_httpx_client().__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for underlying httpx.AsyncClient (see httpx docs)"""
+        await self.get_async_httpx_client().__aexit__(*args, **kwargs)
+
+
+@define
+class AuthenticatedClient:
+    """A Client which has been authenticated for use on secured endpoints
+
+    The following are accepted as keyword arguments and will be used to construct httpx Clients internally:
+
+        ``base_url``: The base URL for the API, all requests are made to a relative path to this URL
+
+        ``cookies``: A dictionary of cookies to be sent with every request
+
+        ``headers``: A dictionary of headers to be sent with every request
+
+        ``provider``: An implementation of httpx.Auth to use for authentication
+
+        ``timeout``: The maximum amount of a time a request can take. API functions will raise
+        httpx.TimeoutException if this is exceeded.
+
+        ``verify_ssl``: Whether or not to verify the SSL certificate of the API server. This should be True in production,
+        but can be set to False for testing purposes.
+
+        ``follow_redirects``: Whether or not to follow redirects. Default value is False.
+
+        ``httpx_args``: A dictionary of additional arguments to be passed to the ``httpx.Client`` and ``httpx.AsyncClient`` constructor.
+
+
+    Attributes:
+        raise_on_unexpected_status: Whether or not to raise an errors.UnexpectedStatus if the API returns a
+            status code that was not documented in the source OpenAPI document. Can also be provided as a keyword
+            argument to the constructor.
+        provider: AuthProvider to use for authentication
+    """
+
+    raise_on_unexpected_status: bool = field(default=True, kw_only=True)
+    _base_url: str = field(alias="base_url", default="")
+    _cookies: dict[str, str] = field(factory=dict, kw_only=True, alias="cookies")
+    _headers: dict[str, str] = field(factory=dict, kw_only=True, alias="headers")
+    _provider: httpx.Auth = field(default=None, alias="provider")
+    _timeout: Optional[httpx.Timeout] = field(default=None, kw_only=True, alias="timeout")
+    _verify_ssl: Union[str, bool, ssl.SSLContext] = field(default=True, kw_only=True, alias="verify_ssl")
+    _follow_redirects: bool = field(default=False, kw_only=True, alias="follow_redirects")
+    _httpx_args: dict[str, Any] = field(factory=dict, kw_only=True, alias="httpx_args")
+    _client: Optional[httpx.Client] = field(default=None, init=False)
+    _async_client: Optional[httpx.AsyncClient] = field(default=None, init=False)
+
+    def __post_init__(self):
+        from .common.settings import get_settings
+
+        settings = get_settings()
+        self._base_url = settings.base_url
+
+    def with_headers(self, headers: dict[str, str]) -> "AuthenticatedClient":
+        """Get a new client matching this one with additional headers"""
+        if self._client is not None:
+            self._client.headers.update(headers)
+        if self._async_client is not None:
+            self._async_client.headers.update(headers)
+        return evolve(self, headers={**self._headers, **headers})
+
+    def with_cookies(self, cookies: dict[str, str]) -> "AuthenticatedClient":
+        """Get a new client matching this one with additional cookies"""
+        if self._client is not None:
+            self._client.cookies.update(cookies)
+        if self._async_client is not None:
+            self._async_client.cookies.update(cookies)
+        return evolve(self, cookies={**self._cookies, **cookies})
+
+    def with_timeout(self, timeout: httpx.Timeout) -> "AuthenticatedClient":
+        """Get a new client matching this one with a new timeout (in seconds)"""
+        if self._client is not None:
+            self._client.timeout = timeout
+        if self._async_client is not None:
+            self._async_client.timeout = timeout
+        return evolve(self, timeout=timeout)
+
+    def set_httpx_client(self, client: httpx.Client) -> "AuthenticatedClient":
+        """Manually set the underlying httpx.Client
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._client = client
+        return self
+
+    def get_httpx_client(self) -> httpx.Client:
+        """Get the underlying httpx.Client, constructing a new one if not previously set"""
+        if self._client is None:
+            self._client = httpx.Client(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                auth=self._provider,
+                **self._httpx_args,
+            )
+        return self._client
+
+    def __enter__(self) -> "AuthenticatedClient":
+        """Enter a context manager for self.client—you cannot enter twice (see httpx docs)"""
+        self.get_httpx_client().__enter__()
+        return self
+
+    def __exit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for internal httpx.Client (see httpx docs)"""
+        self.get_httpx_client().__exit__(*args, **kwargs)
+
+    def set_async_httpx_client(self, async_client: httpx.AsyncClient) -> "AuthenticatedClient":
+        """Manually the underlying httpx.AsyncClient
+
+        **NOTE**: This will override any other settings on the client, including cookies, headers, and timeout.
+        """
+        self._async_client = async_client
+        return self
+
+    def get_async_httpx_client(self) -> httpx.AsyncClient:
+        """Get the underlying httpx.AsyncClient, constructing a new one if not previously set"""
+        if self._async_client is None:
+            self._async_client = httpx.AsyncClient(
+                base_url=self._base_url,
+                cookies=self._cookies,
+                headers=self._headers,
+                timeout=self._timeout,
+                verify=self._verify_ssl,
+                follow_redirects=self._follow_redirects,
+                **self._httpx_args,
+            )
+        return self._async_client
+
+    async def __aenter__(self) -> "AuthenticatedClient":
+        """Enter a context manager for underlying httpx.AsyncClient—you cannot enter twice (see httpx docs)"""
+        await self.get_async_httpx_client().__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        """Exit a context manager for underlying httpx.AsyncClient (see httpx docs)"""
+        await self.get_async_httpx_client().__aexit__(*args, **kwargs)

blaxel/common/__init__.py

@@ -0,0 +1,17 @@
+from .error import HTTPError
+from .logger import init as init_logger
+from .secrets import Secret
+from .settings import Settings, get_settings, init
+from .slugify import slugify
+from .utils import copy_folder
+
+__all__ = [
+    "Secret",
+    "Settings",
+    "get_settings",
+    "init",
+    "copy_folder",
+    "init_logger",
+    "HTTPError",
+    "slugify"
+]

blaxel/common/error.py

@@ -0,0 +1,27 @@
+"""
+This module defines custom exception classes used for handling HTTP-related errors within Blaxel.
+"""
+
+class HTTPError(Exception):
+    """
+    A custom exception class for HTTP errors.
+
+    Attributes:
+        status_code (int): The HTTP status code associated with the error.
+        message (str): A descriptive message explaining the error.
+    """
+
+    def __init__(self, status_code: int, message: str):
+        self.status_code = status_code
+        self.message = message
+        super().__init__(self.message)
+
+    def __str__(self):
+        """
+        Returns a string representation of the HTTPError.
+
+        Returns:
+            str: A string in the format "status_code message".
+        """
+        return f"{self.status_code} {self.message}"
+

blaxel/common/instrumentation.py

@@ -0,0 +1,317 @@
+"""
+This module provides utilities for setting up and managing OpenTelemetry instrumentation within Blaxel.
+It includes classes and functions for configuring tracers, meters, loggers, and integrating with FastAPI applications.
+"""
+
+import importlib
+import logging
+from typing import Any, Optional, Type
+
+from fastapi import FastAPI
+from opentelemetry import _logs, metrics, trace
+from opentelemetry._logs import set_logger_provider
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor  # type: ignore
+from opentelemetry.metrics import NoOpMeterProvider
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace import NoOpTracerProvider
+from typing_extensions import Dict
+
+from blaxel.authentication import get_authentication_headers
+
+from .settings import get_settings
+
+tracer: trace.Tracer | None = None
+meter: metrics.Meter | None = None
+logger: LoggerProvider | None = None
+
+log = logging.getLogger(__name__)
+
+
+def auth_headers() -> Dict[str, str]:
+    """
+    Retrieves authentication headers based on the current settings.
+
+    Returns:
+        Dict[str, str]: A dictionary containing authentication headers.
+    """
+    settings = get_settings()
+    headers = get_authentication_headers(settings)
+    return {
+        "x-blaxel-authorization": headers.get("X-Blaxel-Authorization", ""),
+        "x-blaxel-workspace": headers.get("X-Blaxel-Workspace", ""),
+    }
+
+
+def get_logger() -> LoggerProvider:
+    """
+    Retrieves the current logger provider.
+
+    Returns:
+        LoggerProvider: The active logger provider.
+
+    Raises:
+        Exception: If the logger has not been initialized.
+    """
+    if logger is None:
+        raise Exception("Logger is not initialized")
+    return logger
+
+
+def get_resource_attributes() -> Dict[str, Any]:
+    resources = Resource.create()
+    resources_dict: Dict[str, Any] = {}
+    for key in resources.attributes:
+        resources_dict[key] = resources.attributes[key]
+    settings = get_settings()
+    resources_dict["workspace"] = settings.workspace
+    resources_dict["service.name"] = settings.name
+    return resources_dict
+
+
+def get_metrics_exporter() -> OTLPMetricExporter | None:
+    settings = get_settings()
+    if not settings.enable_opentelemetry:
+        return None
+    return OTLPMetricExporter(headers=auth_headers())
+
+
+def get_span_exporter() -> OTLPSpanExporter | None:
+    settings = get_settings()
+    if not settings.enable_opentelemetry:
+        return None
+    return OTLPSpanExporter(headers=auth_headers())
+
+
+def get_log_exporter() -> OTLPLogExporter | None:
+    settings = get_settings()
+    if not settings.enable_opentelemetry:
+        return None
+    return OTLPLogExporter(headers=auth_headers())
+
+
+def _import_class(module_path: str, class_name: str) -> Optional[Type]:  # type: ignore
+    """Dynamically import a class from a module path."""
+    try:
+        module = importlib.import_module(module_path)
+        return getattr(module, class_name)
+    except (ImportError, AttributeError) as e:
+        log.error(f"Could not import {class_name} from {module_path}: {str(e)}")
+        return None
+
+
+# Define mapping of instrumentor info: (module path, class name, required package)
+INSTRUMENTOR_CONFIGS = {
+    "httpx": (
+        "opentelemetry.instrumentation.httpx",
+        "HTTPXClientInstrumentor",
+        "httpx",
+    ),
+    "anthropic": (
+        "opentelemetry.instrumentation.anthropic",
+        "AnthropicInstrumentor",
+        "anthropic",
+    ),
+    "chroma": (
+        "opentelemetry.instrumentation.chroma",
+        "ChromaInstrumentor",
+        "chromadb",
+    ),
+    "cohere": (
+        "opentelemetry.instrumentation.cohere",
+        "CohereInstrumentor",
+        "cohere",
+    ),
+    "groq": ("opentelemetry.instrumentation.groq", "GroqInstrumentor", "groq"),
+    "lance": (
+        "opentelemetry.instrumentation.lance",
+        "LanceInstrumentor",
+        "pylance",
+    ),
+    "langchain": (
+        "opentelemetry.instrumentation.langchain",
+        "LangchainInstrumentor",
+        "langchain",
+    ),
+    "llama_index": (
+        "opentelemetry.instrumentation.llama_index",
+        "LlamaIndexInstrumentor",
+        "llama_index",
+    ),
+    "marqo": (
+        "opentelemetry.instrumentation.marqo",
+        "MarqoInstrumentor",
+        "marqo",
+    ),
+    "milvus": (
+        "opentelemetry.instrumentation.milvus",
+        "MilvusInstrumentor",
+        "pymilvus",
+    ),
+    "mistralai": (
+        "opentelemetry.instrumentation.mistralai",
+        "MistralAiInstrumentor",
+        "mistralai",
+    ),
+    "ollama": (
+        "opentelemetry.instrumentation.ollama",
+        "OllamaInstrumentor",
+        "ollama",
+    ),
+    "openai": (
+        "opentelemetry.instrumentation.openai",
+        "OpenAIInstrumentor",
+        "openai",
+    ),
+    "pinecone": (
+        "opentelemetry.instrumentation.pinecone",
+        "PineconeInstrumentor",
+        "pinecone",
+    ),
+    "qdrant": (
+        "opentelemetry.instrumentation.qdrant",
+        "QdrantInstrumentor",
+        "qdrant_client",
+    ),
+    "replicate": (
+        "opentelemetry.instrumentation.replicate",
+        "ReplicateInstrumentor",
+        "replicate",
+    ),
+    "together": (
+        "opentelemetry.instrumentation.together",
+        "TogetherAiInstrumentor",
+        "together",
+    ),
+    "watsonx": (
+        "opentelemetry.instrumentation.watsonx",
+        "WatsonxInstrumentor",
+        "ibm_watson_machine_learning",
+    ),
+    "weaviate": (
+        "opentelemetry.instrumentation.weaviate",
+        "WeaviateInstrumentor",
+        "weaviate",
+    ),
+}
+
+
+def _is_package_installed(package_name: str) -> bool:
+    """Check if a package is installed."""
+    try:
+        importlib.import_module(package_name)
+        return True
+    except (ImportError, ModuleNotFoundError):
+        return False
+
+
+def instrument_app(app: FastAPI):
+    """
+    Instruments the given FastAPI application with OpenTelemetry.
+
+    This includes setting up tracer and meter providers, configuring exporters, and instrumenting
+    various modules based on available packages.
+
+    Parameters:
+        app (FastAPI): The FastAPI application to instrument.
+    """
+    global tracer
+    global meter
+    settings = get_settings()
+    if not settings.enable_opentelemetry:
+        # Use NoOp implementations to stub tracing and metrics
+        trace.set_tracer_provider(NoOpTracerProvider())
+        tracer = trace.get_tracer(__name__)
+
+        metrics.set_meter_provider(NoOpMeterProvider())
+        meter = metrics.get_meter(__name__)
+        return
+
+    resource = Resource.create(
+        {
+            "service.name": settings.name,
+            "service.namespace": settings.workspace,
+            "service.workspace": settings.workspace,
+        }
+    )
+
+    # Set up the TracerProvider if not already set
+    if not isinstance(trace.get_tracer_provider(), TracerProvider):
+        trace_provider = TracerProvider(resource=resource)
+        span_processor = BatchSpanProcessor(get_span_exporter())  # type: ignore
+        trace_provider.add_span_processor(span_processor)
+        trace.set_tracer_provider(trace_provider)
+        tracer = trace_provider.get_tracer(__name__)
+    else:
+        tracer = trace.get_tracer(__name__)
+
+    # Set up the MeterProvider if not already set
+    if not isinstance(metrics.get_meter_provider(), MeterProvider):
+        metrics_exporter = PeriodicExportingMetricReader(get_metrics_exporter())  # type: ignore
+        meter_provider = MeterProvider(
+            resource=resource, metric_readers=[metrics_exporter]
+        )
+        metrics.set_meter_provider(meter_provider)
+        meter = meter_provider.get_meter(__name__)
+    else:
+        meter = metrics.get_meter(__name__)
+
+    if not isinstance(_logs.get_logger_provider(), LoggerProvider):
+        logger_provider = LoggerProvider(resource=resource)
+        set_logger_provider(logger_provider)
+        logger_provider.add_log_record_processor(
+            BatchLogRecordProcessor(get_log_exporter())  # type: ignore
+        )
+        handler = LoggingHandler(
+            level=logging.NOTSET, logger_provider=logger_provider
+        )
+        logging.getLogger().addHandler(handler)
+    else:
+        logger_provider = _logs.get_logger_provider()
+
+    # Only instrument the app when OpenTelemetry is enabled
+    FastAPIInstrumentor.instrument_app(app)  # type: ignore
+
+    for name, (
+        module_path,
+        class_name,
+        required_package,
+    ) in INSTRUMENTOR_CONFIGS.items():
+        if _is_package_installed(required_package):
+            instrumentor_class = _import_class(module_path, class_name)  # type: ignore
+            if instrumentor_class:
+                try:
+                    instrumentor_class().instrument()
+                    log.debug(f"Successfully instrumented {name}")
+                except Exception as e:
+                    log.debug(f"Failed to instrument {name}: {str(e)}")
+            else:
+                log.debug(f"Could not load instrumentor for {name}")
+        else:
+            log.debug(
+                f"Skipping {name} instrumentation - required package '{required_package}' not installed"
+            )
+
+
+def shutdown_instrumentation():
+    """
+    Shuts down the OpenTelemetry instrumentation providers gracefully.
+
+    This ensures that all spans and metrics are properly exported before the application exits.
+    """
+    if tracer is not None:
+        trace_provider = trace.get_tracer_provider()
+        if isinstance(trace_provider, TracerProvider):
+            trace_provider.shutdown()
+    if meter is not None:
+        meter_provider = metrics.get_meter_provider()
+        if isinstance(meter_provider, MeterProvider):
+            meter_provider.shutdown()