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
« 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
5from openai import BaseModel
6from pydantic import EmailStr
7from sqlalchemy.ext.asyncio import AsyncAttrs
8from sqlmodel import Field, Relationship, SQLModel
10from ._utils import now, random_id_with_prefix
13class OrganizationUserRole(StrEnum):
14 OWNER = auto()
15 READER = auto()
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
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)
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 )
46class User(AsyncAttrs, SQLModel, table=True):
47 """User model."""
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")
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")
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"
71class APIKey(AsyncAttrs, SQLModel, table=True):
72 """API key model.
74 API key is used for authenticating the user.
75 """
77 __tablename__ = "api_key" # type: ignore
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)