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

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/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/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