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

39 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.message import ( 

5 Attachment, 

6 IncompleteDetails, 

7) 

8from openai.types.beta.threads.message import Message as _Message 

9from openai.types.beta.threads.message_content import MessageContent 

10from pydantic import field_serializer 

11from sqlalchemy.ext.asyncio import AsyncAttrs 

12from sqlmodel import Enum, Field, Relationship 

13 

14from .._metadata import WithMetadata 

15from .._types import as_sa_type 

16from .._utils import now, random_id_with_prefix 

17from .assistant import Assistant 

18from .run import Run 

19from .thread import Thread 

20 

21 

22class Message(AsyncAttrs, WithMetadata, table=True): 

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

24 assistant_id: Annotated[str | None, Field(foreign_key='assistant.id', nullable=True)] = None 

25 attachments: Annotated[list[Attachment] | None, Field(sa_type=as_sa_type(list[Attachment]), nullable=True)] = None 

26 completed_at: datetime | None = None 

27 content: Annotated[list[MessageContent], Field(default_factory=list, sa_type=as_sa_type(list[MessageContent]))] 

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

29 incomplete_at: datetime | None = None 

30 incomplete_details: Annotated[IncompleteDetails | None, Field(sa_type=as_sa_type(IncompleteDetails), nullable=True)] = None 

31 role: Annotated[Literal['user', 'assistant'], Field(sa_type=Enum('user', 'assistant'))] 

32 run_id: Annotated[str | None, Field(foreign_key='run.id', nullable=True)] = None 

33 status: Annotated[Literal['in_progress', 'incomplete', 'completed'], Field(sa_type=Enum('in_progress', 'incomplete', 'completed'))] 

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

35 

36 async def to_openai_model(self) -> _Message: 

37 value = self.model_dump(by_alias=True) 

38 value['object'] = 'thread.message' 

39 return _Message.model_validate(value) 

40 

41 @field_serializer('completed_at', 'created_at', 'incomplete_at') 

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

43 if dt is None: 

44 return None 

45 return int(dt.timestamp()) 

46 assistant: Assistant | None = Relationship(back_populates='messages') 

47 run: Run | None = Relationship(back_populates='messages') 

48 thread: Thread = Relationship(back_populates='messages')