futurehouse-client 0.3.18.dev109__tar.gz → 0.3.18.dev184__tar.gz

{futurehouse_client-0.3.18.dev109 → futurehouse_client-0.3.18.dev184}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.3.18.dev109
+Version: 0.3.18.dev184
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 Classifier: Operating System :: OS Independent

{futurehouse_client-0.3.18.dev109 → futurehouse_client-0.3.18.dev184}/futurehouse_client/clients/rest_client.py

@@ -12,8 +12,7 @@ import sys
 import tempfile
 import time
 import uuid
-from collections.abc import Collection, Mapping
-from datetime import datetime
+from collections.abc import Collection
 from pathlib import Path
 from types import ModuleType
 from typing import Any, ClassVar, assert_never, cast
@@ -34,7 +33,6 @@ from httpx import (
     RemoteProtocolError,
 )
 from ldp.agent import AgentConfig
-from pydantic import BaseModel, ConfigDict, model_validator
 from requests.exceptions import RequestException, Timeout
 from tenacity import (
     retry,
@@ -50,10 +48,18 @@ from futurehouse_client.models.app import (
     APIKeyPayload,
     AuthType,
     JobDeploymentConfig,
+    PQATaskResponse,
     Stage,
     TaskRequest,
+    TaskResponse,
+    TaskResponseVerbose,
 )
 from futurehouse_client.models.rest import ExecutionStatus
+from futurehouse_client.utils.auth import (
+    AUTH_ERRORS_TO_RETRY_ON,
+    AuthError,
+    refresh_token_on_auth_error,
+)
 from futurehouse_client.utils.general import gather_with_concurrency
 from futurehouse_client.utils.module_utils import (
     OrganizationSelector,
@@ -65,7 +71,7 @@ from futurehouse_client.utils.monitoring import (
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.WARNING,
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
     stream=sys.stdout,
 )
@@ -122,103 +128,11 @@ retry_if_connection_error = retry_if_exception_type((
     FileUploadError,
 ))
 
-
-class SimpleOrganization(BaseModel):
-    id: int
-    name: str
-    display_name: str
-
-
 # 5 minute default for JWTs
 JWT_TOKEN_CACHE_EXPIRY: int = 300  # seconds
 DEFAULT_AGENT_TIMEOUT: int = 2400  # seconds
 
 
-class TaskResponse(BaseModel):
-    """Base class for task responses. This holds attributes shared over all futurehouse jobs."""
-
-    model_config = ConfigDict(extra="ignore")
-
-    status: str
-    query: str
-    user: str | None = None
-    created_at: datetime
-    job_name: str
-    public: bool
-    shared_with: list[SimpleOrganization] | None = None
-    build_owner: str | None = None
-    environment_name: str | None = None
-    agent_name: str | None = None
-    task_id: UUID | None = None
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
-        # Extract fields from environment frame state
-        if not isinstance(data, dict):
-            return data
-        # TODO: We probably want to remove these two once we define the final names.
-        data["job_name"] = data.get("crow")
-        data["query"] = data.get("task")
-        data["task_id"] = cast(UUID, data.get("id")) if data.get("id") else None
-        if not (metadata := data.get("metadata", {})):
-            return data
-        data["environment_name"] = metadata.get("environment_name")
-        data["agent_name"] = metadata.get("agent_name")
-        return data
-
-
-class PQATaskResponse(TaskResponse):
-    model_config = ConfigDict(extra="ignore")
-
-    answer: str | None = None
-    formatted_answer: str | None = None
-    answer_reasoning: str | None = None
-    has_successful_answer: bool | None = None
-    total_cost: float | None = None
-    total_queries: int | None = None
-
-    @model_validator(mode="before")
-    @classmethod
-    def validate_pqa_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
-        if not isinstance(data, dict):
-            return data
-        if not (env_frame := data.get("environment_frame", {})):
-            return data
-        state = env_frame.get("state", {}).get("state", {})
-        response = state.get("response", {})
-        answer = response.get("answer", {})
-        usage = state.get("info", {}).get("usage", {})
-
-        # Add additional PQA specific fields to data so that pydantic can validate the model
-        data["answer"] = answer.get("answer")
-        data["formatted_answer"] = answer.get("formatted_answer")
-        data["answer_reasoning"] = answer.get("answer_reasoning")
-        data["has_successful_answer"] = answer.get("has_successful_answer")
-        data["total_cost"] = cast(float, usage.get("total_cost"))
-        data["total_queries"] = cast(int, usage.get("total_queries"))
-
-        return data
-
-    def clean_verbose(self) -> "TaskResponse":
-        """Clean the verbose response from the server."""
-        self.request = None
-        self.response = None
-        return self
-
-
-class TaskResponseVerbose(TaskResponse):
-    """Class for responses to include all the fields of a task response."""
-
-    model_config = ConfigDict(extra="allow")
-
-    public: bool
-    agent_state: list[dict[str, Any]] | None = None
-    environment_frame: dict[str, Any] | None = None
-    metadata: dict[str, Any] | None = None
-    shared_with: list[SimpleOrganization] | None = None
-
-
 class RestClient:
     REQUEST_TIMEOUT: ClassVar[float] = 30.0  # sec
     MAX_RETRY_ATTEMPTS: ClassVar[int] = 3
@@ -236,7 +150,13 @@ class RestClient:
         api_key: str | None = None,
         jwt: str | None = None,
         headers: dict[str, str] | None = None,
+        verbose_logging: bool = False,
     ):
+        if verbose_logging:
+            logger.setLevel(logging.INFO)
+        else:
+            logger.setLevel(logging.WARNING)
+
         self.base_url = service_uri or stage.value
         self.stage = stage
         self.auth_type = auth_type
@@ -360,6 +280,7 @@ class RestClient:
         except Exception as e:
             raise RestClientError(f"Error authenticating: {e!s}") from e
 
+    @refresh_token_on_auth_error()
     def _check_job(self, name: str, organization: str) -> dict[str, Any]:
         try:
             response = self.client.get(
@@ -367,9 +288,19 @@ class RestClient:
             )
             response.raise_for_status()
             return response.json()
+        except HTTPStatusError as e:
+            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
+                raise AuthError(
+                    e.response.status_code,
+                    f"Authentication failed: {e}",
+                    request=e.request,
+                    response=e.response,
+                ) from e
+            raise
         except Exception as e:
             raise JobFetchError(f"Error checking job: {e!s}") from e
 
+    @refresh_token_on_auth_error()
     def _fetch_my_orgs(self) -> list[str]:
         response = self.client.get(f"/v0.1/organizations?filter={True}")
         response.raise_for_status()
@@ -432,6 +363,7 @@ class RestClient:
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
+    @refresh_token_on_auth_error()
     def get_task(
         self, task_id: str | None = None, history: bool = False, verbose: bool = False
     ) -> "TaskResponse":
@@ -467,8 +399,15 @@ class RestClient:
             ):
                 return PQATaskResponse(**data)
             return TaskResponse(**data)
