File size: 5,164 Bytes
01d5a5d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
"""
Role realted data structure
"""
from dataclasses import dataclass
from typing import Dict, Any, Optional
from datetime import datetime
import uuid
import secrets
import string
from sqlalchemy import Column, Integer, String, DateTime, Boolean, Text
from lpm_kernel.common.repository.database_session import Base


def generate_role_uuid(prefix: str = "role") -> str:
    """
    generate role's uuid, combination of 6 letters and digits, not case sensitive

    Args:
        prefix: UUID's prefix, default is 'role'

    Returns:
        str: generated UUID, with the format of 'prefix_{6 letters and digits, not case sensitive}'
    """
    # only use lowercase letters and digits
    alphabet = string.ascii_lowercase + string.digits
    # generate a random string
    random_str = "".join(secrets.choice(alphabet) for _ in range(6))
    return f"{prefix}_{random_str}"


class Role(Base):
    """Role database structure"""

    __tablename__ = "roles"

    id = Column(Integer, primary_key=True)
    uuid = Column(
        String(64), nullable=False, unique=True, default=generate_role_uuid
    )  # add uuid field
    name = Column(String(100), nullable=False, unique=True)  # role name
    description = Column(String(500))  # role description
    system_prompt = Column(Text, nullable=False)  # role system prompt
    icon = Column(String(100))  # role icon
    is_active = Column(Boolean, default=True)
    enable_l0_retrieval = Column(Boolean, default=True)
    enable_l1_retrieval = Column(Boolean, default=True)
    create_time = Column(DateTime, nullable=False, default=datetime.now)
    update_time = Column(
        DateTime, nullable=False, default=datetime.now, onupdate=datetime.now
    )


@dataclass
class RoleDTO:
    """Role's data transfer object"""

    id: int
    uuid: str  # add uuid field
    name: str
    description: str
    system_prompt: str
    icon: Optional[str] = None
    is_active: bool = True
    enable_l0_retrieval: bool = True
    enable_l1_retrieval: bool = True
    create_time: datetime = None
    update_time: datetime = None

    @classmethod
    def from_model(cls, model: Role) -> "RoleDTO":
        """Create DTO from database model"""
        return cls(
            id=model.id,
            uuid=model.uuid,  # add uuid field
            name=model.name,
            description=model.description,
            system_prompt=model.system_prompt,
            icon=model.icon,
            is_active=model.is_active,
            enable_l0_retrieval=model.enable_l0_retrieval,
            enable_l1_retrieval=model.enable_l1_retrieval,
            create_time=model.create_time,
            update_time=model.update_time,
        )

    def to_dict(self) -> Dict[str, Any]:
        """transfer to dict"""
        return {
            "id": self.id,
            "uuid": self.uuid,  # add uuid field
            "name": self.name,
            "description": self.description,
            "system_prompt": self.system_prompt,
            "icon": self.icon,
            "is_active": self.is_active,
            "enable_l0_retrieval": self.enable_l0_retrieval,
            "enable_l1_retrieval": self.enable_l1_retrieval,
            "create_time": self.create_time.strftime("%Y-%m-%d %H:%M:%S")
            if self.create_time
            else None,
            "update_time": self.update_time.strftime("%Y-%m-%d %H:%M:%S")
            if self.update_time
            else None,
        }


@dataclass
class CreateRoleRequest:
    """create Role request object"""

    name: str
    description: str
    system_prompt: str
    icon: Optional[str] = None
    enable_l0_retrieval: Optional[bool] = True
    enable_l1_retrieval: Optional[bool] = True

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CreateRoleRequest":
        return cls(
            name=data["name"],
            description=data["description"],
            system_prompt=data["system_prompt"],
            icon=data.get("icon"),
            enable_l0_retrieval=data.get("enable_l0_retrieval", True),
            enable_l1_retrieval=data.get("enable_l1_retrieval", True),
        )


@dataclass
class UpdateRoleRequest:
    """update Role request object"""

    name: Optional[str] = None
    description: Optional[str] = None
    system_prompt: Optional[str] = None
    icon: Optional[str] = None
    is_active: Optional[bool] = None
    enable_l0_retrieval: Optional[bool] = None
    enable_l1_retrieval: Optional[bool] = None

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "UpdateRoleRequest":
        return cls(
            name=data.get("name"),
            description=data.get("description"),
            system_prompt=data.get("system_prompt"),
            icon=data.get("icon"),
            is_active=data.get("is_active"),
            enable_l0_retrieval=data.get("enable_l0_retrieval"),
            enable_l1_retrieval=data.get("enable_l1_retrieval"),
        )

@dataclass
class ShareRoleRequest:
    """share Role request object"""
    role_id: str
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ShareRoleRequest":
        return cls(
            role_id=data["role_id"]
        )