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

41 statements  

« prev     ^ index     » next       coverage.py v7.6.8, created at 2024-12-06 09:34 +0800

1from datetime import datetime 

2from typing import Annotated, Literal 

3 

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

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

6from pydantic import field_serializer 

7from sqlalchemy.ext.asyncio import AsyncAttrs 

8from sqlmodel import Enum, Field, Relationship 

9 

10from .._metadata import WithMetadata 

11from .._types import as_sa_type 

12from .._utils import now, random_id_with_prefix 

13from .assistant import Assistant 

14from .run import Run 

15from .thread import Thread 

16 

17 

18class 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')