-        except ValueError as e:
-            raise ValueError("Invalid task ID format. Must be a valid UUID.") from e
+        except HTTPStatusError as e:
+            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
+                raise AuthError(
+                    e.response.status_code,
+                    f"Authentication failed: {e}",
+                    request=e.request,
+                    response=e.response,
+                ) from e
+            raise
         except Exception as e:
             raise TaskFetchError(f"Error getting task: {e!s}") from e
 
@@ -477,6 +416,7 @@ class RestClient:
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
+    @refresh_token_on_auth_error()
     async def aget_task(
         self, task_id: str | None = None, history: bool = False, verbose: bool = False
     ) -> "TaskResponse":
@@ -515,11 +455,19 @@ class RestClient:
             ):
                 return PQATaskResponse(**data)
             return TaskResponse(**data)
-        except ValueError as e:
-            raise ValueError("Invalid task ID format. Must be a valid UUID.") from e
+        except HTTPStatusError as e:
+            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
+                raise AuthError(
+                    e.response.status_code,
+                    f"Authentication failed: {e}",
+                    request=e.request,
+                    response=e.response,
+                ) from e
+            raise
         except Exception as e:
             raise TaskFetchError(f"Error getting task: {e!s}") from e
 
+    @refresh_token_on_auth_error()
     @retry(
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
@@ -543,6 +491,15 @@ class RestClient:
             response.raise_for_status()
             trajectory_id = response.json()["trajectory_id"]
             self.trajectory_id = trajectory_id
+        except HTTPStatusError as e:
+            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
+                raise AuthError(
+                    e.response.status_code,
+                    f"Authentication failed: {e}",
+                    request=e.request,
+                    response=e.response,
+                ) from e
+            raise
         except Exception as e:
             raise TaskFetchError(f"Error creating task: {e!s}") from e
         return trajectory_id
@@ -552,6 +509,7 @@ class RestClient:
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
+    @refresh_token_on_auth_error()
     async def acreate_task(self, task_data: TaskRequest | dict[str, Any]):
         """Create a new futurehouse task."""
         if isinstance(task_data, dict):
@@ -570,6 +528,15 @@ class RestClient:
             response.raise_for_status()
             trajectory_id = response.json()["trajectory_id"]
             self.trajectory_id = trajectory_id
+        except HTTPStatusError as e:
+            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
+                raise AuthError(
+                    e.response.status_code,
+                    f"Authentication failed: {e}",
+                    request=e.request,
+                    response=e.response,
+                ) from e
+            raise
         except Exception as e:
             raise TaskFetchError(f"Error creating task: {e!s}") from e
         return trajectory_id
@@ -720,11 +687,22 @@ class RestClient:
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
+    @refresh_token_on_auth_error()
     def get_build_status(self, build_id: UUID | None = None) -> dict[str, Any]:
         """Get the status of a build."""
-        build_id = build_id or self.build_id
-        response = self.client.get(f"/v0.1/builds/{build_id}")
-        response.raise_for_status()
+        try:
+            build_id = build_id or self.build_id
+            response = self.client.get(f"/v0.1/builds/{build_id}")
+            response.raise_for_status()
+        except HTTPStatusError as e:
+            if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
+                raise AuthError(
+                    e.response.status_code,
+                    f"Authentication failed: {e}",
+                    request=e.request,
+                    response=e.response,
+                ) from e
+            raise
         return response.json()
 
     # TODO: Refactor later so we don't have to ignore PLR0915
@@ -733,6 +711,7 @@ class RestClient:
         wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
         retry=retry_if_connection_error,
     )
+    @refresh_token_on_auth_error()
     def create_job(self, config: JobDeploymentConfig) -> dict[str, Any]:  # noqa: PLR0915
         """Creates a futurehouse job deployment from the environment and environment files.
 
@@ -907,6 +886,13 @@ class RestClient:
                 build_context = response.json()
                 self.build_id = build_context["build_id"]
             except HTTPStatusError as e:
+                if e.response.status_code in AUTH_ERRORS_TO_RETRY_ON:
+                    raise AuthError(
+                        e.response.status_code,
+                        f"Authentication failed: {e}",
+                        request=e.request,
+                        response=e.response,
+                    ) from e
                 error_detail = response.json()
                 error_message = error_detail.get("detail", str(e))
                 raise JobCreationError(

{futurehouse_client-0.3.18.dev109 → futurehouse_client-0.3.18.dev184}/futurehouse_client/models/__init__.py

@@ -3,10 +3,15 @@ from .app import (
     DockerContainerConfiguration,
     FramePath,
     JobDeploymentConfig,
+    PQATaskResponse,
     RuntimeConfig,
     Stage,
     Step,
+    TaskQueue,
+    TaskQueuesConfig,
     TaskRequest,
+    TaskResponse,
+    TaskResponseVerbose,
 )
 
 __all__ = [
@@ -14,8 +19,13 @@ __all__ = [
     "DockerContainerConfiguration",
     "FramePath",
     "JobDeploymentConfig",
+    "PQATaskResponse",
     "RuntimeConfig",
     "Stage",
     "Step",
+    "TaskQueue",
+    "TaskQueuesConfig",
     "TaskRequest",
+    "TaskResponse",
+    "TaskResponseVerbose",
 ]

{futurehouse_client-0.3.18.dev109 → futurehouse_client-0.3.18.dev184}/futurehouse_client/models/app.py

@@ -1,6 +1,8 @@
 import json
 import os
 import re
+from collections.abc import Mapping
+from datetime import datetime
 from enum import StrEnum, auto
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Self, cast
@@ -646,3 +648,94 @@ class TaskRequest(BaseModel):
     runtime_config: RuntimeConfig | None = Field(
         default=None, description="All optional runtime parameters for the job"
     )
+
+
+class SimpleOrganization(BaseModel):
+    id: int
+    name: str
+    display_name: str
+
+
+class TaskResponse(BaseModel):
+    """Base class for task responses. This holds attributes shared over all futurehouse jobs."""
+
+    model_config = ConfigDict(extra="ignore")
+
+    status: str
+    query: str
+    user: str | None = None
+    created_at: datetime
+    job_name: str
+    public: bool
+    shared_with: list[SimpleOrganization] | None = None
+    build_owner: str | None = None
+    environment_name: str | None = None
+    agent_name: str | None = None
+    task_id: UUID | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
+        # Extract fields from environment frame state
+        if not isinstance(data, dict):
+            return data
+        # TODO: We probably want to remove these two once we define the final names.
+        data["job_name"] = data.get("crow")
+        data["query"] = data.get("task")
+        data["task_id"] = cast(UUID, data.get("id")) if data.get("id") else None
+        if not (metadata := data.get("metadata", {})):
+            return data
+        data["environment_name"] = metadata.get("environment_name")
+        data["agent_name"] = metadata.get("agent_name")
+        return data
+
+
+class PQATaskResponse(TaskResponse):
+    model_config = ConfigDict(extra="ignore")
+
+    answer: str | None = None
+    formatted_answer: str | None = None
+    answer_reasoning: str | None = None
+    has_successful_answer: bool | None = None
+    total_cost: float | None = None
+    total_queries: int | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_pqa_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
+        if not isinstance(data, dict):
+            return data
+        if not (env_frame := data.get("environment_frame", {})):
+            return data
+        state = env_frame.get("state", {}).get("state", {})
+        response = state.get("response", {})
+        answer = response.get("answer", {})
+        usage = state.get("info", {}).get("usage", {})
+
+        # Add additional PQA specific fields to data so that pydantic can validate the model
+        data["answer"] = answer.get("answer")
+        data["formatted_answer"] = answer.get("formatted_answer")
+        data["answer_reasoning"] = answer.get("answer_reasoning")
+        data["has_successful_answer"] = answer.get("has_successful_answer")
+        data["total_cost"] = cast(float, usage.get("total_cost"))
+        data["total_queries"] = cast(int, usage.get("total_queries"))
+
+        return data
+
+    def clean_verbose(self) -> "TaskResponse":
+        """Clean the verbose response from the server."""
+        self.request = None
+        self.response = None
+        return self
+
+
+class TaskResponseVerbose(TaskResponse):
+    """Class for responses to include all the fields of a task response."""
+
+    model_config = ConfigDict(extra="allow")
+
+    public: bool
+    agent_state: list[dict[str, Any]] | None = None
+    environment_frame: dict[str, Any] | None = None
+    metadata: dict[str, Any] | None = None
+    shared_with: list[SimpleOrganization] | None = None

futurehouse_client-0.3.18.dev184/futurehouse_client/utils/auth.py

@@ -0,0 +1,107 @@
+import asyncio
+import logging
+from collections.abc import Callable, Coroutine
+from functools import wraps
+from typing import Any, Final, Optional, ParamSpec, TypeVar, overload
+
+import httpx
+from httpx import HTTPStatusError
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+AUTH_ERRORS_TO_RETRY_ON: Final[set[int]] = {
+    httpx.codes.UNAUTHORIZED,
+    httpx.codes.FORBIDDEN,
+}
+
+
+class AuthError(Exception):
+    """Raised when authentication fails with 401/403 status."""
+
+    def __init__(self, status_code: int, message: str, request=None, response=None):
+        self.status_code = status_code
+        self.request = request
+        self.response = response
+        super().__init__(message)
+
+
+def is_auth_error(e: Exception) -> bool:
+    if isinstance(e, AuthError):
+        return True
+    if isinstance(e, HTTPStatusError):
+        return e.response.status_code in AUTH_ERRORS_TO_RETRY_ON
+    return False
+
+
+def get_status_code(e: Exception) -> Optional[int]:
+    if isinstance(e, AuthError):
+        return e.status_code
+    if isinstance(e, HTTPStatusError):
+        return e.response.status_code
+    return None
+
+
+@overload
+def refresh_token_on_auth_error(
+    func: Callable[P, Coroutine[Any, Any, T]],
+) -> Callable[P, Coroutine[Any, Any, T]]: ...
+
+
+@overload
+def refresh_token_on_auth_error(
+    func: None = None, *, max_retries: int = ...
+) -> Callable[[Callable[P, T]], Callable[P, T]]: ...
+
+
+def refresh_token_on_auth_error(func=None, max_retries=1):
+    """Decorator that refreshes JWT token on 401/403 auth errors."""
+
+    def decorator(fn):
+        @wraps(fn)
+        def sync_wrapper(self, *args, **kwargs):
+            retries = 0
+            while True:
+                try:
+                    return fn(self, *args, **kwargs)
+                except Exception as e:
+                    if is_auth_error(e) and retries < max_retries:
+                        retries += 1
+                        status = get_status_code(e) or "Unknown"
+                        logger.info(
+                            f"Received auth error {status}, "
+                            f"refreshing token and retrying (attempt {retries}/{max_retries})..."
+                        )
+                        self.auth_jwt = self._run_auth()
+                        self._clients = {}
+                        continue
+                    raise
+
+        @wraps(fn)
+        async def async_wrapper(self, *args, **kwargs):
+            retries = 0
+            while True:
+                try:
+                    return await fn(self, *args, **kwargs)
+                except Exception as e:
+                    if is_auth_error(e) and retries < max_retries:
+                        retries += 1
+                        status = get_status_code(e) or "Unknown"
+                        logger.info(
+                            f"Received auth error {status}, "
+                            f"refreshing token and retrying (attempt {retries}/{max_retries})..."
+                        )
+                        self.auth_jwt = self._run_auth()
+                        self._clients = {}
+                        continue
+                    raise
+
+        if asyncio.iscoroutinefunction(fn):
+            return async_wrapper
+        return sync_wrapper
+
+    if callable(func):
+        return decorator(func)
+    return decorator

{futurehouse_client-0.3.18.dev109 → futurehouse_client-0.3.18.dev184}/futurehouse_client.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.3.18.dev109
+Version: 0.3.18.dev184
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 Classifier: Operating System :: OS Independent

{futurehouse_client-0.3.18.dev109 → futurehouse_client-0.3.18.dev184}/futurehouse_client.egg-info/SOURCES.txt

@@ -18,7 +18,7 @@ futurehouse_client/models/app.py
 futurehouse_client/models/client.py
 futurehouse_client/models/rest.py
 futurehouse_client/utils/__init__.py
-futurehouse_client/utils/context.py
+futurehouse_client/utils/auth.py
 futurehouse_client/utils/general.py
 futurehouse_client/utils/module_utils.py
 futurehouse_client/utils/monitoring.py

futurehouse_client-0.3.18.dev184/tests/test_rest.py

@@ -0,0 +1,260 @@
+# ruff: noqa: ARG001
+import asyncio
+import os
+import time
+from unittest.mock import patch
+
+import pytest
+from futurehouse_client.clients import (
+    JobNames,
+    PQATaskResponse,
+    TaskResponseVerbose,
+)
+from futurehouse_client.clients.rest_client import RestClient
+from futurehouse_client.models.app import Stage, TaskRequest
+from futurehouse_client.models.rest import ExecutionStatus
+from futurehouse_client.utils.auth import AuthError, refresh_token_on_auth_error
+from pytest_subtests import SubTests
+
+ADMIN_API_KEY = os.environ["PLAYWRIGHT_ADMIN_API_KEY"]
+PUBLIC_API_KEY = os.environ["PLAYWRIGHT_PUBLIC_API_KEY"]
+TEST_MAX_POLLS = 100
+
+
+@pytest.fixture
+def admin_client():
+    """Create a RestClient for testing."""
+    return RestClient(
+        stage=Stage.DEV,
+        api_key=ADMIN_API_KEY,
+    )
+
+
+@pytest.fixture
+def pub_client():
+    """Create a RestClient for testing."""
+    return RestClient(
+        stage=Stage.DEV,
+        api_key=PUBLIC_API_KEY,
+    )
+
+
+@pytest.fixture
+def task_data():
+    """Create a sample task request."""
+    return TaskRequest(
+        name=JobNames.from_string("dummy"),
+        query="How many moons does earth have?",
+    )
+
+
+@pytest.fixture
+def pqa_task_data():
+    return TaskRequest(
+        name=JobNames.from_string("crow"),
+        query="How many moons does earth have?",
+    )
+
+
+@pytest.mark.timeout(300)
+@pytest.mark.flaky(reruns=3)
+def test_futurehouse_dummy_env_crow(admin_client: RestClient, task_data: TaskRequest):
+    admin_client.create_task(task_data)
+    while (task_status := admin_client.get_task().status) in {"queued", "in progress"}:
+        time.sleep(5)
+    assert task_status == "success"
+
+
+def test_insufficient_permissions_request(
+    pub_client: RestClient, task_data: TaskRequest
+):
+    # Create a new instance so that cached credentials aren't reused
+    with pytest.raises(AuthError) as exc_info:
+        pub_client.create_task(task_data)
+
+    assert "403 Forbidden" in str(exc_info.value)
+
+
+@pytest.mark.timeout(300)
+@pytest.mark.asyncio
+async def test_job_response(  # noqa: PLR0915
+    subtests: SubTests, admin_client: RestClient, pqa_task_data: TaskRequest
+):
+    task_id = admin_client.create_task(pqa_task_data)
+    atask_id = await admin_client.acreate_task(pqa_task_data)
+
+    with subtests.test("Test TaskResponse with queued task"):
+        task_response = admin_client.get_task(task_id)
+        assert task_response.status in {"queued", "in progress"}
+        assert task_response.job_name == pqa_task_data.name
+        assert task_response.query == pqa_task_data.query
+        task_response = await admin_client.aget_task(atask_id)
+        assert task_response.status in {"queued", "in progress"}
+        assert task_response.job_name == pqa_task_data.name
+        assert task_response.query == pqa_task_data.query
+
+    for _ in range(TEST_MAX_POLLS):
+        task_response = admin_client.get_task(task_id)
+        if task_response.status in ExecutionStatus.terminal_states():
+            break
+        await asyncio.sleep(5)
+
+    for _ in range(TEST_MAX_POLLS):
+        task_response = await admin_client.aget_task(atask_id)
+        if task_response.status in ExecutionStatus.terminal_states():
+            break
+        await asyncio.sleep(5)
+
+    with subtests.test("Test PQA job response"):
+        task_response = admin_client.get_task(task_id)
+        assert isinstance(task_response, PQATaskResponse)
+        # assert it has general fields
+        assert task_response.status == "success"
+        assert task_response.task_id is not None
+        assert pqa_task_data.name in task_response.job_name
+        assert pqa_task_data.query in task_response.query
+        # assert it has PQA specific fields
+        assert task_response.answer is not None
+        # assert it's not verbose
+        assert not hasattr(task_response, "environment_frame")
+        assert not hasattr(task_response, "agent_state")
+
+    with subtests.test("Test async PQA job response"):
+        task_response = await admin_client.aget_task(atask_id)
+        assert isinstance(task_response, PQATaskResponse)
+        # assert it has general fields
+        assert task_response.status == "success"
+        assert task_response.task_id is not None
+        assert pqa_task_data.name in task_response.job_name
+        assert pqa_task_data.query in task_response.query
+        # assert it has PQA specific fields
+        assert task_response.answer is not None
+        # assert it's not verbose
+        assert not hasattr(task_response, "environment_frame")
+        assert not hasattr(task_response, "agent_state")
+
+    with subtests.test("Test task response with verbose"):
+        task_response = admin_client.get_task(task_id, verbose=True)
+        assert isinstance(task_response, TaskResponseVerbose)
+        assert task_response.status == "success"
+        assert task_response.environment_frame is not None
+        assert task_response.agent_state is not None
+
+    with subtests.test("Test task async response with verbose"):
+        task_response = await admin_client.aget_task(atask_id, verbose=True)
+        assert isinstance(task_response, TaskResponseVerbose)
+        assert task_response.status == "success"
+        assert task_response.environment_frame is not None
+        assert task_response.agent_state is not None
+
+
+@pytest.mark.timeout(300)
+@pytest.mark.flaky(reruns=3)
+def test_run_until_done_futurehouse_dummy_env_crow(
+    admin_client: RestClient, task_data: TaskRequest
+):
+    tasks_to_do = [task_data, task_data]
+
+    results = admin_client.run_tasks_until_done(tasks_to_do)
+
+    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
+    assert all(task.status == "success" for task in results)
+
+
+@pytest.mark.timeout(300)
+@pytest.mark.flaky(reruns=3)
+@pytest.mark.asyncio
+async def test_arun_until_done_futurehouse_dummy_env_crow(
+    admin_client: RestClient, task_data: TaskRequest
+):
+    tasks_to_do = [task_data, task_data]
+
+    results = await admin_client.arun_tasks_until_done(tasks_to_do)
+
+    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
+    assert all(task.status == "success" for task in results)
+
+
+@pytest.mark.timeout(300)
+@pytest.mark.flaky(reruns=3)
+@pytest.mark.asyncio
+async def test_timeout_run_until_done_futurehouse_dummy_env_crow(
+    admin_client: RestClient, task_data: TaskRequest
+):
+    tasks_to_do = [task_data, task_data]
+
+    results = await admin_client.arun_tasks_until_done(
+        tasks_to_do, verbose=True, timeout=5, progress_bar=True
+    )
+
+    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
+    assert all(task.status != "success" for task in results), "Should not be success."
+    assert all(not isinstance(task, PQATaskResponse) for task in results), (
+        "Should be verbose."
+    )
+
+    results = admin_client.run_tasks_until_done(
+        tasks_to_do, verbose=True, timeout=5, progress_bar=True
+    )
+
+    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
+    assert all(task.status != "success" for task in results), "Should not be success."
+    assert all(not isinstance(task, PQATaskResponse) for task in results), (
+        "Should be verbose."
+    )
+
+
+def test_auth_refresh_flow(admin_client: RestClient):
+    refresh_calls = 0
+    func_calls = 0
+
+    def mock_run_auth(*args, **kwargs):
+        nonlocal refresh_calls
+        refresh_calls += 1
+        return f"fresh-token-{refresh_calls}"
+
+    @refresh_token_on_auth_error()
+    def test_func(self, *args):
+        nonlocal func_calls
+        func_calls += 1
+
+        if func_calls == 1:
+            raise AuthError(401, "Auth failed", None, None)
+        return "success"
+
+    with patch.object(admin_client, "_run_auth", mock_run_auth):
+        result = test_func(admin_client)
+
+        assert result == "success"
+        assert func_calls == 2, "Function should be called twice"
+        assert refresh_calls == 1, "Auth should be refreshed once"
+        assert admin_client.auth_jwt == "fresh-token-1"
+
+
+@pytest.mark.asyncio
+async def test_async_auth_refresh_flow(admin_client: RestClient):
+    refresh_calls = 0
+    func_calls = 0
+
+    def mock_run_auth(*args, **kwargs):
+        nonlocal refresh_calls
+        refresh_calls += 1
+        return f"fresh-token-{refresh_calls}"
+
+    @refresh_token_on_auth_error()
+    async def test_async_func(self, *args):
+        nonlocal func_calls
+        func_calls += 1
+
+        if func_calls == 1:
+            raise AuthError(401, "Auth failed", None, None)
+        await asyncio.sleep(1)
+        return "success"
+
+    with patch.object(admin_client, "_run_auth", mock_run_auth):
+        result = await test_async_func(admin_client)
+
+        assert result == "success"
+        assert func_calls == 2, "Function should be called twice"
+        assert refresh_calls == 1, "Auth should be refreshed once"
+        assert admin_client.auth_jwt == "fresh-token-1"

futurehouse_client-0.3.18.dev109/futurehouse_client/utils/__init__.py

@@ -1,3 +0,0 @@
-from .context import UserContext
-
-__all__ = ["UserContext"]

futurehouse_client-0.3.18.dev109/futurehouse_client/utils/context.py

@@ -1,16 +0,0 @@
-class UserContext:
-    """A context manager for storing user information from the initial request."""
-
-    _user_jwt = None
-
-    @classmethod
-    def set_user_jwt(cls, jwt: str) -> None:
-        cls._user_jwt = jwt
-
-    @classmethod
-    def get_user_jwt(cls) -> str | None:
-        return cls._user_jwt
-
-    @classmethod
-    def clear_user_jwt(cls) -> None:
-        cls._user_jwt = None

futurehouse_client-0.3.18.dev109/tests/test_rest.py

@@ -1,214 +0,0 @@
-import asyncio
-import os
-import time
-
-import pytest
-from futurehouse_client.clients import (
-    JobNames,
-    PQATaskResponse,
-    TaskResponseVerbose,
-)
-from futurehouse_client.clients.rest_client import RestClient, TaskFetchError
-from futurehouse_client.models.app import Stage, TaskRequest
-from futurehouse_client.models.rest import ExecutionStatus
-from pytest_subtests import SubTests
-
-ADMIN_API_KEY = os.environ["PLAYWRIGHT_ADMIN_API_KEY"]
-PUBLIC_API_KEY = os.environ["PLAYWRIGHT_PUBLIC_API_KEY"]
-TEST_MAX_POLLS = 100
-
-
-@pytest.mark.timeout(300)
-@pytest.mark.flaky(reruns=3)
-def test_futurehouse_dummy_env_crow():
-    client = RestClient(
-        stage=Stage.DEV,
-        api_key=ADMIN_API_KEY,
-    )
-
-    task_data = TaskRequest(
-        name=JobNames.from_string("dummy"),
-        query="How many moons does earth have?",
-    )
-    client.create_task(task_data)
-
-    while (task_status := client.get_task().status) in {"queued", "in progress"}:
-        time.sleep(5)
-
-    assert task_status == "success"
-
-
-def test_insufficient_permissions_request():
-    # Create a new instance so that cached credentials aren't reused
-    client = RestClient(
-        stage=Stage.DEV,
-        api_key=PUBLIC_API_KEY,
-    )
-    task_data = TaskRequest(
-        name=JobNames.from_string("dummy"),
-        query="How many moons does earth have?",
-    )
-
-    with pytest.raises(TaskFetchError) as exc_info:
-        client.create_task(task_data)
-
-    assert "Error creating task" in str(exc_info.value)
-
-
-@pytest.mark.timeout(300)
-@pytest.mark.asyncio
-async def test_job_response(subtests: SubTests):  # noqa: PLR0915
-    client = RestClient(
-        stage=Stage.DEV,
-        api_key=ADMIN_API_KEY,
-    )
-    task_data = TaskRequest(
-        name=JobNames.from_string("crow"),
-        query="How many moons does earth have?",
-    )
-    task_id = client.create_task(task_data)
-    atask_id = await client.acreate_task(task_data)
-
-    with subtests.test("Test TaskResponse with queued task"):
-        task_response = client.get_task(task_id)
-        assert task_response.status in {"queued", "in progress"}
-        assert task_response.job_name == task_data.name
-        assert task_response.query == task_data.query
-        task_response = await client.aget_task(atask_id)
-        assert task_response.status in {"queued", "in progress"}
-        assert task_response.job_name == task_data.name
-        assert task_response.query == task_data.query
-
-    for _ in range(TEST_MAX_POLLS):
-        task_response = client.get_task(task_id)
-        if task_response.status in ExecutionStatus.terminal_states():
-            break
-        await asyncio.sleep(5)
-
-    for _ in range(TEST_MAX_POLLS):
-        task_response = await client.aget_task(atask_id)
-        if task_response.status in ExecutionStatus.terminal_states():
-            break
-        await asyncio.sleep(5)
-
-    with subtests.test("Test PQA job response"):
-        task_response = client.get_task(task_id)
-        assert isinstance(task_response, PQATaskResponse)
-        # assert it has general fields
-        assert task_response.status == "success"
-        assert task_response.task_id is not None
-        assert task_data.name in task_response.job_name
-        assert task_data.query in task_response.query
-        # assert it has PQA specific fields
-        assert task_response.answer is not None
-        # assert it's not verbose
-        assert not hasattr(task_response, "environment_frame")
-        assert not hasattr(task_response, "agent_state")
-
-    with subtests.test("Test async PQA job response"):
-        task_response = await client.aget_task(atask_id)
-        assert isinstance(task_response, PQATaskResponse)
-        # assert it has general fields
-        assert task_response.status == "success"
-        assert task_response.task_id is not None
-        assert task_data.name in task_response.job_name
-        assert task_data.query in task_response.query
-        # assert it has PQA specific fields
-        assert task_response.answer is not None
-        # assert it's not verbose
-        assert not hasattr(task_response, "environment_frame")
-        assert not hasattr(task_response, "agent_state")
-
-    with subtests.test("Test task response with verbose"):
-        task_response = client.get_task(task_id, verbose=True)
-        assert isinstance(task_response, TaskResponseVerbose)
-        assert task_response.status == "success"
-        assert task_response.environment_frame is not None
-        assert task_response.agent_state is not None
-
-    with subtests.test("Test task async response with verbose"):
-        task_response = await client.aget_task(atask_id, verbose=True)
-        assert isinstance(task_response, TaskResponseVerbose)
-        assert task_response.status == "success"
-        assert task_response.environment_frame is not None
-        assert task_response.agent_state is not None
-
-
-@pytest.mark.timeout(300)
-@pytest.mark.flaky(reruns=3)
-def test_run_until_done_futurehouse_dummy_env_crow():
-    client = RestClient(
-        stage=Stage.DEV,
-        api_key=ADMIN_API_KEY,
-    )
-
-    task_data = TaskRequest(
-        name=JobNames.from_string("dummy"),
-        query="How many moons does earth have?",
-    )
-
-    tasks_to_do = [task_data, task_data]
-
-    results = client.run_tasks_until_done(tasks_to_do)
-
-    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
-    assert all(task.status == "success" for task in results)
-
-
-@pytest.mark.timeout(300)
-@pytest.mark.flaky(reruns=3)
-@pytest.mark.asyncio
-async def test_arun_until_done_futurehouse_dummy_env_crow():
-    client = RestClient(
-        stage=Stage.DEV,
-        api_key=ADMIN_API_KEY,
-    )
-
-    task_data = TaskRequest(
-        name=JobNames.from_string("dummy"),
-        query="How many moons does earth have?",
-    )
-
-    tasks_to_do = [task_data, task_data]
-
-    results = await client.arun_tasks_until_done(tasks_to_do)
-
-    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
-    assert all(task.status == "success" for task in results)
-
-
-@pytest.mark.timeout(300)
-@pytest.mark.flaky(reruns=3)
-@pytest.mark.asyncio
-async def test_timeout_run_until_done_futurehouse_dummy_env_crow():
-    client = RestClient(
-        stage=Stage.DEV,
-        api_key=ADMIN_API_KEY,
-    )
-
-    task_data = TaskRequest(
-        name=JobNames.from_string("dummy"),
-        query="How many moons does earth have?",
-    )
-
-    tasks_to_do = [task_data, task_data]
-
-    results = await client.arun_tasks_until_done(
-        tasks_to_do, verbose=True, timeout=5, progress_bar=True
-    )
-
-    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
-    assert all(task.status != "success" for task in results), "Should not be success."
-    assert all(not isinstance(task, PQATaskResponse) for task in results), (
-        "Should be verbose."
-    )
-
-    results = client.run_tasks_until_done(
-        tasks_to_do, verbose=True, timeout=5, progress_bar=True
-    )
-
-    assert len(results) == len(tasks_to_do), "Should return 2 tasks."
-    assert all(task.status != "success" for task in results), "Should not be success."
-    assert all(not isinstance(task, PQATaskResponse) for task in results), (
-        "Should be verbose."
-    )