You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

48 lines
1.9 KiB
Python

from __future__ import annotations
import sqlalchemy
from api.model.base import BaseHelper, BaseModel
import sqlalchemy
from sqlalchemy import select, update
from sqlalchemy.orm import mapped_column, relationship, Mapped
class BotPersonaModel(BaseModel):
__tablename__ = "chat_complete_bot_persona"
id: Mapped[int] = mapped_column(sqlalchemy.Integer, primary_key=True, autoincrement=True)
bot_id: Mapped[str] = mapped_column(sqlalchemy.String(60), index=True)
bot_name: Mapped[str] = mapped_column(sqlalchemy.String(60), index=True)
bot_avatar: Mapped[str] = mapped_column(sqlalchemy.String, nullable=True)
bot_description: Mapped[str] = mapped_column(sqlalchemy.String, nullable=True)
system_prompt: Mapped[str] = mapped_column(sqlalchemy.String)
message_log: Mapped[list] = mapped_column(sqlalchemy.JSON)
default_question: Mapped[str] = mapped_column(sqlalchemy.String, nullable=True)
class BotPersonaHelper(BaseHelper):
async def add(self, obj: BotPersonaModel):
self.session.add(obj)
await self.session.commit()
await self.session.refresh(obj)
return obj
async def update(self, obj: BotPersonaModel):
obj = await self.session.merge(obj)
await self.session.commit()
return obj
async def get_list(self):
stmt = select(BotPersonaModel).with_only_columns([
BotPersonaModel.id,
BotPersonaModel.bot_name,
BotPersonaModel.bot_avatar,
BotPersonaModel.bot_description
])
return await self.session.scalars(stmt)
async def find_by_id(self, id: int):
stmt = select(BotPersonaModel).where(BotPersonaModel.id == id)
return await self.session.scalar(stmt)
async def find_by_bot_id(self, bot_id: str):
stmt = select(BotPersonaModel).where(BotPersonaModel.bot_id == bot_id)
return await self.session.scalar(stmt)