fundamental-client 0.2.3__py3-none-any.whl

fundamental/utils/http.py

@@ -0,0 +1,294 @@
+"""Simple HTTP utilities for NEXUS client."""
+
+import json as json_lib
+import logging
+import random
+import time
+from typing import Any, Dict, Optional
+
+import httpx
+
+from fundamental.clients.base import BaseClient
+from fundamental.config import Config
+from fundamental.constants import DEFAULT_TIMEOUT_SECONDS, SIGV4_SERVICE_NAME
+from fundamental.exceptions import (
+    AuthenticationError,
+    AuthorizationError,
+    HTTPError,
+    NetworkError,
+    NotFoundError,
+    RateLimitError,
+    RequestTimeoutError,
+    ServerError,
+    ValidationError,
+)
+
+logger = logging.getLogger(__name__)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+
+
+def _sign_request_sigv4(
+    method: str,
+    url: str,
+    headers: Dict[str, str],
+    body: Optional[bytes],
+    region: str,
+) -> Dict[str, str]:
+    try:
+        import boto3  # pyright: ignore[reportMissingImports]
+        from botocore.auth import SigV4Auth  # pyright: ignore[reportMissingImports]
+        from botocore.awsrequest import AWSRequest  # pyright: ignore[reportMissingImports]
+    except ImportError as e:
+        raise ImportError(
+            "boto3 is required for EC2/SigV4 authentication. "
+            'Install it with: pip install "fundamental-client[ec2]" '
+            'or: uv add "fundamental-client[ec2]"'
+        ) from e
+
+    session = boto3.Session()
+    credentials = session.get_credentials()
+
+    if credentials is None:
+        raise AuthenticationError(
+            message="AWS credentials not found. Configure AWS credentials via "
+            "environment variables, AWS config file, or IAM role."
+        )
+
+    aws_request = AWSRequest(method=method, url=url, headers=headers, data=body or b"")
+
+    SigV4Auth(credentials, SIGV4_SERVICE_NAME, region).add_auth(aws_request)
+
+    return dict(aws_request.headers)
+
+
+def _prepare_request(
+    method: str,
+    url: str,
+    client_config: Config,
+    trace_dict: Optional[Dict[str, str]] = None,
+    extra_headers: Optional[Dict[str, str]] = None,
+    content: Optional[bytes] = None,
+    json: Optional[Dict[str, Any]] = None,
+) -> tuple[Dict[str, str], Optional[bytes]]:
+    """Prepare headers and body for an API request.
+
+    Returns:
+        Tuple of (headers, body_bytes)
+    """
+    headers: Dict[str, str] = {
+        "User-Agent": f"NEXUS-Client-SDK/{client_config.get_version()}",
+    }
+
+    if trace_dict:
+        trace_id = trace_dict.get("trace_id")
+        span_id = trace_dict.get("span_id")
+        if trace_id:
+            headers["x-datadog-trace-id"] = trace_id
+        if span_id:
+            headers["x-datadog-parent-id"] = span_id
+
+    if extra_headers:
+        headers.update(extra_headers)
+
+    # Serialize JSON to bytes
+    body = content
+    if json is not None:
+        body = json_lib.dumps(json).encode("utf-8")
+        headers["Content-Type"] = "application/json"
+
+    if client_config.use_sigv4_auth:
+        headers = _sign_request_sigv4(
+            method=method,
+            url=url,
+            headers=headers,
+            body=body,
+            region=client_config.aws_region,
+        )
+    else:
+        headers["x-api-key"] = f"{client_config.get_api_key()}"
+
+    return headers, body
+
+
+def _calculate_backoff(attempt: int) -> float:
+    """Calculate exponential backoff delay with jitter."""
+    delay = min(2**attempt, 10)  # Cap at 10 seconds
+    jitter = random.uniform(0.1, 0.3) * delay
+    return float(delay + jitter)
+
+
+def _handle_response_error(response: httpx.Response, trace_id: Optional[str] = None) -> None:
+    """Handle HTTP error responses by raising appropriate exceptions."""
+    try:
+        error_data = response.json()
+        detail = error_data.get("detail", response.text[:200])
+    except ValueError:
+        detail = response.text[:200] or f"HTTP error {response.status_code}"
+
+    if response.status_code == 400:
+        raise ValidationError(
+            message=f"Invalid request. Error: {detail}",
+            trace_id=trace_id,
+        )
+    if response.status_code == 401:
+        raise AuthenticationError(
+            message="Invalid API key. Check your credentials.",
+            trace_id=trace_id,
+        )
+    if response.status_code == 403:
+        raise AuthorizationError(
+            message="Access denied. Check your API permissions.",
+            trace_id=trace_id,
+        )
+    if response.status_code == 404:
+        raise NotFoundError(
+            message=f"Resource not found. Error: {detail}",
+            trace_id=trace_id,
+        )
+    if response.status_code == 429:
+        retry_after = response.headers.get("Retry-After")
+        raise RateLimitError(
+            message=f"Rate limit exceeded. Retry after {retry_after} seconds.",
+            trace_id=trace_id,
+        )
+    if response.status_code >= 500:
+        raise ServerError(
+            message="Internal server error.",
+            trace_id=trace_id,
+        )
+    http_error = HTTPError(message=f"Request failed: {detail}", trace_id=trace_id)
+    http_error.status_code = response.status_code
+    raise http_error
+
+
+def api_call(
+    method: str,
+    full_url: str,
+    client: BaseClient,
+    files: Optional[Dict] = None,
+    content: Optional[bytes] = None,
+    data: Optional[Dict[str, Any]] = None,
+    json: Optional[Dict[str, Any]] = None,
+    headers: Optional[Dict[str, str]] = None,
+    timeout: Optional[float] = DEFAULT_TIMEOUT_SECONDS,
+    max_retries: Optional[int] = None,
+) -> httpx.Response:
+    """Make HTTP request to the api.
+
+    Args:
+        method: HTTP method (GET, POST, etc.)
+        full_url: Complete URL for the request
+        client: Client instance
+        files: Files to upload
+        content: Raw content bytes
+        data: Form data
+        json: JSON payload
+        headers: Additional headers
+        timeout: Request timeout in seconds
+        max_retries: Number of retries
+
+    Returns:
+        httpx.Response object
+
+    Raises:
+        Various HTTP exceptions based on response status
+    """
+    client_config = client.config
+    trace_dict = client.get_trace_dict()
+    if max_retries is None:
+        max_retries = client_config.retries
+
+    # Only add auth for requests to our API, not external URLs (e.g., S3)
+    is_api_request = full_url.startswith(client_config.api_url)
+    merged_headers = headers or {}
+
+    if is_api_request:
+        # _prepare_request serializes JSON to bytes (content) so SigV4 can sign the exact body.
+        # For API key auth this is harmless - we're just doing what httpx does behind the scenes.
+        # We set json=None so httpx doesn't re-serialize it.
+        merged_headers, content = _prepare_request(
+            method=method,
+            url=full_url,
+            client_config=client_config,
+            trace_dict=trace_dict,
+            extra_headers=headers,
+            content=content,
+            json=json,
+        )
+        json = None
+
+    trace_id = trace_dict.get("trace_id") if trace_dict else None
+    with httpx.Client(
+        timeout=httpx.Timeout(timeout, connect=5.0),
+        follow_redirects=True,
+    ) as http_client:
+        for attempt in range(max_retries + 1):
+            try:
+                logger.debug(f"Attempt {attempt + 1} of {max_retries + 1} to {method} {full_url}")
+                response = http_client.request(
+                    method=method,
+                    url=full_url,
+                    files=files,
+                    headers=merged_headers,
+                    content=content,
+                    data=data,
+                    json=json,
+                )
+
+                if 200 <= response.status_code < 300:
+                    return response
+
+                # Handle non-retryable errors immediately
+                if 400 <= response.status_code < 500 and response.status_code != 429:
+                    _handle_response_error(response, trace_id=trace_id)
+
+                # Handle retryable errors (429, 5xx)
+                if attempt < max_retries and response.status_code in (
+                    429,
+                    500,
+                    502,
+                    503,
+                    504,
+                ):
+                    wait_time = _calculate_backoff(attempt)
+
+                    # Use Retry-After header for rate limiting
+                    if response.status_code == 429:
+                        retry_after = response.headers.get("Retry-After")
+                        if retry_after:
+                            try:
+                                wait_time = float(retry_after)
+                            except ValueError:
+                                pass
+
+                    time.sleep(wait_time)
+                    continue
+
+                # Final attempt failed
+                _handle_response_error(response, trace_id=trace_id)
+
+            except (httpx.TimeoutException, httpx.RequestError) as e:
+                if attempt < max_retries:
+                    wait_time = _calculate_backoff(attempt)
+                    logger.debug(f"Network error, retrying in {wait_time:.1f}s")
+                    time.sleep(wait_time)
+                    continue
+
+                # Final attempt failed
+                if isinstance(e, httpx.TimeoutException):
+                    raise RequestTimeoutError(
+                        message="Request timed out. Check your connection or try again.",
+                        trace_id=trace_id,
+                    ) from e
+                if "Connection refused" in str(e):
+                    raise NetworkError(
+                        message="Connection refused, cannot connect to api server. \n"
+                        "Check your connection or try again.",
+                        trace_id=trace_id,
+                    ) from e
+                raise NetworkError(
+                    message="Network error occurred. Please try again.",
+                    trace_id=trace_id,
+                ) from e
+
+    raise RuntimeError("Network error occurred. Please try again.")

