futurehouse-client 0.4.5.dev119__tar.gz → 0.4.5.dev160__tar.gz

{futurehouse_client-0.4.5.dev119/src/futurehouse_client.egg-info → futurehouse_client-0.4.5.dev160}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.4.5.dev119
+Version: 0.4.5.dev160
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 License:                                  Apache License

{futurehouse_client-0.4.5.dev119 → futurehouse_client-0.4.5.dev160}/src/futurehouse_client/__init__.py

@@ -8,6 +8,13 @@ from .models.app import (
     TaskResponse,
     TaskResponseVerbose,
 )
+from .models.job_event import (
+    CostComponent,
+    ExecutionType,
+    JobEventCreateRequest,
+    JobEventCreateResponse,
+    JobEventUpdateRequest,
+)
 from .utils.world_model_tools import (
     create_world_model_tool,
     make_world_model_tools,
@@ -15,9 +22,14 @@ from .utils.world_model_tools import (
 )
 
 __all__ = [
+    "CostComponent",
+    "ExecutionType",
     "FinchTaskResponse",
     "FutureHouseClient",
     "JobClient",
+    "JobEventCreateRequest",
+    "JobEventCreateResponse",
+    "JobEventUpdateRequest",
     "JobNames",
     "PQATaskResponse",
     "PhoenixTaskResponse",

{futurehouse_client-0.4.5.dev119 → futurehouse_client-0.4.5.dev160}/src/futurehouse_client/clients/rest_client.py

@@ -52,6 +52,11 @@ from futurehouse_client.models.app import (
     TaskResponseVerbose,
     TrajectoryQueryParams,
 )
+from futurehouse_client.models.job_event import (
+    JobEventCreateRequest,
+    JobEventCreateResponse,
+    JobEventUpdateRequest,
+)
 from futurehouse_client.models.rest import (
     DiscoveryResponse,
     ExecutionStatus,
@@ -160,6 +165,18 @@ class FileUploadError(RestClientError):
     """Raised when there's an error uploading a file."""
 
 
+class JobEventClientError(RestClientError):
+    """Raised when there's an error with job event operations."""
+
+
+class JobEventCreationError(JobEventClientError):
+    """Raised when there's an error creating a job event."""
+
+
+class JobEventUpdateError(JobEventClientError):
+    """Raised when there's an error updating a job event."""
+
+
 retry_if_connection_error = create_retry_if_connection_error(FileUploadError)
 
 DEFAULT_AGENT_TIMEOUT: int = 2400  # seconds
@@ -2609,6 +2626,176 @@ class RestClient(DataStorageMethods):
                 f"Error fetching discoveries for project: {e!r}"
             ) from e
 
+    @retry(
+        stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def create_job_event(
+        self, request: JobEventCreateRequest
+    ) -> JobEventCreateResponse:
+        """Create a new job event.
+
+        Args:
+            request: Job event creation request
+
+        Returns:
+            Job event creation response
+
+        Raises:
+            JobEventCreationError: If the API call fails
+        """
+        try:
+            response = self.client.post(
+                "/v0.1/job-events",
+                json=request.model_dump(exclude_none=True, mode="json"),
+            )
+            response.raise_for_status()
+            return JobEventCreateResponse(**response.json())
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.BAD_REQUEST:
+                raise JobEventCreationError(
+                    f"Invalid job event creation request: {e.response.text}."
+                ) from e
+            if e.response.status_code == codes.NOT_FOUND:
+                raise JobEventCreationError(
+                    f"Execution not found for job event creation: {e.response.text}."
+                ) from e
+            raise JobEventCreationError(
+                f"Error creating job event: {e.response.status_code} - {e.response.text}."
+            ) from e
+        except Exception as e:
+            raise JobEventCreationError(
+                f"An unexpected error occurred during job event creation: {e!r}."
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def acreate_job_event(
+        self, request: JobEventCreateRequest
+    ) -> JobEventCreateResponse:
+        """Asynchronously create a new job event.
+
+        Args:
+            request: Job event creation request
+
+        Returns:
+            Job event creation response
+
+        Raises:
+            JobEventCreationError: If the API call fails
+        """
+        try:
+            response = await self.async_client.post(
+                "/v0.1/job-events",
+                json=request.model_dump(exclude_none=True, mode="json"),
+            )
+            response.raise_for_status()
+            return JobEventCreateResponse(**response.json())
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.BAD_REQUEST:
+                raise JobEventCreationError(
+                    f"Invalid job event creation request: {e.response.text}."
+                ) from e
+            if e.response.status_code == codes.NOT_FOUND:
+                raise JobEventCreationError(
+                    f"Execution not found for job event creation: {e.response.text}."
+                ) from e
+            raise JobEventCreationError(
+                f"Error creating job event: {e.response.status_code} - {e.response.text}."
+            ) from e
+        except Exception as e:
+            raise JobEventCreationError(
+                f"An unexpected error occurred during job event creation: {e!r}."
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    def update_job_event(
+        self, job_event_id: UUID, request: JobEventUpdateRequest
+    ) -> None:
+        """Update an existing job event.
+
+        Args:
+            job_event_id: ID of the job event to update
+            request: Job event update request
+
+        Raises:
+            JobEventUpdateError: If the API call fails
+        """
+        try:
+            response = self.client.patch(
+                f"/v0.1/job-events/{job_event_id}",
+                json=request.model_dump(exclude_none=True, mode="json"),
+            )
+            response.raise_for_status()
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.NOT_FOUND:
+                raise JobEventUpdateError(
+                    f"Job event with ID {job_event_id} not found."
+                ) from e
+            if e.response.status_code == codes.BAD_REQUEST:
+                raise JobEventUpdateError(
+                    f"Invalid job event update request: {e.response.text}."
+                ) from e
+            raise JobEventUpdateError(
+                f"Error updating job event: {e.response.status_code} - {e.response.text}."
+            ) from e
+        except Exception as e:
+            raise JobEventUpdateError(
+                f"An unexpected error occurred during job event update: {e!r}."
+            ) from e
+
+    @retry(
+        stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=RETRY_MULTIPLIER, max=MAX_RETRY_WAIT),
+        retry=retry_if_connection_error,
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+    async def aupdate_job_event(
+        self, job_event_id: UUID, request: JobEventUpdateRequest
+    ) -> None:
+        """Asynchronously update an existing job event.
+
+        Args:
+            job_event_id: ID of the job event to update
+            request: Job event update request
+
+        Raises:
+            JobEventUpdateError: If the API call fails
+        """
+        try:
+            response = await self.async_client.patch(
+                f"/v0.1/job-events/{job_event_id}",
+                json=request.model_dump(exclude_none=True, mode="json"),
+            )
+            response.raise_for_status()
+        except HTTPStatusError as e:
+            if e.response.status_code == codes.NOT_FOUND:
+                raise JobEventUpdateError(
+                    f"Job event with ID {job_event_id} not found."
+                ) from e
+            if e.response.status_code == codes.BAD_REQUEST:
+                raise JobEventUpdateError(
+                    f"Invalid job event update request: {e.response.text}."
+                ) from e
+            raise JobEventUpdateError(
+                f"Error updating job event: {e.response.status_code} - {e.response.text}."
+            ) from e
+        except Exception as e:
+            raise JobEventUpdateError(
+                f"An unexpected error occurred during job event update: {e!r}."
+            ) from e
+
 
 def get_installed_packages() -> dict[str, str]:
     """Returns a dictionary of installed packages and their versions."""

{futurehouse_client-0.4.5.dev119 → futurehouse_client-0.4.5.dev160}/src/futurehouse_client/models/__init__.py

@@ -13,13 +13,25 @@ from .app import (
     TaskResponse,
     TaskResponseVerbose,
 )
+from .job_event import (
+    CostComponent,
+    ExecutionType,
+    JobEventCreateRequest,
+    JobEventCreateResponse,
+    JobEventUpdateRequest,
+)
 from .rest import TrajectoryPatchRequest, WorldModel, WorldModelResponse
 
 __all__ = [
     "AuthType",
+    "CostComponent",
     "DockerContainerConfiguration",
+    "ExecutionType",
     "FramePath",
     "JobDeploymentConfig",
+    "JobEventCreateRequest",
+    "JobEventCreateResponse",
+    "JobEventUpdateRequest",
     "PQATaskResponse",
     "RuntimeConfig",
     "Stage",

futurehouse_client-0.4.5.dev160/src/futurehouse_client/models/client.py

@@ -0,0 +1,72 @@
+from typing import Any, Generic, TypeAlias, TypeVar
+
+from aviary.message import Message
+from aviary.tools.base import Tool
+from ldp.data_structures import Transition
+from ldp.graph.ops import OpResult
+from pydantic import BaseModel, ConfigDict, Field, field_serializer
+
+T = TypeVar("T")
+
+
+# TODO: revisit this
+# unsure what crow states will return
+# need to revisit after we get more crows deployed
+class BaseState(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
+
+
+class BeforeTransitionState(BaseState):
+    current_state: Any = Field()
+    observations: list[Message] = Field()
+
+
+class InitialState(BaseState):
+    initial_state: Any = Field()
+
+
+class ASVState(BaseState, Generic[T]):
+    action: OpResult[T] = Field()
+    next_state: Any = Field()
+    value: float = Field()
+
+    @field_serializer("action")
+    def serialize_action(self, action: OpResult[T]) -> dict:
+        return action.to_dict()
+
+    @field_serializer("next_state")
+    def serialize_next_state(self, state: Any) -> str:
+        return str(state)
+
+
+class EnvResetState(BaseState):
+    observations: list[Message] = Field()
+    tools: list[Tool] = Field()
+
+
+class EnvStepState(BaseState):
+    observations: list[Message] = Field()
+    reward: float = Field()
+    done: bool = Field()
+    trunc: bool = Field()
+
+
+class TransitionState(BaseState):
+    transition: Transition = Field()
+
+    @field_serializer("transition")
+    def serialize_transition(self, transition: Transition) -> dict:
+        transition_data = transition.model_dump()
+        return transition_data | {
+            "action": transition.action.to_dict() if transition.action else None,
+        }
+
+
+StateType: TypeAlias = (
+    BeforeTransitionState
+    | InitialState
+    | ASVState
+    | EnvResetState
+    | EnvStepState
+    | TransitionState
+)

futurehouse_client-0.4.5.dev160/src/futurehouse_client/models/job_event.py

@@ -0,0 +1,75 @@
+"""Job event models for cost and usage tracking."""
+
+from datetime import datetime
+from enum import StrEnum, auto
+from typing import Any
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+
+
+class ExecutionType(StrEnum):
+    """Type of execution for job events."""
+
+    TRAJECTORY = auto()
+    SESSION = auto()
+
+
+class CostComponent(StrEnum):
+    """Cost component types for job events."""
+
+    LLM_USAGE = auto()
+    EXTERNAL_SERVICE = auto()
+    STEP = auto()
+
+
+class JobEventCreateRequest(BaseModel):
+    """Request model for creating a job event matching crow-service schema."""
+
+    execution_id: UUID = Field(description="UUID for trajectory_id or session_id")
+    execution_type: ExecutionType = Field(
+        description="Either 'TRAJECTORY' or 'SESSION'"
+    )
+    cost_component: CostComponent = Field(
+        description="Cost component: 'LLM_USAGE', 'EXTERNAL_SERVICE', or 'STEP'"
+    )
+    started_at: datetime = Field(description="Start time of the job event")
+    ended_at: datetime = Field(description="End time of the job event")
+    crow: str | None = Field(default=None, description="unique identifier for the crow")
+    amount_acu: float | None = Field(default=None, description="Cost amount in ACUs")
+    amount_usd: float | None = Field(default=None, description="Cost amount in USD")
+    rate: float | None = Field(default=None, description="Rate per token/call in USD")
+    input_token_count: int | None = Field(
+        default=None, description="Input token count for LLM calls"
+    )
+    completion_token_count: int | None = Field(
+        default=None, description="Completion token count for LLM calls"
+    )
+    metadata: dict[str, Any] | None = Field(default=None)
+
+
+class JobEventUpdateRequest(BaseModel):
+    """Request model for updating a job event matching crow-service schema."""
+
+    amount_acu: float | None = Field(default=None, description="Cost amount in ACUs")
+    amount_usd: float | None = Field(default=None, description="Cost amount in USD")
+    rate: float | None = Field(default=None, description="Rate per token/call in USD")
+    input_token_count: int | None = Field(
+        default=None, description="Input token count for LLM calls"
+    )
+    completion_token_count: int | None = Field(
+        default=None, description="Completion token count for LLM calls"
+    )
+    metadata: dict[str, Any] | None = Field(default=None)
+    started_at: datetime | None = Field(
+        default=None, description="Start time of the job event"
+    )
+    ended_at: datetime | None = Field(
+        default=None, description="End time of the job event"
+    )
+
+
+class JobEventCreateResponse(BaseModel):
+    """Response model for job event creation."""
+
+    id: UUID = Field(description="UUID of the created job event")

{futurehouse_client-0.4.5.dev119 → futurehouse_client-0.4.5.dev160}/src/futurehouse_client/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.4.5.dev119'
-__version_tuple__ = version_tuple = (0, 4, 5, 'dev119')
+__version__ = version = '0.4.5.dev160'
+__version_tuple__ = version_tuple = (0, 4, 5, 'dev160')
 
-__commit_id__ = commit_id = 'g1d92eec8f'
+__commit_id__ = commit_id = 'g7b8da3bce'

{futurehouse_client-0.4.5.dev119 → futurehouse_client-0.4.5.dev160/src/futurehouse_client.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: futurehouse-client
-Version: 0.4.5.dev119
+Version: 0.4.5.dev160
 Summary: A client for interacting with endpoints of the FutureHouse service.
 Author-email: FutureHouse technical staff <hello@futurehouse.org>
 License:                                  Apache License

{futurehouse_client-0.4.5.dev119 → futurehouse_client-0.4.5.dev160}/src/futurehouse_client.egg-info/SOURCES.txt

@@ -20,6 +20,7 @@ src/futurehouse_client/models/__init__.py
 src/futurehouse_client/models/app.py
 src/futurehouse_client/models/client.py
 src/futurehouse_client/models/data_storage_methods.py
+src/futurehouse_client/models/job_event.py
 src/futurehouse_client/models/rest.py
 src/futurehouse_client/utils/__init__.py
 src/futurehouse_client/utils/auth.py

{futurehouse_client-0.4.5.dev119 → futurehouse_client-0.4.5.dev160}/tests/test_rest.py

@@ -30,6 +30,8 @@ from futurehouse_client.clients import (
 )
 from futurehouse_client.clients.rest_client import (
     FileUploadError,
+    JobEventCreationError,
+    JobEventUpdateError,
     ProjectError,
     RestClient,
     RestClientError,
@@ -43,6 +45,12 @@ from futurehouse_client.models.app import (
     TaskRequest,
     TaskResponseVerbose,
 )
+from futurehouse_client.models.job_event import (
+    CostComponent,
+    ExecutionType,
+    JobEventCreateRequest,
+    JobEventUpdateRequest,
+)
 from futurehouse_client.models.rest import (
     ExecutionStatus,
     UserAgentRequestPostPayload,
@@ -1229,3 +1237,144 @@ def test_retry_logic_conditions(exception, should_retry, test_description):
     assert result == should_retry, (
         f"Expected {should_retry} for {test_description}, got {result}"
     )
+
+
+class TestJobEventOperations:
+    @pytest.fixture
+    def test_trajectory_id(
+        self, admin_client: RestClient, task_req: TaskRequest
+    ) -> str:
+        """Create a real trajectory for job event testing."""
+        return admin_client.create_task(task_req)
+
+    @pytest.fixture
+    def job_event_create_request(self, test_trajectory_id: str):
+        from datetime import datetime
+
+        return JobEventCreateRequest(
+            execution_id=UUID(test_trajectory_id),
+            execution_type=ExecutionType.TRAJECTORY,
+            cost_component=CostComponent.LLM_USAGE,
+            started_at=datetime.now(),
+            ended_at=datetime.now(),
+            crow=None,
+            amount_usd=0.005,
+            rate=0.0001,
+            input_token_count=100,
+            completion_token_count=50,
+            metadata={"model": "gpt-4o", "temperature": 0.7},
+        )
+
+    @pytest.fixture
+    def job_event_update_request(self):
+        return JobEventUpdateRequest(
+            amount_usd=0.007,
+            rate=0.00015,
+            input_token_count=120,
+            completion_token_count=60,
+            metadata={"model": "gpt-4o", "temperature": 0.5, "updated": True},
+        )
+
+    @pytest.mark.timeout(300)
+    def test_create_job_event_success(
+        self, admin_client: RestClient, job_event_create_request
+    ):
+        response = admin_client.create_job_event(job_event_create_request)
+        assert response.id is not None
+        assert isinstance(response.id, UUID)
+
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(300)
+    async def test_acreate_job_event_success(
+        self, admin_client: RestClient, job_event_create_request
+    ):
+        response = await admin_client.acreate_job_event(job_event_create_request)
+        assert response.id is not None
+        assert isinstance(response.id, UUID)
+
+    def test_create_job_event_execution_not_found(
+        self, admin_client: RestClient, job_event_create_request
+    ):
+        # Create a new request with a non-existent execution ID
+        invalid_request = JobEventCreateRequest(
+            execution_id=uuid4(),  # Non-existent ID
+            execution_type=job_event_create_request.execution_type,
+            cost_component=job_event_create_request.cost_component,
+            started_at=job_event_create_request.started_at,
+            ended_at=job_event_create_request.ended_at,
+            crow=job_event_create_request.crow,
+            amount_usd=job_event_create_request.amount_usd,
+            rate=job_event_create_request.rate,
+            input_token_count=job_event_create_request.input_token_count,
+            completion_token_count=job_event_create_request.completion_token_count,
+            metadata=job_event_create_request.metadata,
+        )
+        with pytest.raises(JobEventCreationError, match="Execution not found"):
+            admin_client.create_job_event(invalid_request)
+
+    @pytest.mark.asyncio
+    async def test_acreate_job_event_execution_not_found(
+        self, admin_client: RestClient, job_event_create_request
+    ):
+        # Create a new request with a non-existent execution ID
+        invalid_request = JobEventCreateRequest(
+            execution_id=uuid4(),  # Non-existent ID
+            execution_type=job_event_create_request.execution_type,
+            cost_component=job_event_create_request.cost_component,
+            started_at=job_event_create_request.started_at,
+            ended_at=job_event_create_request.ended_at,
+            crow=job_event_create_request.crow,
+            amount_usd=job_event_create_request.amount_usd,
+            rate=job_event_create_request.rate,
+            input_token_count=job_event_create_request.input_token_count,
+            completion_token_count=job_event_create_request.completion_token_count,
+            metadata=job_event_create_request.metadata,
+        )
+        with pytest.raises(JobEventCreationError, match="Execution not found"):
+            await admin_client.acreate_job_event(invalid_request)
+
+    @pytest.mark.timeout(300)
+    def test_update_job_event_success(
+        self,
+        admin_client: RestClient,
+        job_event_create_request,
+        job_event_update_request,
+    ):
+        create_response = admin_client.create_job_event(job_event_create_request)
+        job_event_id = create_response.id
+
+        # Should not raise an exception and return None
+        result = admin_client.update_job_event(job_event_id, job_event_update_request)
+        assert result is None
+
+    @pytest.mark.asyncio
+    @pytest.mark.timeout(300)
+    async def test_aupdate_job_event_success(
+        self,
+        admin_client: RestClient,
+        job_event_create_request,
+        job_event_update_request,
+    ):
+        create_response = await admin_client.acreate_job_event(job_event_create_request)
+        job_event_id = create_response.id
+
+        await admin_client.aupdate_job_event(job_event_id, job_event_update_request)
+
+    def test_update_job_event_not_found(
+        self, admin_client: RestClient, job_event_update_request
+    ):
+        job_event_id = uuid4()
+        with pytest.raises(
+            JobEventUpdateError, match=r"Job event with ID .* not found"
+        ):
+            admin_client.update_job_event(job_event_id, job_event_update_request)
+
+    @pytest.mark.asyncio
+    async def test_aupdate_job_event_not_found(
+        self, admin_client: RestClient, job_event_update_request
+    ):
+        job_event_id = uuid4()
+        with pytest.raises(
+            JobEventUpdateError, match=r"Job event with ID .* not found"
+        ):
+            await admin_client.aupdate_job_event(job_event_id, job_event_update_request)

futurehouse_client-0.4.5.dev119/src/futurehouse_client/models/client.py

@@ -1,128 +0,0 @@
-from typing import Any, Generic, TypeAlias, TypeVar
-
-from aviary.message import Message
-from aviary.tools.base import Tool
-from ldp.agent import Agent
-from ldp.data_structures import Transition
-from ldp.graph.ops import OpResult
-from pydantic import BaseModel, ConfigDict, Field, field_serializer
-
-T = TypeVar("T")
-
-
-# TODO: revisit this
-# unsure what crow states will return
-# need to revisit after we get more crows deployed
-class BaseState(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow")
-
-
-class BeforeTransitionState(BaseState):
-    current_state: Any = Field()
-    observations: list[Message] = Field()
-
-
-class InitialState(BaseState):
-    initial_state: Any = Field()
-
-
-class ASVState(BaseState, Generic[T]):
-    action: OpResult[T] = Field()
-    next_state: Any = Field()
-    value: float = Field()
-
-    @field_serializer("action")
-    def serialize_action(self, action: OpResult[T]) -> dict:
-        return action.to_dict()
-
-
-class EnvResetState(BaseState):
-    observations: list[Message] = Field()
-    tools: list[Tool] = Field()
-
-
-class EnvStepState(BaseState):
-    observations: list[Message] = Field()
-    reward: float = Field()
-    done: bool = Field()
-    trunc: bool = Field()
-
-
-class TransitionState(BaseState):
-    transition: Transition = Field()
-
-    @field_serializer("transition")
-    def serialize_transition(self, transition: Transition) -> dict:
-        transition_data = transition.model_dump()
-        return transition_data | {
-            "action": transition.action.to_dict() if transition.action else None,
-        }
-
-
-class GlobalState(BaseState):
-    agent: Agent | None = None
-    env: Any | None = None
-    agent_state: Any | None = None
-    next_agent_state: Any | None = None
-    observations: list = []
-    action: Any | None = None
-    value: float = 0.0
-    last_step_state: Transition | None = None
-
-    def update_observations(self, obs: list[Message]) -> list[Message]:
-        previous_observations = self.observations or []
-        self.observations = obs
-        return previous_observations
-
-    def store_step_state(self, step_state: Transition) -> None:
-        self.last_step_state = step_state
-
-    def update_trajectory_data(self, **kwargs) -> None:
-        for key, value in kwargs.items():
-            setattr(self, key, value)
-
-    def _get_safe_previous_observations(
-        self, current_obs: list[Message] | None = None
-    ) -> list[Message]:
-        if self.last_step_state:
-            last_step_state = self.last_step_state
-            if last_step_state.next_observation:
-                return last_step_state.next_observation
-        if self.observations:
-            return self.observations
-        return current_obs or []
-
-    def create_step_state(self, callback_type: str, **kwargs) -> Transition:
-        defaults = {
-            "timestep": getattr(self.agent, "_timestep", 0) if self.agent else 0,
-            "agent_state": self.agent_state,
-            "next_agent_state": self.next_agent_state or self.agent_state,
-            "observation": self._get_safe_previous_observations(),
-            "next_observation": self.observations or [],
-            "action": self.action,
-            "reward": 0.0,
-            "truncated": False,
-            "done": False,
-            "value": self.value or 0.0,
-            "metadata": {"callback_type": callback_type},
-        }
-
-        for key, value in kwargs.items():
-            if key == "metadata" and isinstance(value, dict):
-                if isinstance(defaults["metadata"], dict):
-                    defaults["metadata"].update(value)
-            else:
-                defaults[key] = value
-
-        return Transition(**defaults)
-
-
-StateType: TypeAlias = (
-    BeforeTransitionState
-    | InitialState
-    | ASVState
-    | EnvResetState
-    | EnvStepState
-    | TransitionState
-    | GlobalState
-)