diff --git a/api/apps/canvas_app.py b/api/apps/canvas_app.py index 550ef04f1..22d99e3e8 100644 --- a/api/apps/canvas_app.py +++ b/api/apps/canvas_app.py @@ -24,7 +24,7 @@ from flask import request, Response from flask_login import login_required, current_user from agent.component import LLM -from api.db import FileType +from api.db import CanvasCategory, FileType from api.db.services.canvas_service import CanvasTemplateService, UserCanvasService, API4ConversationService from api.db.services.document_service import DocumentService from api.db.services.file_service import FileService @@ -45,14 +45,14 @@ from rag.utils.redis_conn import REDIS_CONN @manager.route('/templates', methods=['GET']) # noqa: F821 @login_required def templates(): - return get_json_result(data=[c.to_dict() for c in CanvasTemplateService.get_all()]) + return get_json_result(data=[c.to_dict() for c in CanvasTemplateService.query(canvas_category=CanvasCategory.Agent)]) @manager.route('/list', methods=['GET']) # noqa: F821 @login_required def canvas_list(): return get_json_result(data=sorted([c.to_dict() for c in \ - UserCanvasService.query(user_id=current_user.id)], key=lambda x: x["update_time"]*-1) + UserCanvasService.query(user_id=current_user.id, canvas_category=CanvasCategory.Agent)], key=lambda x: x["update_time"]*-1) ) @@ -79,7 +79,7 @@ def save(): req["dsl"] = json.loads(req["dsl"]) if "id" not in req: req["user_id"] = current_user.id - if UserCanvasService.query(user_id=current_user.id, title=req["title"].strip()): + if UserCanvasService.query(user_id=current_user.id, title=req["title"].strip(), canvas_category=CanvasCategory.Agent): return get_data_error_result(message=f"{req['title'].strip()} already exists.") req["id"] = get_uuid() if not UserCanvasService.save(**req): @@ -91,7 +91,7 @@ def save(): code=RetCode.OPERATING_ERROR) UserCanvasService.update_by_id(req["id"], req) # save version - UserCanvasVersionService.insert( user_canvas_id=req["id"], dsl=req["dsl"], title="{0}_{1}".format(req["title"], time.strftime("%Y_%m_%d_%H_%M_%S"))) + UserCanvasVersionService.insert(user_canvas_id=req["id"], dsl=req["dsl"], title="{0}_{1}".format(req["title"], time.strftime("%Y_%m_%d_%H_%M_%S"))) UserCanvasVersionService.delete_all_versions(req["id"]) return get_json_result(data=req) @@ -395,7 +395,7 @@ def list_canvas(): tenants = TenantService.get_joined_tenants_by_user_id(current_user.id) canvas, total = UserCanvasService.get_by_tenant_ids( [m["tenant_id"] for m in tenants], current_user.id, page_number, - items_per_page, orderby, desc, keywords) + items_per_page, orderby, desc, keywords, canvas_category=CanvasCategory.Agent) return get_json_result(data={"canvas": canvas, "total": total}) except Exception as e: return server_error_response(e) diff --git a/api/db/__init__.py b/api/db/__init__.py index a7fb046f8..186e9bcba 100644 --- a/api/db/__init__.py +++ b/api/db/__init__.py @@ -74,8 +74,10 @@ class TaskStatus(StrEnum): DONE = "3" FAIL = "4" + VALID_TASK_STATUS = {TaskStatus.UNSTART, TaskStatus.RUNNING, TaskStatus.CANCEL, TaskStatus.DONE, TaskStatus.FAIL} + class ParserType(StrEnum): PRESENTATION = "presentation" LAWS = "laws" @@ -105,10 +107,19 @@ class CanvasType(StrEnum): DocBot = "docbot" +class CanvasCategory(StrEnum): + Agent = "agent_canvas" + DataFlow = "dataflow_canvas" + +VALID_CAVAS_CATEGORIES = {CanvasCategory.Agent, CanvasCategory.DataFlow} + + class MCPServerType(StrEnum): SSE = "sse" STREAMABLE_HTTP = "streamable-http" + VALID_MCP_SERVER_TYPES = {MCPServerType.SSE, MCPServerType.STREAMABLE_HTTP} + KNOWLEDGEBASE_FOLDER_NAME=".knowledgebase" diff --git a/api/db/db_models.py b/api/db/db_models.py index bfae37838..cda279f22 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -245,22 +245,21 @@ class JsonSerializedField(SerializedField): class RetryingPooledMySQLDatabase(PooledMySQLDatabase): def __init__(self, *args, **kwargs): - self.max_retries = kwargs.pop('max_retries', 5) - self.retry_delay = kwargs.pop('retry_delay', 1) + self.max_retries = kwargs.pop("max_retries", 5) + self.retry_delay = kwargs.pop("retry_delay", 1) super().__init__(*args, **kwargs) def execute_sql(self, sql, params=None, commit=True): from peewee import OperationalError + for attempt in range(self.max_retries + 1): try: return super().execute_sql(sql, params, commit) except OperationalError as e: if e.args[0] in (2013, 2006) and attempt < self.max_retries: - logging.warning( - f"Lost connection (attempt {attempt+1}/{self.max_retries}): {e}" - ) + logging.warning(f"Lost connection (attempt {attempt + 1}/{self.max_retries}): {e}") self._handle_connection_loss() - time.sleep(self.retry_delay * (2 ** attempt)) + time.sleep(self.retry_delay * (2**attempt)) else: logging.error(f"DB execution failure: {e}") raise @@ -272,16 +271,15 @@ class RetryingPooledMySQLDatabase(PooledMySQLDatabase): def begin(self): from peewee import OperationalError + for attempt in range(self.max_retries + 1): try: return super().begin() except OperationalError as e: if e.args[0] in (2013, 2006) and attempt < self.max_retries: - logging.warning( - f"Lost connection during transaction (attempt {attempt+1}/{self.max_retries})" - ) + logging.warning(f"Lost connection during transaction (attempt {attempt + 1}/{self.max_retries})") self._handle_connection_loss() - time.sleep(self.retry_delay * (2 ** attempt)) + time.sleep(self.retry_delay * (2**attempt)) else: raise @@ -815,6 +813,7 @@ class UserCanvas(DataBaseModel): permission = CharField(max_length=16, null=False, help_text="me|team", default="me", index=True) description = TextField(null=True, help_text="Canvas description") canvas_type = CharField(max_length=32, null=True, help_text="Canvas type", index=True) + canvas_category = CharField(max_length=32, null=False, default="agent_canvas", help_text="Canvas category: agent_canvas|dataflow_canvas", index=True) dsl = JSONField(null=True, default={}) class Meta: @@ -827,6 +826,7 @@ class CanvasTemplate(DataBaseModel): title = JSONField(null=True, default=dict, help_text="Canvas title") description = JSONField(null=True, default=dict, help_text="Canvas description") canvas_type = CharField(max_length=32, null=True, help_text="Canvas type", index=True) + canvas_category = CharField(max_length=32, null=False, default="agent_canvas", help_text="Canvas category: agent_canvas|dataflow_canvas", index=True) dsl = JSONField(null=True, default={}) class Meta: @@ -1029,4 +1029,12 @@ def migrate_db(): migrate(migrator.alter_column_type("canvas_template", "description", JSONField(null=True, default=dict, help_text="Canvas description"))) except Exception: pass + try: + migrate(migrator.add_column("user_canvas", "canvas_category", CharField(max_length=32, null=False, default="agent_canvas", help_text="agent_canvas|dataflow_canvas", index=True))) + except Exception: + pass + try: + migrate(migrator.add_column("canvas_template", "canvas_category", CharField(max_length=32, null=False, default="agent_canvas", help_text="agent_canvas|dataflow_canvas", index=True))) + except Exception: + pass logging.disable(logging.NOTSET) diff --git a/api/db/services/canvas_service.py b/api/db/services/canvas_service.py index 562ce53f4..ddb00ac11 100644 --- a/api/db/services/canvas_service.py +++ b/api/db/services/canvas_service.py @@ -18,7 +18,7 @@ import logging import time from uuid import uuid4 from agent.canvas import Canvas -from api.db import TenantPermission +from api.db import CanvasCategory, TenantPermission from api.db.db_models import DB, CanvasTemplate, User, UserCanvas, API4Conversation from api.db.services.api_service import API4ConversationService from api.db.services.common_service import CommonService @@ -31,6 +31,12 @@ from peewee import fn class CanvasTemplateService(CommonService): model = CanvasTemplate +class DataFlowTemplateService(CommonService): + """ + Alias of CanvasTemplateService + """ + model = CanvasTemplate + class UserCanvasService(CommonService): model = UserCanvas @@ -38,13 +44,14 @@ class UserCanvasService(CommonService): @classmethod @DB.connection_context() def get_list(cls, tenant_id, - page_number, items_per_page, orderby, desc, id, title): + page_number, items_per_page, orderby, desc, id, title, canvas_category=CanvasCategory.Agent): agents = cls.model.select() if id: agents = agents.where(cls.model.id == id) if title: agents = agents.where(cls.model.title == title) agents = agents.where(cls.model.user_id == tenant_id) + agents = agents.where(cls.model.canvas_category == canvas_category) if desc: agents = agents.order_by(cls.model.getter_by(orderby).desc()) else: @@ -71,6 +78,7 @@ class UserCanvasService(CommonService): cls.model.create_time, cls.model.create_date, cls.model.update_date, + cls.model.canvas_category, User.nickname, User.avatar.alias('tenant_avatar'), ] @@ -87,7 +95,7 @@ class UserCanvasService(CommonService): @DB.connection_context() def get_by_tenant_ids(cls, joined_tenant_ids, user_id, page_number, items_per_page, - orderby, desc, keywords, + orderby, desc, keywords, canvas_category=CanvasCategory.Agent, ): fields = [ cls.model.id, @@ -98,7 +106,8 @@ class UserCanvasService(CommonService): cls.model.permission, User.nickname, User.avatar.alias('tenant_avatar'), - cls.model.update_time + cls.model.update_time, + cls.model.canvas_category, ] if keywords: agents = cls.model.select(*fields).join(User, on=(cls.model.user_id == User.id)).where( @@ -113,6 +122,7 @@ class UserCanvasService(CommonService): TenantPermission.TEAM.value)) | ( cls.model.user_id == user_id)) ) + agents = agents.where(cls.model.canvas_category == canvas_category) if desc: agents = agents.order_by(cls.model.getter_by(orderby).desc()) else: