futurehouse-client 0.3.18.dev186__py3-none-any.whl → 0.3.19__py3-none-any.whl

futurehouse_client/models/app.py

@@ -1,3 +1,4 @@
+import copy
 import json
 import os
 import re
@@ -675,7 +676,8 @@ class TaskResponse(BaseModel):
 
     @model_validator(mode="before")
     @classmethod
-    def validate_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
+    def validate_fields(cls, original_data: Mapping[str, Any]) -> Mapping[str, Any]:
+        data = copy.deepcopy(original_data)  # Avoid mutating the original data
         # Extract fields from environment frame state
         if not isinstance(data, dict):
             return data
@@ -690,7 +692,72 @@ class TaskResponse(BaseModel):
         return data
 
 
+class PhoenixTaskResponse(TaskResponse):
+    """
+    Response scheme for tasks executed with Phoenix.
+
+    Additional fields:
+        answer: Final answer from Phoenix
+    """
+
+    model_config = ConfigDict(extra="ignore")
+    answer: str | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_phoenix_fields(
+        cls, original_data: Mapping[str, Any]
+    ) -> Mapping[str, Any]:
+        data = copy.deepcopy(original_data)
+        if not isinstance(data, dict):
+            return data
+        if not (env_frame := data.get("environment_frame", {})):
+            return data
+        state = env_frame.get("state", {}).get("state", {})
+        data["answer"] = state.get("answer")
+        return data
+
+
+class FinchTaskResponse(TaskResponse):
+    """
+    Response scheme for tasks executed with Finch.
+
+    Additional fields:
+        answer: Final answer from Finch
+        notebook: a dictionary with `cells` and `metadata` regarding the notebook content
+    """
+
+    model_config = ConfigDict(extra="ignore")
+    answer: str | None = None
+    notebook: dict[str, Any] | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_finch_fields(
+        cls, original_data: Mapping[str, Any]
+    ) -> Mapping[str, Any]:
+        data = copy.deepcopy(original_data)
+        if not isinstance(data, dict):
+            return data
+        if not (env_frame := data.get("environment_frame", {})):
+            return data
+        state = env_frame.get("state", {}).get("state", {})
+        data["answer"] = state.get("answer")
+        data["notebook"] = state.get("nb_state")
+        return data
+
+
 class PQATaskResponse(TaskResponse):
+    """
+    Response scheme for tasks executed with PQA.
+
+    Additional fields:
+        answer: Final answer from PQA
+        formatted_answer: Formatted answer from PQA
+        answer_reasoning: Reasoning used to generate the final answer, if available
+        has_successful_answer: Whether the answer is successful
+    """
+
     model_config = ConfigDict(extra="ignore")
 
     answer: str | None = None
@@ -702,7 +769,8 @@ class PQATaskResponse(TaskResponse):
 
     @model_validator(mode="before")
     @classmethod
-    def validate_pqa_fields(cls, data: Mapping[str, Any]) -> Mapping[str, Any]:
+    def validate_pqa_fields(cls, original_data: Mapping[str, Any]) -> Mapping[str, Any]:
+        data = copy.deepcopy(original_data)  # Avoid mutating the original data
         if not isinstance(data, dict):
             return data
         if not (env_frame := data.get("environment_frame", {})):

futurehouse_client/utils/auth.py

@@ -1,107 +1,92 @@
-import asyncio
 import logging
-from collections.abc import Callable, Coroutine
-from functools import wraps
-from typing import Any, Final, Optional, ParamSpec, TypeVar, overload
+from collections.abc import Collection, Generator
+from typing import ClassVar, Final
 
 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
+from futurehouse_client.models.app import APIKeyPayload, AuthType
 
+logger = logging.getLogger(__name__)
 