fundamental/utils/polling.py

@@ -0,0 +1,97 @@
+"""
+Polling utilities for task status checks.
+"""
+
+import time
+
+from fundamental.clients.base import BaseClient
+from fundamental.exceptions import (
+    NEXUSError,
+    RequestTimeoutError,
+    ServerError,
+)
+from fundamental.models import TaskStatus, TaskStatusResponse
+from fundamental.utils.http import api_call
+
+
+def _raise_exception_with_trace(
+    exception_class: type[NEXUSError],
+    message: str,
+    client: BaseClient,
+) -> None:
+    trace_dict = client.get_trace_dict()
+    trace_id = trace_dict.get("trace_id") if trace_dict else None
+    raise exception_class(message=message, trace_id=trace_id)
+
+
+def wait_for_task_status(
+    client: BaseClient,
+    status_url: str,
+    timeout: float,
+    polling_interval: float,
+    polling_requests_without_delay: int = 0,
+    wait_for_completion: bool = True,
+) -> TaskStatusResponse:
+    """
+    Poll for task status until completion or timeout.
+
+    Parameters
+    ----------
+    client : BaseClient
+        The client instance.
+    status_url : str
+        The URL to poll for status.
+    timeout : float
+        Maximum time to wait in seconds.
+    polling_interval : float
+        Time between polling requests in seconds.
+    polling_requests_without_delay : int, default=0
+        Number of initial requests without delay.
+    wait_for_completion : bool, default=True
+        Whether to wait for completion or return immediately on in_progress.
+
+    Returns
+    -------
+    TaskStatusResponse
+        The task status response.
+
+    Raises
+    ------
+    RequestTimeoutError
+        If the request times out.
+    ServerError
+        If the server returns an error.
+    ValidationError
+        If the request fails validation.
+    """
+    start_time = time.perf_counter()
+    request_counter = 0
+
+    while True:
+        if time.perf_counter() - start_time > timeout:
+            _raise_exception_with_trace(
+                RequestTimeoutError,
+                message=f"Request timed out after {timeout}s",
+                client=client,
+            )
+
+        response = api_call(method="GET", full_url=status_url, client=client)
+        status_response = TaskStatusResponse(**response.json())
+
+        if status_response.status == TaskStatus.SUCCESS:
+            return status_response
+
+        if status_response.status == TaskStatus.IN_PROGRESS:
+            if not wait_for_completion:
+                return status_response
+            if request_counter > polling_requests_without_delay:
+                time.sleep(polling_interval)
+        else:
+            # Unexpected state
+            _raise_exception_with_trace(
+                ServerError,
+                message=f"Unexpected task status: {status_response.status}",
+                client=client,
+            )
+
+        request_counter += 1

