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

52 statements  

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

1from datetime import datetime 

2from enum import StrEnum, auto 

3from typing import Literal, cast 

4 

5from openai import BaseModel 

6from pydantic import EmailStr 

7from sqlalchemy.ext.asyncio import AsyncAttrs 

8from sqlmodel import Field, Relationship, SQLModel 

9 

10from ._utils import now, random_id_with_prefix 

11 

12 

13class OrganizationUserRole(StrEnum): 

14 OWNER = auto() 

15 READER = auto() 

16 

17 

18class _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 

27class 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 

46class 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 

60class 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 

66class 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 

71class 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)