ragflow/api/db/__init__.py

#
#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

from enum import IntEnum
from enum import StrEnum

from common.constants import PipelineTaskType


class UserTenantRole(StrEnum):
    OWNER = "owner"
    ADMIN = "admin"
    NORMAL = "normal"
    INVITE = "invite"


class TenantPermission(StrEnum):
    ME = "me"
    TEAM = "team"


class SerializedType(IntEnum):
    PICKLE = 1
    JSON = 2


class FileType(StrEnum):
    PDF = "pdf"
    DOC = "doc"
    VISUAL = "visual"
    AURAL = "aural"
    VIRTUAL = "virtual"
    FOLDER = "folder"
    OTHER = "other"


VALID_FILE_TYPES = {FileType.PDF, FileType.DOC, FileType.VISUAL, FileType.AURAL, FileType.VIRTUAL, FileType.FOLDER, FileType.OTHER}


class InputType(StrEnum):
    LOAD_STATE = "load_state"  # e.g. loading a current full state or a save state, such as from a file
    POLL = "poll"  # e.g. calling an API to get all documents in the last hour
    EVENT = "event"  # e.g. registered an endpoint as a listener, and processing connector events
    SLIM_RETRIEVAL = "slim_retrieval"


class CanvasCategory(StrEnum):
    Agent = "agent_canvas"
    DataFlow = "dataflow_canvas"


VALID_PIPELINE_TASK_TYPES = {
    PipelineTaskType.PARSE,
    PipelineTaskType.DOWNLOAD,
    PipelineTaskType.RAPTOR,
    PipelineTaskType.GRAPH_RAG,
    PipelineTaskType.MINDMAP,
    PipelineTaskType.ARTIFACT,
    PipelineTaskType.SKILL,
}


# KB-level fan-out task types: their Task row uses GRAPH_RAPTOR_FAKE_DOC_ID as a
# sentinel doc_id, and ``task_executor.collect_task`` substitutes the first real
# doc_id from ``msg["doc_ids"]`` before re-running ``TaskService.get_task`` so
# the join through Document → Knowledgebase → Tenant resolves and tenant_id /
# kb_id / language are hydrated onto the task dict. Add new fan-out task types
# here or TaskContext will raise "Task must contain 'tenant_id'".
PIPELINE_SPECIAL_PROGRESS_FREEZE_TASK_TYPES = {
    PipelineTaskType.RAPTOR.lower(),
    PipelineTaskType.GRAPH_RAG.lower(),
    PipelineTaskType.MINDMAP.lower(),
    PipelineTaskType.ARTIFACT.lower(),
    PipelineTaskType.SKILL.lower(),
}


KNOWLEDGEBASE_FOLDER_NAME = ".knowledgebase"
SKILLS_FOLDER_NAME = "skills"