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

1 from datetime import datetime

2 from typing import Annotated, Literal

3

4 from openai.types.beta.threads.message import (

5 Attachment,

6 IncompleteDetails,

7)

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

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

10 from pydantic import field_serializer

11 from sqlalchemy.ext.asyncio import AsyncAttrs

12 from sqlmodel import Enum, Field, Relationship

13

14 from .._metadata import WithMetadata

15 from .._types import as_sa_type

16 from .._utils import now, random_id_with_prefix

17 from .assistant import Assistant

18 from .run import Run

19 from .thread import Thread

20

21

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