Coverage for src/fastoai/models/generated/run_step.py: 85%

1 from datetime import datetime

2 from typing import Annotated, Literal

3

4 from openai.types.beta.threads.runs.run_step import LastError, StepDetails, Usage

5 from openai.types.beta.threads.runs.run_step import RunStep as _RunStep

6 from pydantic import field_serializer

7 from sqlalchemy.ext.asyncio import AsyncAttrs

8 from sqlmodel import Enum, Field, Relationship

9

10 from .._metadata import WithMetadata

11 from .._types import as_sa_type

12 from .._utils import now, random_id_with_prefix

13 from .assistant import Assistant

14 from .run import Run

15 from .thread import Thread

16

17

18 class RunStep(AsyncAttrs, WithMetadata, table=True):

19 __tablename__ = 'step'

20 id: Annotated[str, Field(primary_key=True, default_factory=random_id_with_prefix('step_'))]

21 assistant_id: Annotated[str, Field(foreign_key='assistant.id')]

22 cancelled_at: datetime | None = None

23 completed_at: datetime | None = None

24 created_at: Annotated[datetime, Field(default_factory=now)]

25 expired_at: datetime | None = None

26 failed_at: datetime | None = None

27 last_error: Annotated[LastError | None, Field(sa_type=as_sa_type(LastError), nullable=True)] = None

28 run_id: Annotated[str, Field(foreign_key='run.id')]

29 status: Annotated[Literal['in_progress', 'cancelled', 'failed', 'completed', 'expired'], Field(sa_type=Enum('in_progress', 'cancelled', 'failed', 'completed', 'expired'))]

30 step_details: Annotated[StepDetails, Field(sa_type=as_sa_type(StepDetails))]

31 thread_id: Annotated[str, Field(foreign_key='thread.id')]

32 type: Annotated[Literal['message_creation', 'tool_calls'], Field(sa_type=Enum('message_creation', 'tool_calls'))]

33 usage: Annotated[Usage | None, Field(sa_type=as_sa_type(Usage), nullable=True)] = None

34

35 async def to_openai_model(self) -> _RunStep:

36 value = self.model_dump(by_alias=True)

37 value['object'] = 'thread.run.step'

38 return _RunStep.model_validate(value)

39

40 @field_serializer('cancelled_at', 'completed_at', 'created_at', 'expired_at', 'failed_at')

41 def serialize_datetime(self, dt: datetime | None) -> int | None:

42 if dt is None:

43 return None

44 return int(dt.timestamp())

45 assistant: Assistant = Relationship(back_populates='steps')

46 run: Run = Relationship(back_populates='steps')

47 thread: Thread = Relationship(back_populates='steps')