mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2025-09-15 09:39:42 +00:00
66 lines
2.7 KiB
Python
66 lines
2.7 KiB
Python
from time import time
|
|
from typing import Optional,List
|
|
from uuid import uuid4
|
|
|
|
from ktransformers.server.models.assistants.assistants import Assistant
|
|
from ktransformers.server.schemas.assistants.assistants import AssistantCreate,AssistantObject,AssistantModify
|
|
from ktransformers.server.utils.sql_utils import SQLUtil
|
|
from ktransformers.server.config.log import logger
|
|
from ktransformers.server.schemas.base import Order
|
|
|
|
|
|
class AssistantDatabaseManager:
|
|
def __init__(self) -> None:
|
|
self.sql_util = SQLUtil()
|
|
|
|
def create_assistant_object(self, assistant: AssistantCreate) -> AssistantObject:
|
|
assistant = AssistantObject(
|
|
**assistant.model_dump(mode='json'),
|
|
id=str(uuid4()),
|
|
object='assistant',
|
|
created_at=int(time()),
|
|
)
|
|
return assistant
|
|
|
|
def db_count_assistants(self) -> int:
|
|
with self.sql_util.get_db() as db:
|
|
return db.query(Assistant).count()
|
|
|
|
def db_create_assistant(self, assistant: AssistantCreate):
|
|
ass_obj = self.create_assistant_object(assistant)
|
|
ass_obj.sync_db()
|
|
return ass_obj
|
|
|
|
def db_list_assistants(self, limit: Optional[int], order: Order) -> List[AssistantObject]:
|
|
with self.sql_util.get_db() as db:
|
|
query = db.query(Assistant).order_by(
|
|
order.to_sqlalchemy_order()(Assistant.created_at))
|
|
if limit is not None:
|
|
db_assistants = query.limit(limit)
|
|
else:
|
|
db_assistants = query.all()
|
|
return [AssistantObject.model_validate(a.__dict__) for a in db_assistants]
|
|
|
|
def db_get_assistant_by_id(self, assistant_id: str) -> Optional[AssistantObject]:
|
|
with self.sql_util.get_db() as db:
|
|
db_assistant = db.query(Assistant).filter(
|
|
Assistant.id == assistant_id).first()
|
|
if db_assistant is None:
|
|
logger.debug(f"no assistant with id {str}")
|
|
return None
|
|
return AssistantObject.model_validate(db_assistant.__dict__)
|
|
|
|
def db_update_assistant_by_id(self, assistant_id: str, assistant: AssistantModify):
|
|
with self.sql_util.get_db() as db:
|
|
db_assistant = db.query(Assistant).filter(
|
|
Assistant.id == assistant_id).first()
|
|
self.sql_util.db_update_commit_refresh(db, db_assistant, assistant)
|
|
return AssistantObject.model_validate(db_assistant.__dict__)
|
|
|
|
def db_delete_assistant_by_id(self, assistant_id: str):
|
|
with self.sql_util.get_db() as db:
|
|
db_assistant = db.query(Assistant).filter(
|
|
Assistant.id == assistant_id).first()
|
|
db.delete(db_assistant)
|
|
db.commit()
|
|
|