fundamental/utils/safetensors_deserialize.py

@@ -0,0 +1,98 @@
+"""Safe deserialization for estimator fields using safetensors.
+
+Deserializes estimator_fields.safetensors from the API into a dict.
+This replaces pickle to eliminate RCE vulnerabilities.
+
+Coupling:
+- Input: estimator_fields.safetensors bytes from API (base64 decoded)
+- Producer: Controller's safetensors_serialization.save_estimator_fields()
+- Output: dict matching LTM's sync_fields() structure
+
+Format:
+- Tensors -> numeric numpy arrays
+- Metadata {"value": [...], "type": "string_array"} -> string arrays
+- Metadata {"type": "array_list", "length": N} + field::0, field::1 -> list of arrays
+- Metadata {"value": ...} -> scalars/lists/dicts
+"""
+
+import json
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from safetensors import safe_open
+from safetensors.numpy import load
+
+
+def _deserialize_fields(tensors: dict[str, np.ndarray], metadata: dict[str, str]) -> dict[str, Any]:
+    """Deserialize tensors and metadata back into fields.
+
+    Metadata format:
+    - {"value": ...} -> scalar/list/dict
+    - {"value": [...], "type": "string_array"} -> numpy array (string/object dtype)
+    - {"type": "array_list", "length": N} -> list of arrays (indexed as field::0, field::1, ...)
+
+    Args:
+        tensors: Dictionary of tensor names to numpy arrays
+        metadata: Dictionary of metadata keys to JSON strings
+
+    Returns:
+        Reconstructed fields dictionary
+    """
+    fields: dict[str, Any] = {}
+
+    # Add all tensors that aren't indexed (list items have :: in name)
+    fields = {k: v for k, v in tensors.items() if "::" not in k}
+
+    # Process metadata entries
+    for k, v in metadata.items():
+        if "::" in k:  # Skip indexed list items (processed via parent)
+            continue
+
+        parsed = json.loads(v)
+
+        if parsed.get("type") == "array_list":
+            # Reconstruct list of arrays from indexed entries
+            length = parsed["length"]
+            arrays = []
+            for i in range(length):
+                idx_key = f"{k}::{i}"
+                if idx_key in tensors:
+                    arrays.append(tensors[idx_key])
+                elif idx_key in metadata:
+                    idx_parsed = json.loads(metadata[idx_key])
+                    arrays.append(np.array(idx_parsed["value"]))
+                else:
+                    raise ValueError(f"Missing array_list entry for {idx_key}")
+            fields[k] = arrays
+        elif parsed.get("type") == "string_array":
+            fields[k] = np.array(parsed["value"])
+        else:
+            fields[k] = parsed["value"]
+
+    return fields
+
+
+def load_estimator_fields_from_bytes(data: bytes) -> dict[str, Any]:
+    """Load estimator fields from safetensors bytes.
+
+    Args:
+        data: Bytes containing the safetensors data (from API response)
+
+    Returns:
+        Reconstructed fields dictionary with numpy arrays
+    """
+    tensors = load(data)
+
+    # To get metadata, we need to use safe_open with a file-like object
+    # safetensors supports loading from bytes directly for tensors,
+    # but metadata requires safe_open.
+    # Use a temp directory so Windows can reopen the file without lock issues.
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir) / "estimator_fields.safetensors"
+        tmp_path.write_bytes(data)
+        with safe_open(str(tmp_path), framework="np") as f:
+            metadata = f.metadata() or {}
+
+    return _deserialize_fields(tensors, metadata)