From aae5851178fe4e07e948f3e19528b4a4c03f9d31 Mon Sep 17 00:00:00 2001 From: Jeremy Yin Date: Tue, 19 Mar 2024 20:29:22 +0800 Subject: [PATCH] refactor: model --- simplylab/main.py | 10 ++-- simplylab/model/__init__.py | 0 simplylab/model/entity.py | 7 +++ simplylab/model/req.py | 15 ++++++ simplylab/model/res.py | 30 +++++++++++ simplylab/{entity.py => model/table.py} | 69 ++----------------------- simplylab/providers/chat.py | 3 +- simplylab/providers/user.py | 2 +- simplylab/services/__init__.py | 2 +- simplylab/services/chat.py | 8 ++- 10 files changed, 71 insertions(+), 75 deletions(-) create mode 100644 simplylab/model/__init__.py create mode 100644 simplylab/model/entity.py create mode 100644 simplylab/model/req.py create mode 100644 simplylab/model/res.py rename simplylab/{entity.py => model/table.py} (71%) diff --git a/simplylab/main.py b/simplylab/main.py index 22f9170..b31cfe4 100644 --- a/simplylab/main.py +++ b/simplylab/main.py @@ -8,12 +8,10 @@ from loguru import logger from starlette.responses import JSONResponse from simplylab.database import Database -from simplylab.entity import GetAiChatResponseInput, Context -from simplylab.entity import GetAiChatResponseOutput -from simplylab.entity import GetUserChatHistoryInput -from simplylab.entity import GetUserChatHistoryOutput -from simplylab.entity import GetChatStatusTodayInput -from simplylab.entity import GetChatStatusTodayOutput +from simplylab.model.entity import Context +from simplylab.model.req import GetAiChatResponseInput, GetUserChatHistoryInput, GetChatStatusTodayInput +from simplylab.model.res import GetAiChatResponseOutput, GetChatStatusTodayOutput +from simplylab.model.res import GetUserChatHistoryOutput from simplylab.error import Error, UserNotFoundError from simplylab.providers import Providers from simplylab.services import Services diff --git a/simplylab/model/__init__.py b/simplylab/model/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/simplylab/model/entity.py b/simplylab/model/entity.py new file mode 100644 index 0000000..fa2b954 --- /dev/null +++ b/simplylab/model/entity.py @@ -0,0 +1,7 @@ +from pydantic import BaseModel + +from simplylab.model.table import User + + +class Context(BaseModel): + user: User diff --git a/simplylab/model/req.py b/simplylab/model/req.py new file mode 100644 index 0000000..7143033 --- /dev/null +++ b/simplylab/model/req.py @@ -0,0 +1,15 @@ +from pydantic import BaseModel + + +class GetAiChatResponseInput(BaseModel): + message: str + user_name: str + + +class GetUserChatHistoryInput(BaseModel): + user_name: str + last_n: int + + +class GetChatStatusTodayInput(BaseModel): + user_name: str diff --git a/simplylab/model/res.py b/simplylab/model/res.py new file mode 100644 index 0000000..715f727 --- /dev/null +++ b/simplylab/model/res.py @@ -0,0 +1,30 @@ +from pydantic import BaseModel + + +class GetAiChatResponseOutput(BaseModel): + response: str + + +class GetChatStatusTodayOutput(BaseModel): + user_name: str + chat_cnt: int + + +class UserChatMessage(BaseModel): + type: str + text: str + + +type GetUserChatHistoryOutput = list[UserChatMessage] + + +# class UserConversationMessages(BaseModel): +# user_id: ObjectIdField = Field() +# user_name: str = Field() +# conversation_id: ObjectIdField = Field(default=None) +# title: str = Field() +# created_at: datetime.datetime = Field(default_factory=datetime.datetime.now) +# created_by: ObjectIdField = Field() +# updated_at: Optional[datetime.datetime] = Field(default=None) +# updated_by: Optional[ObjectIdField] = Field(default=None) +# messages: list[Message] diff --git a/simplylab/entity.py b/simplylab/model/table.py similarity index 71% rename from simplylab/entity.py rename to simplylab/model/table.py index bdfc43c..d20e156 100644 --- a/simplylab/entity.py +++ b/simplylab/model/table.py @@ -1,49 +1,13 @@ import datetime from enum import Enum -from typing import Optional, Annotated, Any +from typing import Annotated, Any, Optional -from bson import ObjectId -from pydantic import BaseModel, Field, BeforeValidator, ConfigDict - - -class GetAiChatResponseInput(BaseModel): - message: str - user_name: str - - -class GetAiChatResponseOutput(BaseModel): - response: str - - -class GetUserChatHistoryInput(BaseModel): - user_name: str - last_n: int - - -class UserChatMessage(BaseModel): - type: str - text: str - - -type GetUserChatHistoryOutput = list[UserChatMessage] - - -class GetChatStatusTodayInput(BaseModel): - user_name: str - - -class GetChatStatusTodayOutput(BaseModel): - user_name: str - chat_cnt: int - - -# === mongodb documents start === - - -from pydantic import (AfterValidator, GetPydanticSchema, - PlainSerializer, WithJsonSchema) +from pydantic import AfterValidator, PlainSerializer, WithJsonSchema, GetPydanticSchema, BaseModel, Field, ConfigDict from pydantic_mongo import ObjectIdField as _objectIdField +# ObjectIdField = Annotated[str, BeforeValidator(ObjectId)] + + ObjectIdField = Annotated[ _objectIdField, AfterValidator(lambda id: _objectIdField(id)), @@ -58,9 +22,6 @@ class Test(BaseModel): id: ObjectIdField = Field(default_factory=ObjectIdField, alias='_id', title='_id') -# ObjectIdField = Annotated[str, BeforeValidator(ObjectId)] - - class MessageRoleType(str, Enum): User = "user" Ai = "ai" @@ -117,7 +78,6 @@ class Message(BaseModel): }, ) - # class Conversation(BaseModel): # id: ObjectIdField = Field(default_factory=ObjectIdField, alias="_id", title='_id') # user_id: ObjectIdField = Field() @@ -141,22 +101,3 @@ class Message(BaseModel): # } # }, # ) - - -# === mongodb documents end === - - -# class UserConversationMessages(BaseModel): -# user_id: ObjectIdField = Field() -# user_name: str = Field() -# conversation_id: ObjectIdField = Field(default=None) -# title: str = Field() -# created_at: datetime.datetime = Field(default_factory=datetime.datetime.now) -# created_by: ObjectIdField = Field() -# updated_at: Optional[datetime.datetime] = Field(default=None) -# updated_by: Optional[ObjectIdField] = Field(default=None) -# messages: list[Message] - - -class Context(BaseModel): - user: User diff --git a/simplylab/providers/chat.py b/simplylab/providers/chat.py index 3ea77a9..6ab1067 100644 --- a/simplylab/providers/chat.py +++ b/simplylab/providers/chat.py @@ -5,7 +5,8 @@ import pymongo from loguru import logger from simplylab.database import Database -from simplylab.entity import ObjectIdField, Message +from simplylab.model.table import ObjectIdField +from simplylab.model.table import Message class ChatProvider: diff --git a/simplylab/providers/user.py b/simplylab/providers/user.py index e230219..c55a727 100644 --- a/simplylab/providers/user.py +++ b/simplylab/providers/user.py @@ -4,7 +4,7 @@ from typing import Optional from loguru import logger from simplylab.database import Database -from simplylab.entity import User +from simplylab.model.table import User class UserProvider: diff --git a/simplylab/services/__init__.py b/simplylab/services/__init__.py index e8ce1ec..7512483 100644 --- a/simplylab/services/__init__.py +++ b/simplylab/services/__init__.py @@ -1,6 +1,6 @@ from typing import Any -from simplylab.entity import Context +from simplylab.model.entity import Context from simplylab.providers import Providers from simplylab.services.chat import ChatService diff --git a/simplylab/services/chat.py b/simplylab/services/chat.py index 3838c65..8d1f4da 100644 --- a/simplylab/services/chat.py +++ b/simplylab/services/chat.py @@ -2,8 +2,12 @@ from typing import Any from loguru import logger -from simplylab.entity import GetAiChatResponseInput, GetUserChatHistoryInput, GetChatStatusTodayInput, UserChatMessage, \ - GetChatStatusTodayOutput, GetAiChatResponseOutput, GetUserChatHistoryOutput, Context, Message, MessageRoleType +from simplylab.model.res import UserChatMessage, \ + GetUserChatHistoryOutput +from simplylab.model.entity import Context +from simplylab.model.res import GetAiChatResponseOutput, GetChatStatusTodayOutput +from simplylab.model.req import GetAiChatResponseInput, GetUserChatHistoryInput, GetChatStatusTodayInput +from simplylab.model.table import MessageRoleType, Message from simplylab.error import MessageLimitedInDailyError, MessageLimitedIn30SecondsError from simplylab.providers import Providers