-@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
+INVALID_REFRESH_TYPE_MSG: Final[str] = (
+    "API key auth is required to refresh auth tokens."
+)
+JWT_TOKEN_CACHE_EXPIRY: int = 300  # seconds
+
+
+def _run_auth(
+    client: httpx.Client,
+    auth_type: AuthType = AuthType.API_KEY,
+    api_key: str | None = None,
+    jwt: str | None = None,
+) -> str:
+    auth_payload: APIKeyPayload | None
+    if auth_type == AuthType.API_KEY:
+        auth_payload = APIKeyPayload(api_key=api_key)
+    elif auth_type == AuthType.JWT:
+        auth_payload = None
+    try:
+        if auth_payload:
+            response = client.post("/auth/login", json=auth_payload.model_dump())
+            response.raise_for_status()
+            token_data = response.json()
+        elif jwt:
+            token_data = {"access_token": jwt, "expires_in": JWT_TOKEN_CACHE_EXPIRY}
+        else:
+            raise ValueError("JWT token required for JWT authentication.")
+
+        return token_data["access_token"]
+    except Exception as e:
+        raise Exception("Failed to authenticate") from e  # noqa: TRY002
+
+
+class RefreshingJWT(httpx.Auth):
+    """Automatically (re-)inject a JWT and transparently retry exactly once when we hit a 401/403."""
+
+    RETRY_STATUSES: ClassVar[Collection[httpx.codes]] = {
+        httpx.codes.UNAUTHORIZED,
+        httpx.codes.FORBIDDEN,
+    }
+
+    def __init__(
+        self,
+        auth_client: httpx.Client,
+        auth_type: AuthType = AuthType.API_KEY,
+        api_key: str | None = None,
+        jwt: str | None = None,
+    ):
+        self.auth_type = auth_type
+        self.auth_client = auth_client
+        self.api_key = api_key
+        self._jwt = _run_auth(
+            client=auth_client,
+            jwt=jwt,
+            auth_type=auth_type,
+            api_key=api_key,
+        )
+
+    def refresh_token(self) -> None:
+        if self.auth_type == AuthType.JWT:
+            logger.error(INVALID_REFRESH_TYPE_MSG)
+            raise ValueError(INVALID_REFRESH_TYPE_MSG)
+        self._jwt = _run_auth(
+            client=self.auth_client,
+            auth_type=self.auth_type,
+            api_key=self.api_key,
+        )
+
+    def auth_flow(
+        self, request: httpx.Request
+    ) -> Generator[httpx.Request, httpx.Response, None]:
+        request.headers["Authorization"] = f"Bearer {self._jwt}"
+        response = yield request
+
+        # If it failed, refresh once and replay the request
+        if response.status_code in self.RETRY_STATUSES:
+            logger.info(
+                "Received %s, refreshing token and retrying …",
+                response.status_code,
+            )
+            self.refresh_token()
+            request.headers["Authorization"] = f"Bearer {self._jwt}"
+            yield request  # second (and final) attempt, again or use a while loop

{futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.3.18.dev186
+Version: 0.3.19
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 Classifier: Operating System :: OS Independent
@@ -8,10 +8,9 @@ Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python
-Requires-Python: <3.13,>=3.11
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 Requires-Dist: cloudpickle
-Requires-Dist: dm-tree<0.1.9
 Requires-Dist: fhaviary
 Requires-Dist: httpx
 Requires-Dist: ldp>=0.22.0

futurehouse_client-0.3.19.dist-info/RECORD

@@ -0,0 +1,18 @@
+futurehouse_client/__init__.py,sha256=BztM_ntbgmIEjzvnBWcvPhvLjM8xGDFCK0Upf3-nIn8,488
+futurehouse_client/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+futurehouse_client/clients/__init__.py,sha256=-HXNj-XJ3LRO5XM6MZ709iPs29YpApss0Q2YYg1qMZw,280
+futurehouse_client/clients/job_client.py,sha256=JgB5IUAyCmnhGRsYc3bgKldA-lkM1JLwHRwwUeOCdus,11944
+futurehouse_client/clients/rest_client.py,sha256=3wfVz6d2KuRQUr_nms7P25yVR6aTjsRrSkqmVs55soA,54552
+futurehouse_client/models/__init__.py,sha256=5x-f9AoM1hGzJBEHcHAXSt7tPeImST5oZLuMdwp0mXc,554
+futurehouse_client/models/app.py,sha256=VCtg0ygd-TSrR6DtfljTBt9jnl1eBNal8UXHFdkDg88,28587
+futurehouse_client/models/client.py,sha256=n4HD0KStKLm6Ek9nL9ylP-bkK10yzAaD1uIDF83Qp_A,1828
+futurehouse_client/models/rest.py,sha256=lgwkMIXz0af-49BYSkKeS7SRqvN3motqnAikDN4YGTc,789
+futurehouse_client/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+futurehouse_client/utils/auth.py,sha256=tgWELjKfg8eWme_qdcRmc8TjQN9DVZuHHaVXZNHLchk,2960
+futurehouse_client/utils/general.py,sha256=A_rtTiYW30ELGEZlWCIArO7q1nEmqi8hUlmBRYkMQ_c,767
+futurehouse_client/utils/module_utils.py,sha256=aFyd-X-pDARXz9GWpn8SSViUVYdSbuy9vSkrzcVIaGI,4955
+futurehouse_client/utils/monitoring.py,sha256=UjRlufe67kI3VxRHOd5fLtJmlCbVA2Wqwpd4uZhXkQM,8728
+futurehouse_client-0.3.19.dist-info/METADATA,sha256=FbtQGStv4salVccxR5wtpdlGbufSqxoiCtM44qDOHJs,12731
+futurehouse_client-0.3.19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+futurehouse_client-0.3.19.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
+futurehouse_client-0.3.19.dist-info/RECORD,,

{futurehouse_client-0.3.18.dev186.dist-info → futurehouse_client-0.3.19.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.7.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

futurehouse_client-0.3.18.dev186.dist-info/RECORD

@@ -1,17 +0,0 @@
-futurehouse_client/__init__.py,sha256=ddxO7JE97c6bt7LjNglZZ2Ql8bYCGI9laSFeh9MP6VU,344
-futurehouse_client/clients/__init__.py,sha256=tFWqwIAY5PvwfOVsCje4imjTpf6xXNRMh_UHIKVI1_0,320
-futurehouse_client/clients/job_client.py,sha256=uNkqQbeZw7wbA0qDWcIOwOykrosza-jev58paJZ_mbA,11150
-futurehouse_client/clients/rest_client.py,sha256=W9ASP1ZKYS7UL5J9b-Km77YXEiDQ9hCf4X_9PqaZZZc,47914
-futurehouse_client/models/__init__.py,sha256=5x-f9AoM1hGzJBEHcHAXSt7tPeImST5oZLuMdwp0mXc,554
-futurehouse_client/models/app.py,sha256=w_1e4F0IiC-BKeOLqYkABYo4U-Nka1S-F64S_eHB2KM,26421
-futurehouse_client/models/client.py,sha256=n4HD0KStKLm6Ek9nL9ylP-bkK10yzAaD1uIDF83Qp_A,1828
-futurehouse_client/models/rest.py,sha256=lgwkMIXz0af-49BYSkKeS7SRqvN3motqnAikDN4YGTc,789
-futurehouse_client/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-futurehouse_client/utils/auth.py,sha256=Lq9mjSGc7iuRP6fmLICCS6KjzLHN6-tJUuhYp0XXrkE,3342
-futurehouse_client/utils/general.py,sha256=A_rtTiYW30ELGEZlWCIArO7q1nEmqi8hUlmBRYkMQ_c,767
-futurehouse_client/utils/module_utils.py,sha256=aFyd-X-pDARXz9GWpn8SSViUVYdSbuy9vSkrzcVIaGI,4955
-futurehouse_client/utils/monitoring.py,sha256=UjRlufe67kI3VxRHOd5fLtJmlCbVA2Wqwpd4uZhXkQM,8728
-futurehouse_client-0.3.18.dev186.dist-info/METADATA,sha256=PvjehEQZu2ihl7kG1uDvWJVUxyYbV7J-VmAe42Ml3zo,12767
-futurehouse_client-0.3.18.dev186.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-futurehouse_client-0.3.18.dev186.dist-info/top_level.txt,sha256=TRuLUCt_qBnggdFHCX4O_BoCu1j2X43lKfIZC-ElwWY,19
-futurehouse_client-0.3.18.dev186.dist-info/RECORD,,