Coverage for src/fastoai/models/user.py: 96%

1 from datetime import datetime

2 from enum import StrEnum, auto

3 from typing import Literal, cast

4

5 from openai import BaseModel

6 from pydantic import EmailStr

7 from sqlalchemy.ext.asyncio import AsyncAttrs

8 from sqlmodel import Field, Relationship, SQLModel

9

10 from ._utils import now, random_id_with_prefix

11

12

13 class OrganizationUserRole(StrEnum):

14 OWNER = auto()

15 READER = auto()

16

17

18 class _OrganizationUser(BaseModel):

19 object: Literal["organization.user"] = "organization.user"

20 id: str

21 name: str

22 email: EmailStr

23 role: OrganizationUserRole

24 added_at: int

25

26

27 class OrganizationUser(AsyncAttrs, SQLModel, table=True):

28 organization_id: str = Field(foreign_key="organization.id", primary_key=True)

29 organization: "Organization" = Relationship(back_populates="members")

30 user_id: str = Field(foreign_key="user.id", primary_key=True)

31 user: "User" = Relationship(back_populates="organization_users")

32 role: OrganizationUserRole

33 added_at: datetime = Field(default_factory=now)

34

35 async def to_openai_model(self) -> _OrganizationUser:

36 user = cast(User, await self.awaitable_attrs.user)

37 return _OrganizationUser(

38 id=self.user_id,

39 name=user.name,

40 email=user.email,

41 role=self.role,

42 added_at=int(self.added_at.timestamp()),

43 )

44

45

46 class User(AsyncAttrs, SQLModel, table=True):

47 """User model."""

48

49 id: str = Field(default_factory=random_id_with_prefix("user_"), primary_key=True)

50 name: str

51 password: str

52 email: EmailStr

53 phone: str | None = None

54 is_active: bool = True

55 created_at: datetime = Field(default_factory=now)

56 api_keys: list["APIKey"] = Relationship(back_populates="user")

57 organization_users: OrganizationUser = Relationship(back_populates="user")

58

59

60 class Organization(AsyncAttrs, SQLModel, table=True):

61 id: str = Field(default_factory=random_id_with_prefix("org-"), primary_key=True)

62 name: str = "Personal"

63 members: list[OrganizationUser] = Relationship(back_populates="organization")

64

65

66 class Project(AsyncAttrs, SQLModel, table=True):

67 id: str = Field(default_factory=random_id_with_prefix("proj_"), primary_key=True)

68 name: str = "Default project"

69

70

71 class APIKey(AsyncAttrs, SQLModel, table=True):

72 """API key model.

73

74 API key is used for authenticating the user.

75 """

76

77 __tablename__ = "api_key" # type: ignore

78

79 id: str = Field(default_factory=random_id_with_prefix("sk-"), primary_key=True)

80 user_id: str = Field(foreign_key="user.id")

81 user: User = Relationship(back_populates="api_keys")

82 name: str = Field(default="New API key", max_length=255)

83 created_at: datetime = Field(default_factory=now)