futurehouse-client 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

futurehouse_client/clients/job_client.py

@@ -14,6 +14,10 @@ from futurehouse_client.models.rest import (
     StoreAgentStatePostRequest,
     StoreEnvironmentFrameRequest,
 )
+from futurehouse_client.utils.monitoring import (
+    external_trace,
+    insert_distributed_trace_headers,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -90,14 +94,32 @@ class JobClient:
         data = FinalEnvironmentRequest(status=status)
         try:
             async with httpx.AsyncClient(timeout=self.REQUEST_TIMEOUT) as client:
-                response = await client.patch(
-                    url=f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}/environment-frame",
-                    json=data.model_dump(mode="json"),
-                    headers={
-                        "Authorization": f"Bearer {self.oauth_jwt}",
-                        "x-trajectory-id": self.trajectory_id,
+                url = f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}/environment-frame"
+                headers = {
+                    "Authorization": f"Bearer {self.oauth_jwt}",
+                    "x-trajectory-id": self.trajectory_id,
+                }
+
+                with external_trace(
+                    url=url,
+                    method="PATCH",
+                    library="httpx",
+                    custom_params={
+                        "trajectory_id": self.trajectory_id,
+                        "agent": self.agent,
+                        "environment": self.environment,
+                        "status": status,
+                        "operation": "finalize_environment",
                     },
-                )
+                ):
+                    headers = insert_distributed_trace_headers(headers)
+
+                    response = await client.patch(
+                        url=url,
+                        json=data.model_dump(mode="json"),
+                        headers=headers,
+                    )
+
                 response.raise_for_status()
                 logger.debug(f"Environment updated with status {status}")
         except httpx.HTTPStatusError:
@@ -145,14 +167,32 @@ class JobClient:
 
         try:
             async with httpx.AsyncClient(timeout=self.REQUEST_TIMEOUT) as client:
-                response = await client.post(
-                    url=f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}/agent-state",
-                    json=data.model_dump(mode="json"),
-                    headers={
-                        "Authorization": f"Bearer {self.oauth_jwt}",
-                        "x-trajectory-id": self.trajectory_id,
+                url = f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}/agent-state"
+                headers = {
+                    "Authorization": f"Bearer {self.oauth_jwt}",
+                    "x-trajectory-id": self.trajectory_id,
+                }
+
+                with external_trace(
+                    url=url,
+                    method="POST",
+                    library="httpx",
+                    custom_params={
+                        "trajectory_id": self.trajectory_id,
+                        "agent": self.agent,
+                        "environment": self.environment,
+                        "step": step,
+                        "timestep": self.current_timestep,
+                        "operation": "store_agent_state",
                     },
-                )
+                ):
+                    headers = insert_distributed_trace_headers(headers)
+
+                    response = await client.post(
+                        url=url,
+                        json=data.model_dump(mode="json"),
+                        headers=headers,
+                    )
                 response.raise_for_status()
                 logger.info(f"Successfully stored agent state for step {step}")
                 return response.json()
@@ -198,14 +238,34 @@ class JobClient:
 
         try:
             async with httpx.AsyncClient(timeout=self.REQUEST_TIMEOUT) as client:
-                response = await client.post(
-                    url=f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}/environment-frame",
-                    json=data.model_dump(mode="json"),
-                    headers={
-                        "Authorization": f"Bearer {self.oauth_jwt}",
-                        "x-trajectory-id": self.trajectory_id,
-                    },
-                )
+                url = f"{self.base_uri}/v0.1/trajectories/{self.trajectory_id}/environment-frame"
+                headers = {
+                    "Authorization": f"Bearer {self.oauth_jwt}",
+                    "x-trajectory-id": self.trajectory_id,
+                }
+
+                custom_params = {
+                    "trajectory_id": self.trajectory_id,
+                    "agent": self.agent,
+                    "environment": self.environment,
+                    "timestep": self.current_timestep,
+                    "operation": "store_environment_frame",
+                }
+                if self.current_step:
+                    custom_params["step"] = self.current_step
+                if state_identifier:
+                    custom_params["state_identifier"] = state_identifier
+
+                with external_trace(
+                    url=url, method="POST", library="httpx", custom_params=custom_params
+                ):
+                    headers = insert_distributed_trace_headers(headers)
+
+                    response = await client.post(
+                        url=url,
+                        json=data.model_dump(mode="json"),
+                        headers=headers,
+                    )
                 response.raise_for_status()
                 logger.debug(
                     f"Successfully stored environment frame for state {state_identifier}",

futurehouse_client/clients/rest_client.py

@@ -15,9 +15,19 @@ from uuid import UUID
 
 import cloudpickle
 from aviary.functional import EnvironmentBuilder
-from httpx import Client, HTTPStatusError
+from httpx import (
+    Client,
+    CloseError,
+    ConnectError,
+    ConnectTimeout,
+    HTTPStatusError,
+    NetworkError,
+    ReadError,
+    ReadTimeout,
+    RemoteProtocolError,
+)
 from pydantic import BaseModel, ConfigDict, model_validator
-from requests.exceptions import Timeout
+from requests.exceptions import RequestException, Timeout
 from tenacity import (
     retry,
     retry_if_exception_type,
@@ -37,11 +47,29 @@ from futurehouse_client.utils.module_utils import (
     OrganizationSelector,
     fetch_environment_function_docstring,
 )
+from futurehouse_client.utils.monitoring import (
+    external_trace,
+)
 
 logger = logging.getLogger(__name__)
 
 TaskRequest.model_rebuild()
 
+retry_if_connection_error = retry_if_exception_type((
+    # From requests
+    Timeout,
+    ConnectionError,
+    RequestException,
+    # From httpx
+    ConnectError,
+    ConnectTimeout,
+    ReadTimeout,
+    ReadError,
+    NetworkError,
+    RemoteProtocolError,
+    CloseError,
+))
+
 FILE_UPLOAD_IGNORE_PARTS = {
     ".ruff_cache",
     "__pycache__",
@@ -342,7 +370,7 @@ class RestClient:
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
-        retry=retry_if_exception_type(Timeout),
+        retry=retry_if_connection_error,
     )
     def get_task(
         self, task_id: str | None = None, history: bool = False, verbose: bool = False
@@ -350,10 +378,22 @@ class RestClient:
         """Get details for a specific task."""
         try:
             task_id = task_id or self.trajectory_id
-            response = self.client.get(
-                f"/v0.1/trajectories/{task_id}",
-                params={"history": history},
-            )
+            url = f"/v0.1/trajectories/{task_id}"
+            full_url = f"{self.base_url}{url}"
+
+            with external_trace(
+                url=full_url,
+                method="GET",
+                library="httpx",
+                custom_params={
+                    "operation": "get_job",
+                    "job_id": task_id,
+                },
+            ):
+                response = self.client.get(
+                    url,
+                    params={"history": history},
+                )
             response.raise_for_status()
             verbose_response = TaskResponseVerbose(**response.json())
             if verbose:
@@ -372,7 +412,7 @@ class RestClient:
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
-        retry=retry_if_exception_type(Timeout),
+        retry=retry_if_connection_error,
     )
     def create_task(self, task_data: TaskRequest | dict[str, Any]):
         """Create a new futurehouse task."""
@@ -398,7 +438,7 @@ class RestClient:
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
-        retry=retry_if_exception_type(Timeout),
+        retry=retry_if_connection_error,
     )
     def get_build_status(self, build_id: UUID | None = None) -> dict[str, Any]:
         """Get the status of a build."""
@@ -411,7 +451,7 @@ class RestClient:
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
-        retry=retry_if_exception_type(Timeout),
+        retry=retry_if_connection_error,
     )
     def create_job(self, config: JobDeploymentConfig) -> dict[str, Any]:  # noqa: PLR0915
         """Creates a futurehouse job deployment from the environment and environment files.

futurehouse_client/utils/monitoring.py

@@ -0,0 +1,246 @@
+"""Utilities for monitoring and observability integration.
+
+This module provides utilities for integrating with monitoring and observability tools
+like NewRelic. It handles availability checking and provides wrapper functions that
+conditionally use monitoring tools only when they're available and properly initialized.
+
+NOTE: NewRelic is an optional dependency. To use monitoring functionality, install
+the package with the monitoring extras:
+    pip install futurehouse-client[monitoring]
+
+Environment variables:
+    NEW_RELIC_ENVIRONMENT: The environment to use for NewRelic reporting (dev, staging, prod)
+    NEW_RELIC_CONFIG_FILE: Path to the NewRelic configuration file
+    NEW_RELIC_LICENSE_KEY: Your NewRelic license key
+"""
+
+import contextlib
+import json
+import logging
+import os
+from contextvars import ContextVar
+
+logger = logging.getLogger(__name__)
+
+# Check if NewRelic initialization is enabled (default: False)
+NEWRELIC_AUTO_INIT = (
+    os.environ.get("FUTUREHOUSE_NEWRELIC_AUTO_INIT", "false").lower() == "true"
+)
+
+# Check if NewRelic is installed
+try:
+    import newrelic.agent
+
+    NEWRELIC_INSTALLED = True
+except ImportError:
+    NEWRELIC_INSTALLED = False
+    logger.info("NewRelic package not installed")
+
+# Context variable to track NewRelic initialization state
+newrelic_initialized: ContextVar[bool] = ContextVar(
+    "newrelic_initialized", default=False
+)
+
+
+def ensure_newrelic() -> bool:  # noqa: PLR0911
+    """Check if NewRelic is available and initialize it if auto-init is enabled.
+
+    This will use environment variables:
+    - FUTUREHOUSE_NEWRELIC_AUTO_INIT: Set to "true" to enable automatic initialization (default: "false")
+    - NEW_RELIC_CONFIG_FILE: Path to the NewRelic config file (required)
+    - NEW_RELIC_ENVIRONMENT: Environment (dev, staging, prod)
+
+    Returns:
+        bool: True if NewRelic is available for use, False otherwise
+    """
+    if newrelic_initialized.get():
+        return True
+
+    if not NEWRELIC_AUTO_INIT:
+        return False
+
+    if not NEWRELIC_INSTALLED:
+        logger.info("NewRelic package is not installed")
+        return False
+
+    nr_config = os.environ.get("NEW_RELIC_CONFIG_FILE")
+    if not nr_config:
+        logger.warning("NEW_RELIC_CONFIG_FILE environment variable must be set")
+        return False
+
+    try:
+        nr_env = os.environ.get("NEW_RELIC_ENVIRONMENT", "dev")
+        newrelic.agent.initialize(nr_config, environment=nr_env)
+
+        app = newrelic.agent.application()
+        if app is None:
+            logger.warning("NewRelic initialization failed: no application returned")
+            return False
+
+        newrelic_initialized.set(True)
+        logger.info(f"NewRelic initialized successfully for environment: {nr_env}")
+    except Exception as e:
+        logger.warning(f"NewRelic initialization failed: {e}")
+        return False
+
+    else:
+        return True
+
+
+def insert_distributed_trace_headers(headers: dict[str, str]) -> dict[str, str]:
+    """Insert distributed trace headers if NewRelic is available.
+
+    Args:
+        headers: The headers dictionary to modify.
+
+    Returns:
+        The modified headers dictionary with NewRelic distributed trace headers if available,
+        otherwise the original headers.
+    """
+    if not ensure_newrelic():
+        return headers
+
+    try:
+        nr_headers: list[tuple[str, str]] = []
+        newrelic.agent.insert_distributed_trace_headers(nr_headers)
+        for header in nr_headers:
+            headers[header[0]] = header[1]
+    except Exception as e:
+        logger.info(f"Error inserting distributed trace headers: {e}")
+
+    return headers
+
+
+@contextlib.contextmanager
+def external_trace(
+    url: str,
+    method: str = "GET",
+    library: str = "httpx",
+    custom_params: dict | None = None,
+):
+    """Context manager for NewRelic external traces that works whether NewRelic is available or not.
+
+    Creates an ExternalTrace span in NewRelic for HTTP requests to external services. This provides detailed timing and proper distributed tracing between services.
+    "External" refers to HTTP requests made to services outside of your application (like third-party APIs or other microservices).
+
+    Args:
+        url: The URL being called.
+        method: The HTTP method (GET, POST, etc.).
+        library: The library being used for the HTTP call.
+        custom_params: Optional dictionary of custom parameters to add to the transaction.
+
+    Yields:
+        None: This is a context manager that doesn't yield a value.
+    """
+    if not ensure_newrelic():
+        yield
+        return
+
+    # Proceed with tracing
+    try:
+        with newrelic.agent.ExternalTrace(
+            library=library,
+            url=url,
+            method=method,
+        ):
+            txn = newrelic.agent.current_transaction()
+            if txn:
+                txn.add_custom_parameter("request_url", url)
+                txn.add_custom_parameter("request_method", method)
+
+                if custom_params:
+                    for key, value in custom_params.items():
+                        txn.add_custom_parameter(key, value)
+
+            yield
+    except Exception as e:
+        # If there's an exception in the transaction handling,
+        # log it but don't let it break the client
+        try:
+            txn = newrelic.agent.current_transaction()
+            if txn:
+                txn.add_custom_parameter("external_request_url", url)
+                txn.add_custom_parameter("external_request_method", method)
+                txn.add_custom_parameter("error_type", e.__class__.__name__)
+                txn.add_custom_parameter("error_message", str(e))
+                txn.record_exception(e)
+        except Exception as nr_error:
+            # If even the error handling fails, just log it
+            logger.info(f"Failed to record NewRelic error: {nr_error}")
+
+        # Always re-raise the original exception
+        raise
+
+
+@contextlib.contextmanager
+def monitored_transaction(
+    name: str, group: str = "Task", custom_params: dict | None = None
+):
+    """Context manager for NewRelic transactions that appear in distributed traces.
+
+    This uses WebTransaction for better visibility in distributed traces, even for
+    background job-type workloads.
+
+    Args:
+        name: Name of the transaction (e.g., 'run_crow_job')
+        group: Group for transaction categorization (default: 'Task')
+        custom_params: Optional dictionary of custom parameters to add to the transaction
+
+    Yields:
+        None: This is a context manager that doesn't yield a value.
+    """
+    if not ensure_newrelic():
+        logger.info("NewRelic not available, skipping transaction")
+        yield
+        return
+
+    try:
+        app = newrelic.agent.application()
+        if app is None:
+            logger.warning("No NewRelic application found, skipping transaction")
+            yield
+            return
+
+        parsed_headers = None
+        trace_context = os.environ.get("NEW_RELIC_DISTRIBUTED_TRACING_CONTEXT")
+        if trace_context:
+            try:
+                parsed_headers = json.loads(trace_context)
+            except Exception as e:
+                logger.warning(f"Failed to parse distributed trace context: {e}")
+        else:
+            logger.info("No distributed trace context found")
+
+        with newrelic.agent.WebTransaction(app, name, group=group):
+            if parsed_headers:
+                current_txn = newrelic.agent.current_transaction()
+                if current_txn:
+                    accepted = newrelic.agent.accept_distributed_trace_headers(
+                        parsed_headers
+                    )
+                    if not accepted:
+                        logger.warning("Failed to accept distributed trace headers")
+
+            if custom_params:
+                txn = newrelic.agent.current_transaction()
+                if txn:
+                    for key, value in custom_params.items():
+                        txn.add_custom_parameter(key, value)
+
+            yield
+            logger.info(f"Completed NewRelic transaction: {name}")
+    except Exception as e:
+        # If there's an exception in the transaction handling,
+        # log it but don't let it break the client
+        try:
+            txn = newrelic.agent.current_transaction()
+            if txn:
+                txn.add_custom_parameter("error_type", e.__class__.__name__)
+                txn.add_custom_parameter("error_message", str(e))
+                txn.record_exception(e)
+        except Exception as nr_error:
+            # If even the error handling fails, just log it
+            logger.info(f"Failed to record NewRelic error: {nr_error}")
+
+        logger.warning(f"Error in NewRelic transaction: {e}")
+        yield

{futurehouse_client-0.0.5.dist-info → futurehouse_client-0.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.0.5
+Version: 0.0.6
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 Classifier: Operating System :: OS Independent
@@ -35,6 +35,8 @@ Requires-Dist: pytest-timeout; extra == "dev"
 Requires-Dist: pytest-xdist; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: setuptools_scm; extra == "dev"
+Provides-Extra: monitoring
+Requires-Dist: newrelic>=8.8.0; extra == "monitoring"
 
 # FutureHouse Platform API Documentation
 

{futurehouse_client-0.0.5.dist-info → futurehouse_client-0.0.6.dist-info}/RECORD

@@ -1,14 +1,15 @@
 futurehouse_client/__init__.py,sha256=ddxO7JE97c6bt7LjNglZZ2Ql8bYCGI9laSFeh9MP6VU,344
 futurehouse_client/clients/__init__.py,sha256=tFWqwIAY5PvwfOVsCje4imjTpf6xXNRMh_UHIKVI1_0,320
-futurehouse_client/clients/job_client.py,sha256=RNgdSJVI1vjQSypdfswHX0Gvv_XnKG4bZjklf3WdSuk,8828
-futurehouse_client/clients/rest_client.py,sha256=GtHDTGHdz0f6xzBWaufqh3f76Life66df6abrJrSKvA,25116
+futurehouse_client/clients/job_client.py,sha256=yBFKDNcFnuZDNgoK2d5037rbuzQ7TlSK6MmklEKV8EA,11056
+futurehouse_client/clients/rest_client.py,sha256=ZCGhKCNzDtaLRPjhciKXdjbR_PYHKkBt7SSa8ansuaM,25961
 futurehouse_client/models/__init__.py,sha256=ta3jFLM_LsDz1rKDmx8rja8sT7WtSKoFvMgLF0yFpvA,342
 futurehouse_client/models/app.py,sha256=G8_-I3aQnRIyxFJT3snTSWsPcZZ2nEvYkRaE5sVdeys,22299
 futurehouse_client/models/client.py,sha256=n4HD0KStKLm6Ek9nL9ylP-bkK10yzAaD1uIDF83Qp_A,1828
 futurehouse_client/models/rest.py,sha256=W-wNFTN7HALYFFphw-RQYRMm6_TSa1cl4T-mZ1msk90,393
 futurehouse_client/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 futurehouse_client/utils/module_utils.py,sha256=aFyd-X-pDARXz9GWpn8SSViUVYdSbuy9vSkrzcVIaGI,4955
-futurehouse_client-0.0.5.dist-info/METADATA,sha256=2OxntdouInoOXVuPlmVmKkIwZ7g4E_rilBDZ0HCkobU,8056
-futurehouse_client-0.0.5.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-futurehouse_client-0.0.5.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
-futurehouse_client-0.0.5.dist-info/RECORD,,
+futurehouse_client/utils/monitoring.py,sha256=UjRlufe67kI3VxRHOd5fLtJmlCbVA2Wqwpd4uZhXkQM,8728
+futurehouse_client-0.0.6.dist-info/METADATA,sha256=K0-xlHNzmzmBOpCRJIXu-izmorzx4rik_Zx5giLjfoY,8137
+futurehouse_client-0.0.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+futurehouse_client-0.0.6.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
+futurehouse_client-0.0.6.dist-info/RECORD,,