futurehouse-client 0.3.18.dev186__tar.gz → 0.3.19__tar.gz

{futurehouse_client-0.3.18.dev186 → futurehouse_client-0.3.19}/PKG-INFO

@@ -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.18.dev186 → futurehouse_client-0.3.19}/futurehouse_client/__init__.py

@@ -1,12 +1,22 @@
 from .clients.job_client import JobClient, JobNames
-from .clients.rest_client import PQATaskResponse, TaskResponse, TaskResponseVerbose
 from .clients.rest_client import RestClient as FutureHouseClient
+from .models.app import (
+    FinchTaskResponse,
+    PhoenixTaskResponse,
+    PQATaskResponse,
+    TaskRequest,
+    TaskResponse,
+    TaskResponseVerbose,
+)
 
 __all__ = [
+    "FinchTaskResponse",
     "FutureHouseClient",
     "JobClient",
     "JobNames",
     "PQATaskResponse",
+    "PhoenixTaskResponse",
+    "TaskRequest",
     "TaskResponse",
     "TaskResponseVerbose",
 ]

{futurehouse_client-0.3.18.dev186 → futurehouse_client-0.3.19}/futurehouse_client/clients/__init__.py

@@ -1,12 +1,11 @@
 from .job_client import JobClient, JobNames
-from .rest_client import PQATaskResponse, TaskResponse, TaskResponseVerbose
 from .rest_client import RestClient as FutureHouseClient
+from .rest_client import TaskResponse, TaskResponseVerbose
 
 __all__ = [
     "FutureHouseClient",
     "JobClient",
     "JobNames",
-    "PQATaskResponse",
     "TaskResponse",
     "TaskResponseVerbose",
 ]

{futurehouse_client-0.3.18.dev186 → futurehouse_client-0.3.19}/futurehouse_client/clients/job_client.py

@@ -8,7 +8,13 @@ from aviary.env import Frame
 from pydantic import BaseModel
 from tenacity import before_sleep_log, retry, stop_after_attempt, wait_exponential
 
-from futurehouse_client.models.app import Stage
+from futurehouse_client.models.app import (
+    FinchTaskResponse,
+    PhoenixTaskResponse,
+    PQATaskResponse,
+    Stage,
+    TaskResponse,
+)
 from futurehouse_client.models.rest import (
     FinalEnvironmentRequest,
     StoreAgentStatePostRequest,
@@ -31,6 +37,19 @@ class JobNames(StrEnum):
     DUMMY = "job-futurehouse-dummy-env"
     PHOENIX = "job-futurehouse-phoenix"
     FINCH = "job-futurehouse-data-analysis-crow-high"
+    CHIMP = "job-futurehouse-chimp"
+
+    @classmethod
+    def _get_response_mapping(cls) -> dict[str, type[TaskResponse]]:
+        return {
+            cls.CROW: PQATaskResponse,
+            cls.FALCON: PQATaskResponse,
+            cls.OWL: PQATaskResponse,
+            cls.CHIMP: PQATaskResponse,
+            cls.PHOENIX: PhoenixTaskResponse,
+            cls.FINCH: FinchTaskResponse,
+            cls.DUMMY: TaskResponse,
+        }
 
     @classmethod
     def from_stage(cls, job_name: str, stage: Stage | None = None) -> str:
@@ -52,6 +71,13 @@ class JobNames(StrEnum):
                 f"Invalid job name: {job_name}. \nOptions are: {', '.join([name.name for name in cls])}"
             ) from e
 
+    @staticmethod
+    def get_response_object_from_job(job_name: str) -> type[TaskResponse]:
+        return JobNames._get_response_mapping()[job_name]
+
+    def get_response_object(self) -> type[TaskResponse]:
+        return self._get_response_mapping()[self.name]
+
 
 class JobClient:
     REQUEST_TIMEOUT: ClassVar[float] = 30.0  # sec