zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-11 06:05:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

Feat: Add knowledge compilation workflows (#16515)

Browse Source 
## Summary
- Add knowledge compilation template APIs, services, and builtin
template seed data
- Add advanced knowledge compile structure/artifact/RAPTOR workflow
support
- Update parsing, dataset/document APIs, and supporting services for
compilation workflows
This commit is contained in:
Kevin Hu
2026-07-02 23:22:07 +08:00
committed by GitHub
parent 7d64a78f83
commit 62f94cd59b
57 changed files with 14587 additions and 3094 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
api/db/__init__.py
View File

@@ -21,15 +21,15 @@ from common.constants import PipelineTaskType
class UserTenantRole(StrEnum):
OWNER = 'owner'
ADMIN = 'admin'
NORMAL = 'normal'
INVITE = 'invite'
OWNER = "owner"
ADMIN = "admin"
NORMAL = "normal"
INVITE = "invite"
class TenantPermission(StrEnum):
ME = 'me'
TEAM = 'team'
ME = "me"
TEAM = "team"
class SerializedType(IntEnum):
@@ -38,14 +38,15 @@ class SerializedType(IntEnum):
class FileType(StrEnum):
PDF = 'pdf'
DOC = 'doc'
VISUAL = 'visual'
AURAL = 'aural'
VIRTUAL = 'virtual'
FOLDER = 'folder'
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}
@@ -61,11 +62,31 @@ class CanvasCategory(StrEnum):
DataFlow = "dataflow_canvas"
VALID_PIPELINE_TASK_TYPES = {PipelineTaskType.PARSE, PipelineTaskType.DOWNLOAD, PipelineTaskType.RAPTOR, PipelineTaskType.GRAPH_RAG, PipelineTaskType.MINDMAP}
VALID_PIPELINE_TASK_TYPES = {
PipelineTaskType.PARSE,
PipelineTaskType.DOWNLOAD,
PipelineTaskType.RAPTOR,
PipelineTaskType.GRAPH_RAG,
PipelineTaskType.MINDMAP,
PipelineTaskType.ARTIFACT,
PipelineTaskType.SKILL,
}
PIPELINE_SPECIAL_PROGRESS_FREEZE_TASK_TYPES = {PipelineTaskType.RAPTOR.lower(), PipelineTaskType.GRAPH_RAG.lower(), PipelineTaskType.MINDMAP.lower()}
# 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"
KNOWLEDGEBASE_FOLDER_NAME = ".knowledgebase"
SKILLS_FOLDER_NAME = "skills"

195
api/db/db_models.py
View File

@@ -269,19 +269,13 @@ class RetryingPooledMySQLDatabase(PooledMySQLDatabase):
return super().execute_sql(sql, params, commit)
except (OperationalError, InterfaceError) as e:
error_codes = [2013, 2006]
error_messages = ['', 'Lost connection']
should_retry = (
(hasattr(e, 'args') and e.args and e.args[0] in error_codes) or
(str(e) in error_messages) or
(hasattr(e, '__class__') and e.__class__.__name__ == 'InterfaceError')
)
error_messages = ["", "Lost connection"]
should_retry = (hasattr(e, "args") and e.args and e.args[0] in error_codes) or (str(e) in error_messages) or (hasattr(e, "__class__") and e.__class__.__name__ == "InterfaceError")
if should_retry and attempt < self.max_retries:
logging.warning(
f"Database connection issue (attempt {attempt+1}/{self.max_retries}): {e}"
)
logging.warning(f"Database connection issue (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
@@ -311,20 +305,14 @@ class RetryingPooledMySQLDatabase(PooledMySQLDatabase):
return super().begin()
except (OperationalError, InterfaceError) as e:
error_codes = [2013, 2006]
error_messages = ['', 'Lost connection']
error_messages = ["", "Lost connection"]
should_retry = (
(hasattr(e, 'args') and e.args and e.args[0] in error_codes) or
(str(e) in error_messages) or
(hasattr(e, '__class__') and e.__class__.__name__ == 'InterfaceError')
)
should_retry = (hasattr(e, "args") and e.args and e.args[0] in error_codes) or (str(e) in error_messages) or (hasattr(e, "__class__") and e.__class__.__name__ == "InterfaceError")
if should_retry 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
return None
@@ -348,17 +336,14 @@ class RetryingPooledPostgresqlDatabase(PooledPostgresqlDatabase):
# 08006: connection_failure
# 08003: connection_does_not_exist
# 08000: connection_exception
error_messages = ['connection', 'server closed', 'connection refused',
'no connection to the server', 'terminating connection']
error_messages = ["connection", "server closed", "connection refused", "no connection to the server", "terminating connection"]
should_retry = any(msg in str(e).lower() for msg in error_messages)
if should_retry and attempt < self.max_retries:
logging.warning(
f"PostgreSQL connection issue (attempt {attempt+1}/{self.max_retries}): {e}"
)
logging.warning(f"PostgreSQL connection issue (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"PostgreSQL execution failure: {e}")
raise
@@ -385,17 +370,14 @@ class RetryingPooledPostgresqlDatabase(PooledPostgresqlDatabase):
try:
return super().begin()
except (OperationalError, InterfaceError) as e:
error_messages = ['connection', 'server closed', 'connection refused',
'no connection to the server', 'terminating connection']
error_messages = ["connection", "server closed", "connection refused", "no connection to the server", "terminating connection"]
should_retry = any(msg in str(e).lower() for msg in error_messages)
if should_retry and attempt < self.max_retries:
logging.warning(
f"PostgreSQL connection lost during transaction (attempt {attempt+1}/{self.max_retries})"
)
logging.warning(f"PostgreSQL connection lost 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
return None
@@ -407,6 +389,7 @@ class RetryingPooledOceanBaseDatabase(PooledMySQLDatabase):
OceanBase is compatible with MySQL protocol, so we inherit from PooledMySQLDatabase.
This class provides connection pooling and automatic retry for connection issues.
"""
def __init__(self, *args, **kwargs):
self.max_retries = kwargs.pop("max_retries", 5)
self.retry_delay = kwargs.pop("retry_delay", 1)
@@ -421,20 +404,18 @@ class RetryingPooledOceanBaseDatabase(PooledMySQLDatabase):
# 2013: Lost connection to MySQL server during query
# 2006: MySQL server has gone away
error_codes = [2013, 2006]
error_messages = ['', 'Lost connection', 'gone away']
error_messages = ["", "Lost connection", "gone away"]
should_retry = (
(hasattr(e, 'args') and e.args and e.args[0] in error_codes) or
any(msg in str(e).lower() for msg in error_messages) or
(hasattr(e, '__class__') and e.__class__.__name__ == 'InterfaceError')
(hasattr(e, "args") and e.args and e.args[0] in error_codes)
or any(msg in str(e).lower() for msg in error_messages)
or (hasattr(e, "__class__") and e.__class__.__name__ == "InterfaceError")
)
if should_retry and attempt < self.max_retries:
logging.warning(
f"OceanBase connection issue (attempt {attempt+1}/{self.max_retries}): {e}"
)
logging.warning(f"OceanBase connection issue (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"OceanBase execution failure: {e}")
raise
@@ -462,20 +443,14 @@ class RetryingPooledOceanBaseDatabase(PooledMySQLDatabase):
return super().begin()
except (OperationalError, InterfaceError) as e:
error_codes = [2013, 2006]
error_messages = ['', 'Lost connection']
error_messages = ["", "Lost connection"]
should_retry = (
(hasattr(e, 'args') and e.args and e.args[0] in error_codes) or
(str(e) in error_messages) or
(hasattr(e, '__class__') and e.__class__.__name__ == 'InterfaceError')
)
should_retry = (hasattr(e, "args") and e.args and e.args[0] in error_codes) or (str(e) in error_messages) or (hasattr(e, "__class__") and e.__class__.__name__ == "InterfaceError")
if should_retry 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
return None
@@ -500,13 +475,11 @@ class BaseDataBase:
db_name = database_config.pop("name")
pool_config = {
'max_retries': 5,
'retry_delay': 1,
"max_retries": 5,
"retry_delay": 1,
}
database_config.update(pool_config)
self.database_connection = PooledDatabase[settings.DATABASE_TYPE.upper()].value(
db_name, **database_config
)
self.database_connection = PooledDatabase[settings.DATABASE_TYPE.upper()].value(db_name, **database_config)
# self.database_connection = PooledDatabase[settings.DATABASE_TYPE.upper()].value(db_name, **database_config)
logging.info("init database on cluster mode successfully")
@@ -846,9 +819,7 @@ class TenantLLM(DataBaseModel):
class Meta:
db_table = "tenant_llm"
indexes = (
(("tenant_id", "llm_factory", "llm_name"), True),
)
indexes = ((("tenant_id", "llm_factory", "llm_name"), True),)
class TenantLangfuse(DataBaseModel):
@@ -892,6 +863,10 @@ class Knowledgebase(DataBaseModel):
raptor_task_finish_at = DateTimeField(null=True)
mindmap_task_id = CharField(max_length=32, null=True, help_text="Mindmap task ID", index=True)
mindmap_task_finish_at = DateTimeField(null=True)
artifact_task_id = CharField(max_length=32, null=True, help_text="Artifact compilation task ID", index=True)
artifact_task_finish_at = DateTimeField(null=True)
skill_task_id = CharField(max_length=32, null=True, help_text="Skill generation task ID", index=True)
skill_task_finish_at = DateTimeField(null=True)
status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True)
@@ -964,6 +939,13 @@ class FileCommit(DataBaseModel):
author_id = CharField(max_length=32, null=False, help_text="user who created the commit", index=True)
file_count = IntegerField(default=0, help_text="number of files in this commit")
tree_state = LongTextField(null=True, help_text="JSON snapshot of the full folder tree at this commit")
# ---- Artifact-commit extension ----
# Populated only for commits recorded via
# ``FileCommitService.record_page_edit`` (i.e. artifact-page saves).
# For workspace file commits both fields stay null and the ``message``
# column carries the commit body.
title = CharField(max_length=255, null=True, help_text="commit title (artifact-page edits)")
comments = TextField(null=True, help_text="commit body/description (artifact-page edits)")
class Meta:
db_table = "file_commit"
@@ -980,6 +962,14 @@ class FileCommitItem(DataBaseModel):
new_location = CharField(max_length=255, null=True, help_text="new storage location")
old_name = CharField(max_length=255, null=True, help_text="old file name (for rename)")
new_name = CharField(max_length=255, null=True, help_text="new file name (for rename)")
# ---- Artifact-commit extension ----
# Populated only for artifact-page saves recorded via
# ``FileCommitService.record_page_edit``.
diff = LongTextField(null=True, help_text="pre-computed unified diff (artifact-page edits)")
content_after_storage = CharField(max_length=16, null=True, help_text="'minio' | 'es' — where the post-save blob lives", index=True)
content_after_location = CharField(max_length=512, null=True, help_text="storage key/id for the post-save blob")
slug_kwd = CharField(max_length=512, null=True, help_text="artifact page slug (<page_type>/<name>)", index=True)
page_type_kwd = CharField(max_length=32, null=True, help_text="artifact page type", index=True)
class Meta:
db_table = "file_commit_item"
@@ -988,6 +978,13 @@ class FileCommitItem(DataBaseModel):
)
# ``ArtifactCommit`` retired — artifact page history is now stored under
# ``FileCommit`` + ``FileCommitItem`` via ``FileCommitService.record_page_edit``
# (see the artifact-commit extension columns on those models above).
# Pre-existing ``artifact_commit`` rows are intentionally left in place;
# no code path reads them.
class Task(DataBaseModel):
id = CharField(max_length=32, primary_key=True)
doc_id = CharField(max_length=32, null=False, index=True)
@@ -1146,6 +1143,35 @@ class MCPServer(DataBaseModel):
db_table = "mcp_server"
class CompilationTemplate(DataBaseModel):
id = CharField(max_length=32, primary_key=True)
tenant_id = CharField(max_length=32, null=True, index=True)
group_id = CharField(max_length=32, null=True, index=True)
name = CharField(max_length=128, null=False, index=True)
description = TextField(null=True, default="")
kind = CharField(max_length=64, null=False, index=True)
config = JSONField(null=False, default={})
is_builtin = BooleanField(null=False, default=False, index=True)
status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True)
class Meta:
db_table = "compilation_template"
indexes = ((("tenant_id", "name", "is_builtin", "status"), True),)
class CompilationTemplateGroup(DataBaseModel):
id = CharField(max_length=32, primary_key=True)
tenant_id = CharField(max_length=32, null=False, index=True)
name = CharField(max_length=128, null=False, index=True)
description = TextField(null=True, default="")
scope = CharField(max_length=16, null=False, index=True, help_text="file | dataset")
status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True)
class Meta:
db_table = "compilation_template_group"
indexes = ((("tenant_id", "name", "status"), True),)
class Search(DataBaseModel):
id = CharField(max_length=32, primary_key=True)
avatar = TextField(null=True, help_text="avatar base64 string")
@@ -1264,9 +1290,9 @@ class ChatChannel(DataBaseModel):
class DateTimeTzField(CharField):
field_type = 'VARCHAR'
field_type = "VARCHAR"
def db_value(self, value: datetime|None) -> str|None:
def db_value(self, value: datetime | None) -> str | None:
if value is not None:
if value.tzinfo is not None:
return value.isoformat()
@@ -1274,11 +1300,12 @@ class DateTimeTzField(CharField):
return value.replace(tzinfo=timezone.utc).isoformat()
return value
def python_value(self, value: str|None) -> datetime|None:
def python_value(self, value: str | None) -> datetime | None:
if value is not None:
dt = datetime.fromisoformat(value)
if dt.tzinfo is None:
import pytz
return dt.replace(tzinfo=pytz.UTC)
return dt
return value
@@ -1307,6 +1334,7 @@ class SyncLogs(DataBaseModel):
class EvaluationDataset(DataBaseModel):
"""Ground truth dataset for RAG evaluation"""
id = CharField(max_length=32, primary_key=True)
tenant_id = CharField(max_length=32, null=False, index=True, help_text="tenant ID")
name = CharField(max_length=255, null=False, index=True, help_text="dataset name")
@@ -1323,6 +1351,7 @@ class EvaluationDataset(DataBaseModel):
class EvaluationCase(DataBaseModel):
"""Individual test case in an evaluation dataset"""
id = CharField(max_length=32, primary_key=True)
dataset_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_datasets")
question = TextField(null=False, help_text="test question")
@@ -1338,6 +1367,7 @@ class EvaluationCase(DataBaseModel):
class EvaluationRun(DataBaseModel):
"""A single evaluation run"""
id = CharField(max_length=32, primary_key=True)
dataset_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_datasets")
dialog_id = CharField(max_length=32, null=False, index=True, help_text="dialog configuration being evaluated")
@@ -1355,6 +1385,7 @@ class EvaluationRun(DataBaseModel):
class EvaluationResult(DataBaseModel):
"""Result for a single test case in an evaluation run"""
id = CharField(max_length=32, primary_key=True)
run_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_runs")
case_id = CharField(max_length=32, null=False, index=True, help_text="FK to evaluation_cases")
@@ -1375,7 +1406,7 @@ class Memory(DataBaseModel):
avatar = TextField(null=True, help_text="avatar base64 string")
tenant_id = CharField(max_length=32, null=False, index=True)
memory_type = IntegerField(null=False, default=1, index=True, help_text="Bit flags (LSB->MSB): 1=raw, 2=semantic, 4=episodic, 8=procedural. E.g., 5 enables raw + episodic.")
storage_type = CharField(max_length=32, default='table', null=False, index=True, help_text="table|graph")
storage_type = CharField(max_length=32, default="table", null=False, index=True, help_text="table|graph")
embd_id = CharField(max_length=128, null=False, index=False, help_text="embedding model ID")
tenant_embd_id = IntegerField(null=True, help_text="id in tenant_llm", index=True)
llm_id = CharField(max_length=128, null=False, index=False, help_text="chat model ID")
@@ -1391,14 +1422,17 @@ class Memory(DataBaseModel):
class Meta:
db_table = "memory"
class SystemSettings(DataBaseModel):
name = CharField(max_length=128, primary_key=True)
source = CharField(max_length=32, null=False, index=False)
data_type = CharField(max_length=32, null=False, index=False)
value = TextField(null=False, help_text="Configuration value (JSON, string, etc.)")
class Meta:
db_table = "system_settings"
class TenantModelProvider(DataBaseModel):
id = CharField(max_length=32, primary_key=True)
provider_name = CharField(max_length=128, null=False, index=False, help_text="LLM provider name")
@@ -1406,9 +1440,8 @@ class TenantModelProvider(DataBaseModel):
class Meta:
db_table = "tenant_model_provider"
indexes = (
(("tenant_id", "provider_name"), True),
)
indexes = ((("tenant_id", "provider_name"), True),)
class TenantModelInstance(DataBaseModel):
id = CharField(max_length=32, primary_key=True)
@@ -1444,6 +1477,7 @@ class TenantModelGroup(DataBaseModel):
class Meta:
db_table = "tenant_model_group"
class TenantModelGroupMapping(DataBaseModel):
group_id = CharField(max_length=32, null=False, index=True, help_text="Group ID")
provider_id = CharField(max_length=32, null=False, index=False)
@@ -1462,12 +1496,9 @@ def alter_db_add_column(migrator, table_name, column_name, column_type):
migrate(migrator.add_column(table_name, column_name, column_type))
except OperationalError as ex:
error_codes = [1060]
error_messages = ['Duplicate column name']
error_messages = ["Duplicate column name"]
should_skip_error = (
(hasattr(ex, 'args') and ex.args and ex.args[0] in error_codes) or
(str(ex) in error_messages)
)
should_skip_error = (hasattr(ex, "args") and ex.args and ex.args[0] in error_codes) or (str(ex) in error_messages)
if not should_skip_error:
logging.critical(f"Failed to add {settings.DATABASE_TYPE.upper()}.{table_name} column {column_name}, operation error: {ex}")
@@ -1476,6 +1507,7 @@ def alter_db_add_column(migrator, table_name, column_name, column_type):
logging.critical(f"Failed to add {settings.DATABASE_TYPE.upper()}.{table_name} column {column_name}, error: {ex}")
pass
def alter_db_column_type(migrator, table_name, column_name, new_column_type):
try:
migrate(migrator.alter_column_type(table_name, column_name, new_column_type))
@@ -1483,6 +1515,7 @@ def alter_db_column_type(migrator, table_name, column_name, new_column_type):
logging.critical(f"Failed to alter {settings.DATABASE_TYPE.upper()}.{table_name} column {column_name} type, error: {ex}")
pass
def alter_db_rename_column(migrator, table_name, old_column_name, new_column_name):
try:
migrate(migrator.rename_column(table_name, old_column_name, new_column_name))
@@ -1491,6 +1524,7 @@ def alter_db_rename_column(migrator, table_name, old_column_name, new_column_nam
# logging.critical(f"Failed to rename {settings.DATABASE_TYPE.upper()}.{table_name} column {old_column_name} to {new_column_name}, error: {ex}")
pass
def migrate_add_unique_email(migrator):
"""Deduplicates user emails and add UNIQUE constraint to email column (idempotent)"""
# step 0: check existing index state on user.email and prepare for unique constraint
@@ -1535,13 +1569,7 @@ def migrate_add_unique_email(migrator):
duplicates = User.select(User.email).group_by(User.email).having(fn.COUNT(User.id) > 1).tuples()
for (dup_email,) in duplicates:
# Keep the superuser row, or the oldest row if there is no superuser
rows = list(
User
.select(User.id)
.where(User.email == dup_email)
.order_by(User.is_superuser.desc(), User.create_time.asc())
.tuples()
)
rows = list(User.select(User.id).where(User.email == dup_email).order_by(User.is_superuser.desc(), User.create_time.asc()).tuples())
for (uid,) in rows[1:]:
new_email = f"{dup_email}_DUPLICATE_{uid[:8]}"
User.update(email=new_email).where(User.id == uid).execute()
@@ -1564,7 +1592,6 @@ def migrate_add_unique_email(migrator):
logging.critical("Failed to add UNIQUE constraint on user.email: %s", ex)
def update_tenant_llm_to_id_primary_key():
"""Add ID and set to primary key step by step."""
if settings.DATABASE_TYPE.upper() == "POSTGRES":
@@ -1749,6 +1776,10 @@ def migrate_db():
alter_db_add_column(migrator, "knowledgebase", "raptor_task_finish_at", CharField(null=True))
alter_db_add_column(migrator, "knowledgebase", "mindmap_task_id", CharField(max_length=32, null=True, help_text="Mindmap task ID", index=True))
alter_db_add_column(migrator, "knowledgebase", "mindmap_task_finish_at", CharField(null=True))
alter_db_add_column(migrator, "knowledgebase", "artifact_task_id", CharField(max_length=32, null=True, help_text="Artifact compilation task ID", index=True))
alter_db_add_column(migrator, "knowledgebase", "artifact_task_finish_at", DateTimeField(null=True))
alter_db_add_column(migrator, "knowledgebase", "skill_task_id", CharField(max_length=32, null=True, help_text="Skill generation task ID", index=True))
alter_db_add_column(migrator, "knowledgebase", "skill_task_finish_at", DateTimeField(null=True))
alter_db_column_type(migrator, "tenant_llm", "api_key", TextField(null=True, help_text="API KEY"))
alter_db_add_column(migrator, "tenant_llm", "status", CharField(max_length=1, null=False, help_text="is it validate(0: wasted, 1: validate)", default="1", index=True))
alter_db_add_column(migrator, "connector2kb", "auto_parse", CharField(max_length=1, null=False, default="1", index=False))
@@ -1779,6 +1810,14 @@ def migrate_db():
alter_db_add_column(migrator, "tenant", "ocr_id", CharField(max_length=128, null=True, help_text="default ocr model ID", index=True))
alter_db_column_type(migrator, "chat_channel", "status", IntegerField(default=1, index=True))
alter_db_rename_column(migrator, "chat_channel", "dialog_id", "chat_id")
# ---- FileCommit / FileCommitItem: artifact-page commit extension ----
alter_db_add_column(migrator, "file_commit", "title", CharField(max_length=255, null=True))
alter_db_add_column(migrator, "file_commit", "comments", TextField(null=True))
alter_db_add_column(migrator, "file_commit_item", "diff", LongTextField(null=True))
alter_db_add_column(migrator, "file_commit_item", "content_after_storage", CharField(max_length=16, null=True, index=True))
alter_db_add_column(migrator, "file_commit_item", "content_after_location", CharField(max_length=512, null=True))
alter_db_add_column(migrator, "file_commit_item", "slug_kwd", CharField(max_length=512, null=True, index=True))
alter_db_add_column(migrator, "file_commit_item", "page_type_kwd", CharField(max_length=32, null=True, index=True))
# Drop both the explicit "idx_*" name from later migrations AND the
# Peewee-auto-derived "<table-as-classname>_<col1>_<col2>" name from the
# original TenantModelInstance definition (commit dc4b82523). Databases

24
api/db/init_data.py
View File

@@ -21,11 +21,11 @@ import time
import uuid
from peewee import IntegrityError
from api.db import UserTenantRole
from api.db.db_models import init_database_tables as init_web_db
from api.db.services import UserService
from api.db.services.canvas_service import CanvasTemplateService
from api.db.services.compilation_template_service import CompilationTemplateService
from api.db.services.document_service import DocumentService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.llm_service import LLMBundle
@@ -43,6 +43,7 @@ DEFAULT_SUPERUSER_NICKNAME = os.getenv("DEFAULT_SUPERUSER_NICKNAME", "admin")
DEFAULT_SUPERUSER_EMAIL = os.getenv("DEFAULT_SUPERUSER_EMAIL", "admin@ragflow.io")
DEFAULT_SUPERUSER_PASSWORD = os.getenv("DEFAULT_SUPERUSER_PASSWORD", "admin")
def init_superuser(nickname=DEFAULT_SUPERUSER_NICKNAME, email=DEFAULT_SUPERUSER_EMAIL, password=DEFAULT_SUPERUSER_PASSWORD, role=UserTenantRole.OWNER):
if UserService.query(email=email):
logging.info("User with email %s already exists, skipping initialization.", email)
@@ -67,12 +68,7 @@ def init_superuser(nickname=DEFAULT_SUPERUSER_NICKNAME, email=DEFAULT_SUPERUSER_
"img2txt_id": settings.IMAGE2TEXT_MDL,
"rerank_id": settings.RERANK_MDL,
}
usr_tenant = {
"tenant_id": user_info["id"],
"user_id": user_info["id"],
"invited_by": user_info["id"],
"role": role
}
usr_tenant = {"tenant_id": user_info["id"], "user_id": user_info["id"], "invited_by": user_info["id"], "role": role}
try:
if not UserService.save(**user_info):
@@ -83,15 +79,14 @@ def init_superuser(nickname=DEFAULT_SUPERUSER_NICKNAME, email=DEFAULT_SUPERUSER_
return
TenantService.insert(**tenant)
UserTenantService.insert(**usr_tenant)
logging.info(
f"Super user initialized. email: {email},A default password has been set; changing the password after login is strongly recommended.")
logging.info(f"Super user initialized. email: {email},A default password has been set; changing the password after login is strongly recommended.")
if tenant["llm_id"]:
chat_model_config = get_tenant_default_model_by_type(tenant["id"], LLMType.CHAT)
chat_mdl = LLMBundle(tenant["id"], chat_model_config)
msg = asyncio.run(chat_mdl.async_chat(system="", history=[{"role": "user", "content": "Hello!"}], gen_conf={}))
if msg.find("ERROR: ") == 0:
logging.error("'{}' doesn't work. {}".format( tenant["llm_id"], msg))
logging.error("'{}' doesn't work. {}".format(tenant["llm_id"], msg))
if tenant["embd_id"]:
embd_model_config = get_tenant_default_model_by_type(tenant["id"], LLMType.EMBEDDING)
@@ -111,7 +106,6 @@ def update_document_number_in_init():
KnowledgebaseService.update_document_number_in_init(kb_id=kb_id, doc_num=doc_count.get(kb_id, 0))
def add_graph_templates():
dir = os.path.join(get_project_base_directory(), "agent", "templates")
CanvasTemplateService.filter_delete([1 == 1])
@@ -136,6 +130,10 @@ def add_graph_templates():
logging.exception("Add agent templates error for %s: %s", template_path, e)
def add_compilation_templates():
CompilationTemplateService.seed_builtins_from_files()
def init_web_data():
start_time = time.time()
@@ -147,11 +145,13 @@ def init_web_data():
# init_superuser()
add_graph_templates()
add_compilation_templates()
init_message_id_sequence()
init_memory_size_cache()
fix_missing_tokenized_memory()
logging.info("init web data success:{}".format(time.time() - start_time))
def init_table():
# init system_settings
with open(os.path.join(get_project_base_directory(), "conf", "system_settings.json"), "r") as f:
@@ -178,6 +178,6 @@ def init_table():
raise e
if __name__ == '__main__':
if __name__ == "__main__":
init_web_db()
init_web_data()

136
api/db/init_data/compilation_templates/artifacts.yaml Normal file
View File

@@ -0,0 +1,136 @@
kind: artifacts
display_name: Artifacts — Graph-based wiki
config:
kind: artifacts
example: |
- Each page must be a proper encyclopedic article, NOT a flat bullet list:
- 1. Opening paragraph (2-4 sentences defining what this is). No heading.
- 2. Sections with H2 headings, each starting with prose before sub-bullets.
- 3. Bold key terms on first use; link them with [[ ]] artifactlinks.
- 4. Examples or implications where the source provides them.
- 5. ## See also section at the end with artifactlinks to highly related pages(less than 12).\n
- Page structure could be as following: (Not provided)
entity:
description: >-
You are a robust graph entity extractor for knowledge graphs.
fields:
- type: person
description: A natural person (individual human).
rule: |
- Full name preferred (e.g., "Elon Musk", not "Musk" alone if ambiguous).
- Include titles only if integral to identity (e.g., "Dr. Smith").
- Max length: 60 characters.
- type: org
description: Organization, company, institution, agency, or any collective group.
rule: |
- Use the official name when possible (e.g., "United Nations").
- Abbreviations accepted if widely known (e.g., "UN", "NASA").
- Max length: 80 characters.
- type: product
description: Tangible or intangible product, service, software, or offering.
rule: |
- Include version numbers if relevant (e.g., "iPhone 14").
- Generic categories (e.g., "smartphone") only if no specific name is given.
- Max length: 100 characters.
- type: regulation
description: Law, policy, standard, guideline, or regulatory document.
rule: |
- Use official title or identifier (e.g., "GDPR", "OSHA standard 1910").
- Include jurisdiction if known (e.g., "EU GDPR").
- Max length: 120 characters.
- type: location
description: Geographic place (country, city, address, region, natural feature).
rule: |
- Hierarchical format allowed (e.g., "Paris, France").
- Avoid overly vague terms (e.g., "there") unless resolved.
- Max length: 80 characters.
- type: system
description: Technical system, platform, framework, or infrastructure.
rule: |
- Distinct from product: system implies integrated environment
(e.g., "Linux OS", "power grid").
- Use proper naming.
- Max length: 100 characters.
- type: equipment
description: Physical device, machinery, hardware, or tool.
rule: |
- Specific model preferred (e.g., "Boeing 737").
- Generic allowed only if precise type (e.g., "drill press").
- Max length: 80 characters.
- type: other
description: Entities that do not fit any above category.
rule: |
- Use sparingly; prefer mapping to a defined type when possible.
- Still provide a meaningful label.
- Max length: 80 characters.
relation:
description: >-
You are an expert in extracting semantic relations between entities.
fields:
- type: owns
description: Ownership or possession (legal or de facto).
rule: |
- Direction from owner to owned: (A owns B).
- Example: "Company A owns product B".
- type: part_of
description: Mereological relation — component to whole.
rule: |
- Direction from part to whole: (A part_of B).
- Example: "Engine part_of car".
- type: caused_by
description: Causal relation — event, action, or state leads to another.
rule: |
- Direction from effect to cause: (A caused_by B) meaning B causes A.
- Example: "Accident caused_by brake failure".
- type: regulates
description: Regulatory or governing relation (law/standard controls entity).
rule: |
- Direction from regulator to regulated: (A regulates B).
- Example: "GDPR regulates data processing".
- type: uses
description: Utilization — an entity employs or consumes another entity.
rule: |
- Direction from user to used: (A uses B).
- Example: "System uses equipment".
- type: located_in
description: Spatial containment — entity situated inside a location.
rule: |
- Direction from located entity to containing location: (A located_in B).
- Example: "Office located_in city".
- type: other
description: Any meaningful relation not covered by the above types.
rule: |
- Provide an explicit label in a "relation_label" field.
- Direction must be clear.
claim:
fields:
- statement: >-
A complete factual sentence stated in the source. Any sentence of the form
'X is Y', 'X has Y', 'X does Y', 'X was founded in Y', 'X is located in Y',
'X reported Y', etc. is a claim. Aim for at least 1-3 claims per entity per
chunk that mentions it.
subject: >-
Entity/concept this claim is about (must match one of the entity/concept
names extracted above).
concept:
fields:
- term: >-
Concept name OR a thematic section topic (prefer the source's heading
wording when coherent).
definition_excerpt: >-
Verbatim or near-verbatim defining phrase from the chunk.
global_rules: |
- Each relation links two entities (subject → object) with a predicate type.
- Format: {"subject_id": "<entity_id>", "predicate": "<type>",
"object_id": "<entity_id>", "chunk_id": "<chunk_ID>"}.
- Both subject and object must be previously extracted entities.
- If ambiguous direction, choose the most logical default
(e.g., "part_of" always from part to whole).
- When multiple relations appear in a chunk, list all in order of appearance.
- Keep language consistent; relation type is always English (the given type name).
- Every extracted entity must have exactly one type from the list.
- Entity label (the text representing the entity) is required, non-empty.
- Format: {"entity_id": "<unique_id>", "type": "<type>", "label": "<label>",
"chunk_id": "<chunk_ID>"}.
- If no entity in a chunk, output no entity for that chunk.
- Keep the chunks' original language (Chinese/English etc.) for entities and relations.

17
api/db/init_data/compilation_templates/empty.yaml Normal file
View File

@@ -0,0 +1,17 @@
kind: empty
display_name: Empty
config:
kind: empty
entity:
description: ''
fields:
- type: ''
description: ''
rule: ''
relation:
description: ''
fields:
- type: ''
description: ''
rule: ''
global_rules: ''

74
api/db/init_data/compilation_templates/knowledge_graph.yaml Normal file
View File

@@ -0,0 +1,74 @@
kind: knowledge_graph
display_name: Knowledge graph
config:
kind: knowledge_graph
entity:
description: >-
You are a robust graph entity extractor for knowledge graphs.
fields:
- type: person
description: A natural person (individual human).
rule: |
- Full name preferred (e.g., "Elon Musk", not "Musk" alone if ambiguous).
- Include titles only if integral to identity (e.g., "Dr. Smith").
- Max length: 60 characters.
- type: org
description: Organization, company, institution, agency, or any collective group.
rule: |
- Use the official name when possible (e.g., "United Nations").
- Abbreviations accepted if widely known (e.g., "UN", "NASA").
- Max length: 80 characters.
- type: product
description: Tangible or intangible product, service, software, or offering.
rule: |
- Include version numbers if relevant (e.g., "iPhone 14").
- Generic categories (e.g., "smartphone") only if no specific name is given.
- Max length: 100 characters.
- type: regulation
description: Law, policy, standard, guideline, or regulatory document.
rule: |
- Use official title or identifier (e.g., "GDPR", "OSHA standard 1910").
- Include jurisdiction if known (e.g., "EU GDPR").
- Max length: 120 characters.
- type: location
description: Geographic place (country, city, address, region, natural feature).
rule: |
- Hierarchical format allowed (e.g., "Paris, France").
- Avoid overly vague terms (e.g., "there") unless resolved.
- Max length: 80 characters.
- type: other
description: Entities that do not fit any above category.
rule: |
- Use sparingly; prefer mapping to a defined type when possible.
- Still provide a meaningful label.
- Max length: 80 characters.
relation:
description: >-
You are an expert in extracting semantic relations between entities.
fields:
- type: owns
description: Ownership or possession (legal or de facto).
rule: |
- Direction from owner to owned: (A owns B).
- type: part_of
description: Mereological relation — component to whole.
rule: |
- Direction from part to whole: (A part_of B).
- type: caused_by
description: Causal relation — event, action, or state leads to another.
rule: |
- Direction from effect to cause: (A caused_by B) meaning B causes A.
- type: regulates
description: Regulatory or governing relation (law/standard controls entity).
rule: |
- Direction from regulator to regulated: (A regulates B).
- type: located_in
description: Spatial containment — entity situated inside a location.
rule: |
- Direction from located entity to containing location: (A located_in B).
- type: other
description: Any meaningful relation not covered by the above types.
rule: |
- Provide an explicit label in a "relation_label" field.
- Direction must be clear.
global_rules: ''

69
api/db/init_data/compilation_templates/mind_map.yaml Normal file
View File

@@ -0,0 +1,69 @@
kind: mind_map
display_name: Mind map - Radial concept hierarchy
config:
kind: mind_map
entity:
description: >-
You are a robust mind-map extractor. Extract the central concept,
major branches, and supporting sub-branches needed to build a
non-linear visual hierarchy around the main subject.
fields:
- type: central_topic
description: The main idea, goal, problem, or subject placed at the center of the mind map.
rule: |
- Prefer a concise noun phrase that represents the whole chunk or document section.
- Use the source wording when it clearly names the topic.
- If no central topic can be inferred, output "-1".
- Max length: 80 characters.
- type: branch
description: A major theme, category, or dimension radiating from the central topic.
rule: |
- Use broad categories that organize multiple related details.
- Avoid duplicating the central topic as a branch.
- Keep labels concise and parallel when possible.
- Max length: 80 characters.
- type: sub_branch
description: A specific detail, task, supporting concept, example, or item under a branch.
rule: |
- Attach each sub-branch to the most specific relevant parent branch.
- Preserve important source terms, names, quantities, and constraints.
- Split unrelated details into separate sub-branches.
- Max length: 120 characters.
- type: keyword
description: A compact supporting word or phrase that clarifies a branch or sub-branch.
rule: |
- Use only meaningful terms that improve scanability.
- Do not include filler words or generic labels.
- Max length: 60 characters.
relation:
description: >-
You are an expert hierarchical reasoning assistant specializing in
mind-map structure. Link each concept to its parent so the result
forms an intuitive radial hierarchy.
fields:
- type: has_branch
description: The central topic contains a major branch.
rule: |
- Direction from central topic to branch: (A has_branch B).
- Use only for first-level branches radiating from the center.
- type: has_sub_branch
description: A branch or sub-branch contains a more specific child concept.
rule: |
- Direction from parent to child: (A has_sub_branch B).
- Use recursively for second-level and deeper details.
- type: supports
description: A keyword, example, task, or detail supports a parent concept.
rule: |
- Direction from supporting item to supported concept: (A supports B).
- Use when the item explains, evidences, or clarifies the parent.
- type: related_to
description: A cross-link between two concepts that are associated but not parent-child.
rule: |
- Use sparingly; prefer parent-child hierarchy when possible.
- Direction may follow the stronger explanatory dependency.
global_rules: |
- Build a mind map around one central concept whenever possible.
- Arrange major themes as first-level branches and details as recursive sub-branches.
- Keep labels short enough for visual nodes; avoid sentence-length node names.
- Preserve a clear hierarchy; do not create cycles in parent-child relations.
- Use the same language as the source text unless normalization is necessary.

33
api/db/init_data/compilation_templates/page_index.yaml Normal file
View File

@@ -0,0 +1,33 @@
kind: page_index
display_name: PageIndex — Hierarchical table of contents
config:
kind: page_index
entity:
description: >-
You are a robust Table-of-Contents (TOC) extractor.
fields:
- type: title
description: the heading text (clean, no page numbers or leader dots)
rule: |
- Length restriction:
• Chinese heading: ≤25 characters
• English heading: ≤80 characters
- "title" must be non-empty (or exactly "-1").
- If any part of a chunk has no valid heading, output that part as {"title":"-1", ...}.
- If a chunk contains multiple headings, expand them in order:
- Each heading → {"title":"...","chunk_id":"<chunk_ID>"}.
- When ambiguous, prefer "-1" unless the text strongly looks like a heading.
- Keep language of "title" the same as the input.
- Prefix like following must be titles: 第x章, 第N条, 第N节, 1, 1.1, 1.1.1 ...
relation:
description: >-
You are an expert logical reasoning assistant specializing in hierarchical titles.
fields:
- type: include
description: Upper-level title includes lower-level title.
rule: |
- "-1" is an invalid title; it does not belong to or include any other titles.
- Must follow the hierarchical index/numbering (e.g., "1", "2.1", "3.2.5") when present.
- Keep language of "title" the same as the input.
- 第N章 must include 第N条 or 第N节.
global_rules: ''

48
api/db/init_data/compilation_templates/timeline.yaml Normal file
View File

@@ -0,0 +1,48 @@
kind: timeline
display_name: List (Timeline) — Chronological events / Graph
config:
kind: timeline
entity:
description: >-
You are a robust events-timeline extractor.
fields:
- type: timestamp
description: the date or time reference (clean, no extraneous text)
rule: |
- Format: prefer ISO 8601 (YYYY-MM-DD) or a normalized human-readable form
(e.g., "March 5, 2024").
- If only a relative time (e.g., "yesterday", "next week"), convert to
absolute when the context allows, else keep as is.
- "timestamp" must be non-empty (or exactly "-1" if no valid time/date).
- If a chunk contains multiple events with distinct timestamps, expand them
in chronological order:
- Each event → {"timestamp":"...","event":"...","chunk_id":"<chunk_ID>"}.
- When ambiguous, prefer "-1" unless the text strongly indicates a specific
time/date.
- Keep language and numbering style of "timestamp" consistent with the input.
- type: event
description: the event description associated with the timestamp (concise, no metadata)
rule: |
- Length restriction:
• Chinese event: ≤40 characters
• English event: ≤120 characters
- "event" must be non-empty (or exactly "-1") if no valid event description.
- If no valid event but a timestamp exists, output
{"timestamp":"...","event":"-1", ...}.
- Preserve the core action and key entities; omit redundant phrasing.
relation:
description: >-
You are an expert sequential reasoning assistant specializing in chronological
timelines.
fields:
- type: ordered
description: Events are arranged in strict chronological order (earliest to latest).
rule: |
- "-1" for timestamp or event indicates invalid or missing data; such entries
do not participate in ordering.
- Must follow explicit or inferred temporal indicators
(e.g., "then", "afterwards", "at 3 PM").
- If multiple events share the same timestamp, preserve their textual order
as a sub-list.
- Keep language and date formatting consistent across the timeline.
global_rules: ''

31
api/db/init_data/compilation_templates/tree.yaml Normal file
View File

@@ -0,0 +1,31 @@
kind: tree
display_name: Tree — RAPTOR-based document tree
description: >-
Recursive Abstractive Processing for Tree-Organized Retrieval over a
single document's chunks. Produces a hierarchical summary tree
(root summary → cluster summaries → leaf cluster source_chunk_ids).
config:
kind: tree
raptor:
prompt: |-
Please summarize the following paragraphs. Be careful with the numbers, do not make things up. Paragraphs as following:
{cluster_content}
The above is the content you need to summarize.
max_token: 512
threshold: 0.1
# Entity / relation collections are not used by the tree kind but kept
# as empty stubs so the form schema's shared validators see a stable
# shape regardless of kind.
entity:
description: ''
fields:
- type: ''
description: ''
rule: ''
relation:
description: ''
fields:
- type: ''
description: ''
rule: ''
global_rules: ''

394
api/db/services/compilation_template_group_service.py Normal file
View File

@@ -0,0 +1,394 @@
#
# Copyright 2026 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 peewee import fn
from api.db.db_models import DB, CompilationTemplate, CompilationTemplateGroup
from api.db.services.common_service import CommonService
from common.constants import StatusEnum
from common.misc_utils import get_uuid
SCOPE_FILE = "file"
SCOPE_DATASET = "dataset"
class GroupValidationError(ValueError):
pass
def _derive_scope(templates: list[dict]) -> str:
"""Derive the group's scope from its child templates.
One artifacts child = dataset scope (and must be the only child).
Otherwise file scope, with no artifacts allowed.
"""
if not templates:
raise GroupValidationError("A template group must contain at least one template.")
kinds = [str((t or {}).get("kind") or "").strip() for t in templates]
artifact_count = sum(1 for k in kinds if k == "artifacts")
if artifact_count > 0:
if artifact_count != 1 or len(templates) != 1:
raise GroupValidationError("An artifacts template cannot be combined with other templates in the same group.")
return SCOPE_DATASET
_enforce_single_rechunk_tree(templates)
return SCOPE_FILE
def _enforce_single_rechunk_tree(templates: list[dict]) -> None:
"""At most one tree-kind child in the group may enable re-chunking.
Re-chunking soft-deletes the doc's original chunks via
``available_int=0`` and inserts merged replacements; running two
such templates would race on the same source chunks and produce
non-deterministic output. Per-tenant invariant is enforced
server-side here and mirrored client-side in
``group-interface.ts``.
"""
rechunk_trees = 0
for t in templates:
if str((t or {}).get("kind") or "").strip() != "tree":
continue
cfg = (t or {}).get("config") or {}
raptor = (cfg or {}).get("raptor") or {}
if bool(raptor.get("rechunk")):
rechunk_trees += 1
if rechunk_trees > 1:
raise GroupValidationError("Only one tree template in a group may enable re-chunking.")
class CompilationTemplateGroupService(CommonService):
model = CompilationTemplateGroup
@classmethod
def ensure_table(cls) -> None:
if not cls.model.table_exists():
cls.model.create_table(safe=True)
# ------------------------------------------------------------------
# Read paths
# ------------------------------------------------------------------
@classmethod
def _group_to_dict(cls, group: CompilationTemplateGroup, templates: list[CompilationTemplate]) -> dict:
from api.db.services.compilation_template_service import CompilationTemplateService
return {
"id": group.id,
"name": group.name,
"description": group.description or "",
"scope": group.scope,
"create_time": group.create_time,
"update_time": group.update_time,
"templates": [CompilationTemplateService._to_saved_dict(t) for t in templates],
}
@classmethod
@DB.connection_context()
def list_saved(
cls,
tenant_id: str,
keywords: str = "",
scope: str = "",
orderby: str = "create_time",
desc: bool = True,
) -> list[dict]:
cls.ensure_table()
query = cls.model.select().where(
cls.model.tenant_id == tenant_id,
cls.model.status == StatusEnum.VALID.value,
)
if keywords:
query = query.where(cls.model.name.contains(keywords))
if scope:
query = query.where(cls.model.scope == scope)
if not hasattr(cls.model, orderby):
orderby = "create_time"
order_field = getattr(cls.model, orderby)
query = query.order_by(order_field.desc() if desc else order_field.asc())
groups = list(query)
if not groups:
return []
group_ids = [g.id for g in groups]
children = list(
CompilationTemplate.select()
.where(
CompilationTemplate.group_id.in_(group_ids),
CompilationTemplate.status == StatusEnum.VALID.value,
)
.order_by(CompilationTemplate.create_time.asc())
)
children_by_group: dict[str, list[CompilationTemplate]] = {gid: [] for gid in group_ids}
for child in children:
children_by_group.setdefault(child.group_id, []).append(child)
return [cls._group_to_dict(g, children_by_group.get(g.id, [])) for g in groups]
@classmethod
@DB.connection_context()
def get_saved(cls, group_id: str, tenant_id: str) -> dict | None:
group = cls.model.get_or_none(
cls.model.id == group_id,
cls.model.tenant_id == tenant_id,
cls.model.status == StatusEnum.VALID.value,
)
if not group:
return None
children = list(
CompilationTemplate.select()
.where(
CompilationTemplate.group_id == group_id,
CompilationTemplate.status == StatusEnum.VALID.value,
)
.order_by(CompilationTemplate.create_time.asc())
)
return cls._group_to_dict(group, children)
@classmethod
@DB.connection_context()
def list_for_resolution(cls, tenant_id: str) -> list[dict]:
"""Light list used by frontend pickers (dataset parse-config dropdown).
Returns one row per group with just the fields the picker needs +
the child template ids so the orchestrator can resolve them later.
"""
cls.ensure_table()
groups = list(
cls.model.select().where(
cls.model.tenant_id == tenant_id,
cls.model.status == StatusEnum.VALID.value,
)
)
if not groups:
return []
group_ids = [g.id for g in groups]
kid_pairs = list(
CompilationTemplate.select(
CompilationTemplate.group_id,
CompilationTemplate.id,
CompilationTemplate.kind,
CompilationTemplate.name,
).where(
CompilationTemplate.group_id.in_(group_ids),
CompilationTemplate.status == StatusEnum.VALID.value,
)
)
by_group: dict[str, list[dict]] = {}
for child in kid_pairs:
by_group.setdefault(child.group_id, []).append({"id": child.id, "kind": child.kind, "name": child.name})
return [
{
"id": g.id,
"name": g.name,
"description": g.description or "",
"scope": g.scope,
"templates": by_group.get(g.id, []),
}
for g in groups
]
@classmethod
@DB.connection_context()
def name_exists(cls, tenant_id: str, name: str, exclude_id: str = "") -> bool:
cls.ensure_table()
query = cls.model.select(fn.COUNT(cls.model.id)).where(
cls.model.tenant_id == tenant_id,
cls.model.name == name,
cls.model.status == StatusEnum.VALID.value,
)
if exclude_id:
query = query.where(cls.model.id != exclude_id)
return query.scalar() > 0
@classmethod
@DB.connection_context()
def resolve_template_ids(cls, group_id: str, tenant_id: str) -> list[str]:
"""Resolve a group id to its child template ids. Used by the orchestrator
when reading ``parser_config.compilation_template_group_id``.
"""
cls.ensure_table()
group = cls.model.get_or_none(
cls.model.id == group_id,
cls.model.tenant_id == tenant_id,
cls.model.status == StatusEnum.VALID.value,
)
if not group:
return []
rows = list(
CompilationTemplate.select(CompilationTemplate.id)
.where(
CompilationTemplate.group_id == group_id,
CompilationTemplate.status == StatusEnum.VALID.value,
)
.order_by(CompilationTemplate.create_time.asc())
)
return [r.id for r in rows]
# ------------------------------------------------------------------
# Write paths
# ------------------------------------------------------------------
@classmethod
@DB.connection_context()
def create_group(cls, tenant_id: str, name: str, description: str, templates: list[dict]) -> dict:
cls.ensure_table()
scope = _derive_scope(templates)
group_id = get_uuid()
with DB.atomic():
CompilationTemplateGroup.create(
id=group_id,
tenant_id=tenant_id,
name=name,
description=description or "",
scope=scope,
status=StatusEnum.VALID.value,
)
for i, child in enumerate(templates):
cls._insert_child(group_id, tenant_id, child, index=i)
saved = cls.get_saved(group_id, tenant_id)
assert saved is not None
return saved
@classmethod
@DB.connection_context()
def update_group(
cls,
group_id: str,
tenant_id: str,
name: str | None,
description: str | None,
templates: list[dict] | None,
) -> dict | None:
cls.ensure_table()
existing = cls.model.get_or_none(
cls.model.id == group_id,
cls.model.tenant_id == tenant_id,
cls.model.status == StatusEnum.VALID.value,
)
if not existing:
return None
with DB.atomic():
updates: dict = {}
if name is not None:
updates["name"] = name
if description is not None:
updates["description"] = description
if templates is not None:
updates["scope"] = _derive_scope(templates)
if updates:
cls.model.update(**updates).where(cls.model.id == group_id).execute()
if templates is not None:
# Soft-delete previous children (en-bloc replace). Simpler than
# diffing and acceptable given small N — child IDs are not
# referenced externally (parser_config keys the group, not its
# children).
CompilationTemplate.update(status=StatusEnum.INVALID.value).where(
CompilationTemplate.group_id == group_id,
CompilationTemplate.status == StatusEnum.VALID.value,
).execute()
for i, child in enumerate(templates):
cls._insert_child(group_id, tenant_id, child, index=i)
return cls.get_saved(group_id, tenant_id)
@classmethod
@DB.connection_context()
def delete_group(cls, group_id: str, tenant_id: str) -> bool:
cls.ensure_table()
existing = cls.model.get_or_none(
cls.model.id == group_id,
cls.model.tenant_id == tenant_id,
cls.model.status == StatusEnum.VALID.value,
)
if not existing:
return False
with DB.atomic():
cls.model.update(status=StatusEnum.INVALID.value).where(cls.model.id == group_id).execute()
CompilationTemplate.update(status=StatusEnum.INVALID.value).where(
CompilationTemplate.group_id == group_id,
CompilationTemplate.status == StatusEnum.VALID.value,
).execute()
return True
@classmethod
def _insert_child(
cls,
group_id: str,
tenant_id: str,
child: dict,
*,
index: int,
) -> None:
kind = str((child or {}).get("kind") or "").strip()
name = str((child or {}).get("name") or "").strip()
config = (child or {}).get("config") or {}
if not kind or not name or not isinstance(config, dict):
raise GroupValidationError("Each template must include a name, kind, and config object.")
from api.db.services.compilation_template_service import CompilationTemplateService
config = CompilationTemplateService.fill_config_default_llm(config, tenant_id)
template_id = get_uuid()
CompilationTemplate.create(
id=template_id,
tenant_id=tenant_id,
group_id=group_id,
name=name,
description=str((child or {}).get("description") or ""),
kind=kind,
config=config,
is_builtin=False,
status=StatusEnum.VALID.value,
)
# ------------------------------------------------------------------
# Lookup helpers used by the orchestrator
# ------------------------------------------------------------------
@classmethod
@DB.connection_context()
def get_for_kb(cls, group_id: str, tenant_id: str) -> dict | None:
"""Like ``get_saved`` but returns ``None`` quietly and avoids the
``_to_saved_dict`` LLM-lookup branch — for orchestrator use where
we only need the scope + child rows.
"""
cls.ensure_table()
group = cls.model.get_or_none(
cls.model.id == group_id,
cls.model.tenant_id == tenant_id,
cls.model.status == StatusEnum.VALID.value,
)
if not group:
return None
children = list(
CompilationTemplate.select()
.where(
CompilationTemplate.group_id == group_id,
CompilationTemplate.status == StatusEnum.VALID.value,
)
.order_by(CompilationTemplate.create_time.asc())
)
return {
"id": group.id,
"name": group.name,
"scope": group.scope,
"template_ids": [c.id for c in children],
"templates_by_kind": {c.kind: c.id for c in children},
}

231
api/db/services/compilation_template_service.py Normal file
View File

@@ -0,0 +1,231 @@
#
# Copyright 2026 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.
#
import logging
import os
from peewee import fn
from ruamel.yaml import YAML
from api.db.db_models import CompilationTemplate, DB
from api.db.services.common_service import CommonService
from common.constants import StatusEnum
from common.file_utils import get_project_base_directory
class CompilationTemplateService(CommonService):
model = CompilationTemplate
@classmethod
def fill_config_default_llm(cls, config: dict, tenant_id: str | None) -> dict:
if not isinstance(config, dict) or config.get("llm_id") or not tenant_id:
return config
try:
from api.db.services.user_service import TenantService
ok, tenant = TenantService.get_by_id(tenant_id)
if ok and getattr(tenant, "llm_id", None):
config = dict(config)
config["llm_id"] = tenant.llm_id
except Exception:
logging.exception(
"compilation_template: llm_id default-fill lookup failed for tenant=%s",
tenant_id,
)
return config
@classmethod
def fill_default_llm_for_templates(cls, templates: list[dict], tenant_id: str | None) -> list[dict]:
if not tenant_id:
return templates
filled = []
for template in templates:
item = dict(template)
item["config"] = cls.fill_config_default_llm(item.get("config") or {}, tenant_id)
filled.append(item)
return filled
@classmethod
def _sort_builtins(cls, templates: list[dict]) -> list[dict]:
return sorted(
templates,
key=lambda template: (
template.get("kind") == "empty" or template.get("id") == "empty",
template.get("display_name") or template.get("name") or "",
),
)
@classmethod
@DB.connection_context()
def ensure_table(cls) -> None:
if not cls.model.table_exists():
cls.model.create_table(safe=True)
@classmethod
def _to_saved_dict(cls, template: CompilationTemplate) -> dict:
data = template.to_dict()
config = data.get("config") or {}
# Lazy-fill llm_id with the tenant's default chat model so the
# frontend always sees a value (legacy templates predate the
# field). The DB row is left untouched — this is a read-side
# default. If the tenant has no default chat model set,
# silently leave llm_id absent and let the caller fall back
# however it likes.
if isinstance(config, dict) and not config.get("llm_id"):
tenant_id = data.get("tenant_id")
if tenant_id:
try:
from api.db.services.user_service import TenantService
ok, tenant = TenantService.get_by_id(tenant_id)
if ok and getattr(tenant, "llm_id", None):
config = dict(config)
config["llm_id"] = tenant.llm_id
except Exception:
logging.exception(
"compilation_template: llm_id lazy-fill lookup failed for tenant=%s",
tenant_id,
)
return {
"id": data["id"],
"name": data["name"],
"description": data.get("description") or "",
"kind": data["kind"],
"config": cls.fill_config_default_llm(config, data.get("tenant_id")),
"create_time": data.get("create_time"),
"update_time": data.get("update_time"),
}
@classmethod
def _to_builtin_dict(cls, template: CompilationTemplate) -> dict:
data = template.to_dict()
return {
"id": data["id"],
"kind": data["kind"],
"display_name": data["name"],
"description": data.get("description") or "",
"config": data.get("config") or {},
}
@classmethod
@DB.connection_context()
def list_saved(cls, tenant_id: str, keywords: str = "", kind: str = "", orderby: str = "create_time", desc: bool = True) -> list[dict]:
query = cls.model.select().where(
cls.model.tenant_id == tenant_id,
not cls.model.is_builtin,
cls.model.status == StatusEnum.VALID.value,
)
if keywords:
query = query.where(cls.model.name.contains(keywords))
if kind:
query = query.where(cls.model.kind == kind)
if not hasattr(cls.model, orderby):
orderby = "create_time"
order_field = getattr(cls.model, orderby)
query = query.order_by(order_field.desc() if desc else order_field.asc())
return [cls._to_saved_dict(template) for template in query]
@classmethod
@DB.connection_context()
def list_builtins(cls) -> list[dict]:
cls.ensure_table()
query = cls.model.select().where(cls.model.is_builtin, cls.model.status == StatusEnum.VALID.value).order_by(cls.model.create_time.asc(), cls.model.name.asc())
return cls._sort_builtins([cls._to_builtin_dict(template) for template in query])
@classmethod
@DB.connection_context()
def get_saved(cls, template_id: str, tenant_id: str) -> dict | None:
template = cls.model.get_or_none(
cls.model.id == template_id,
cls.model.tenant_id == tenant_id,
not cls.model.is_builtin,
cls.model.status == StatusEnum.VALID.value,
)
return cls._to_saved_dict(template) if template else None
@classmethod
@DB.connection_context()
def name_exists(cls, tenant_id: str, name: str, exclude_id: str = "") -> bool:
query = cls.model.select(fn.COUNT(cls.model.id)).where(
cls.model.tenant_id == tenant_id,
cls.model.name == name,
not cls.model.is_builtin,
cls.model.status == StatusEnum.VALID.value,
)
if exclude_id:
query = query.where(cls.model.id != exclude_id)
return query.scalar() > 0
@classmethod
@DB.connection_context()
def upsert_builtin(cls, template: dict) -> None:
template_id = template["id"]
existing = cls.model.get_or_none(cls.model.id == template_id)
data = {
"id": template_id,
"tenant_id": None,
"name": template["name"],
"description": template.get("description", ""),
"kind": template["kind"],
"config": template["config"],
"is_builtin": True,
"status": StatusEnum.VALID.value,
}
if existing:
cls.update_by_id(template_id, data)
else:
cls.insert(**data)
@classmethod
def seed_builtins_from_files(cls) -> None:
cls.ensure_table()
for template in cls.load_builtins_from_files():
cls.upsert_builtin(template)
@classmethod
def load_builtins_from_files(cls) -> list[dict]:
template_dir = os.path.join(get_project_base_directory(), "api", "db", "init_data", "compilation_templates")
if not os.path.exists(template_dir):
logging.warning("Missing compilation templates: %s", template_dir)
return []
templates = []
yaml = YAML(typ="safe", pure=True)
for filename in sorted(os.listdir(template_dir)):
if not filename.endswith((".yaml", ".yml")):
continue
template_path = os.path.join(template_dir, filename)
try:
with open(template_path, "r", encoding="utf-8") as f:
template = yaml.load(f) or {}
kind = template.get("kind")
display_name = template.get("display_name")
config = template.get("config")
if not kind or not display_name or not isinstance(config, dict):
logging.warning("Skipping invalid compilation template file: %s", template_path)
continue
templates.append(
{
"id": os.path.splitext(filename)[0],
"name": display_name,
"description": template.get("description", ""),
"kind": kind,
"config": config,
}
)
except Exception as e:
logging.exception("Add compilation template error for %s: %s", template_path, e)
return cls._sort_builtins(templates)

3
api/db/services/dialog_service.py
View File

@@ -42,7 +42,7 @@ from api.utils.reference_metadata_utils import (
from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_type_by_name
from common.time_utils import current_timestamp, datetime_format
from common.text_utils import normalize_arabic_digits
from rag.graphrag.general.mind_map_extractor import MindMapExtractor
from rag.advanced_rag.knowlege_compile.mind_map_extractor import MindMapExtractor
from rag.advanced_rag import DeepResearcher
from rag.app.tag import label_question
from rag.nlp.search import index_name
@@ -761,6 +761,7 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
retrieval_ts = timer()
if not knowledges and prompt_config.get("empty_response"):
empty_res = prompt_config["empty_response"]
yield {"answer": empty_res, "reference": {}, "prompt": "", "audio_binary": None, "final": False}
yield {"answer": empty_res, "reference": kbinfos, "prompt": "\n\n### Query:\n%s" % " ".join(questions), "audio_binary": tts(tts_mdl, empty_res), "final": True}
return

111
api/db/services/document_service.py
View File

@@ -373,10 +373,14 @@ class DocumentService(CommonService):
@DB.connection_context()
def list_doc_headers_by_kb_and_source_type(cls, kb_id, source_type, page_size=500):
fields = [cls.model.id, cls.model.kb_id, cls.model.source_type, cls.model.name]
docs = cls.model.select(*fields).where(
cls.model.kb_id == kb_id,
cls.model.source_type == source_type,
).order_by(cls.model.create_time.asc())
docs = (
cls.model.select(*fields)
.where(
cls.model.kb_id == kb_id,
cls.model.source_type == source_type,
)
.order_by(cls.model.create_time.asc())
)
offset = 0
res = []
while True:
@@ -402,10 +406,14 @@ class DocumentService(CommonService):
rows and the resulting map would silently miss entries.
"""
fields = [cls.model.id, cls.model.content_hash]
docs = cls.model.select(*fields).where(
cls.model.kb_id == kb_id,
cls.model.source_type == source_type,
).order_by(cls.model.create_time.asc())
docs = (
cls.model.select(*fields)
.where(
cls.model.kb_id == kb_id,
cls.model.source_type == source_type,
)
.order_by(cls.model.create_time.asc())
)
offset = 0
result: dict[str, str] = {}
while True:
@@ -489,6 +497,20 @@ class DocumentService(CommonService):
except Exception as e:
logging.error(f"Failed to delete chunks from doc store for document {doc.id}: {e}")
# Prune this doc's line from the KB's tree-kind navigation
# markdown (best-effort — the markdown is a downstream artifact,
# and failure here must not block the document delete).
try:
from rag.advanced_rag.knowlege_compile.dataset_nav import (
remove_dataset_nav_doc_sync,
)
remove_dataset_nav_doc_sync(tenant_id, doc.kb_id, doc.id)
except Exception as e:
logging.warning(
f"Failed to prune dataset_nav for document {doc.id}: {e}",
)
# Delete document metadata (non-critical, log and continue)
try:
DocMetadataService.delete_document_metadata(doc.id, doc.kb_id, tenant_id)
@@ -571,21 +593,20 @@ class DocumentService(CommonService):
@classmethod
@DB.connection_context()
def get_unfinished_docs(cls):
fields = [cls.model.id, cls.model.process_begin_at, cls.model.parser_config, cls.model.progress_msg,
cls.model.run, cls.model.parser_id]
unfinished_task_query = Task.select(Task.doc_id).where(
(Task.progress >= 0) & (Task.progress < 1)
)
fields = [cls.model.id, cls.model.process_begin_at, cls.model.parser_config, cls.model.progress_msg, cls.model.run, cls.model.parser_id]
unfinished_task_query = Task.select(Task.doc_id).where((Task.progress >= 0) & (Task.progress < 1))
docs_with_non_failed_tasks = Task.select(Task.doc_id).where(Task.progress >= 0).distinct()
docs = cls.model.select(*fields).where(
cls.model.status == StatusEnum.VALID.value,
~(cls.model.type == FileType.VIRTUAL.value),
((cls.model.run.is_null(True)) | (cls.model.run != TaskStatus.CANCEL.value)),
(((cls.model.progress < 1) & (cls.model.progress > 0)) |
(cls.model.id.in_(unfinished_task_query)) |
((cls.model.progress == -1) & (cls.model.run == TaskStatus.FAIL.value) &
(cls.model.id.in_(docs_with_non_failed_tasks))))) # including GraphRAG/RAPTOR/Mindmap; re-sync failed docs
(
((cls.model.progress < 1) & (cls.model.progress > 0))
| (cls.model.id.in_(unfinished_task_query))
| ((cls.model.progress == -1) & (cls.model.run == TaskStatus.FAIL.value) & (cls.model.id.in_(docs_with_non_failed_tasks)))
),
) # including GraphRAG/RAPTOR/Mindmap; re-sync failed docs
return list(docs.dicts())
@classmethod
@@ -604,8 +625,7 @@ class DocumentService(CommonService):
)
if num == 0:
logging.error(
"increment_chunk_num: no document matched doc_id=%s kb_id=%s "
"token_num=%s chunk_num=%s duration=%s",
"increment_chunk_num: no document matched doc_id=%s kb_id=%s token_num=%s chunk_num=%s duration=%s",
doc_id,
kb_id,
token_num,
@@ -623,8 +643,7 @@ class DocumentService(CommonService):
)
if num == 0:
logging.error(
"increment_chunk_num: no knowledgebase matched kb_id=%s for doc_id=%s "
"token_num=%s chunk_num=%s duration=%s",
"increment_chunk_num: no knowledgebase matched kb_id=%s for doc_id=%s token_num=%s chunk_num=%s duration=%s",
kb_id,
doc_id,
token_num,
@@ -660,8 +679,7 @@ class DocumentService(CommonService):
)
if num == 0:
logging.error(
"decrement_chunk_num: no knowledgebase matched kb_id=%s for doc_id=%s "
"token_num=%s chunk_num=%s duration=%s",
"decrement_chunk_num: no knowledgebase matched kb_id=%s for doc_id=%s token_num=%s chunk_num=%s duration=%s",
kb_id,
doc_id,
token_num,
@@ -1071,7 +1089,7 @@ def queue_raptor_o_graphrag_tasks(sample_doc, ty, priority, fake_doc_id="", doc_
"""
if doc_ids is None:
doc_ids = []
assert ty in ["graphrag", "raptor", "mindmap"], "type should be graphrag, raptor or mindmap"
assert ty in ["graphrag", "raptor", "mindmap", "artifact", "skill"], "type should be graphrag, raptor, mindmap, artifact or skill"
chunking_config = DocumentService.get_chunking_config(sample_doc["id"])
hasher = xxhash.xxh64()
@@ -1102,6 +1120,51 @@ def queue_raptor_o_graphrag_tasks(sample_doc, ty, priority, fake_doc_id="", doc_
return task["id"]
def queue_per_doc_raptor_task(doc, priority):
"""Queue a doc-scoped RAPTOR task.
Distinct from :func:`queue_raptor_o_graphrag_tasks` (which is KB-scoped
and uses ``GRAPH_RAPTOR_FAKE_DOC_ID`` as the task's ``doc_id`` so it
fans out across the dataset). Here the task's ``doc_id`` is the real
document id, so ``TaskHandler._run_raptor`` runs only on this doc's
chunks and the RAPTOR summaries it produces are scoped to this doc.
Triggered automatically at the tail of standard chunking when the
doc's ``parser_config["raptor"]["use_raptor"]`` is true. No
cross-task dedup — within one chunking-task execution this helper is
called at most once, which is the only invariant the caller needs.
"""
chunking_config = DocumentService.get_chunking_config(doc["id"])
hasher = xxhash.xxh64()
for field in sorted(chunking_config.keys()):
hasher.update(str(chunking_config[field]).encode("utf-8"))
task = {
"id": get_uuid(),
"doc_id": doc["id"],
"from_page": MAXIMUM_TASK_PAGE_NUMBER,
"to_page": MAXIMUM_TASK_PAGE_NUMBER,
"task_type": "raptor",
"progress_msg": datetime.now().strftime("%H:%M:%S") + " created task raptor",
"begin_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
}
for field in ["doc_id", "from_page", "to_page"]:
hasher.update(str(task[field]).encode("utf-8"))
hasher.update(b"raptor")
task["digest"] = hasher.hexdigest()
bulk_insert_into_db(Task, [task], True)
# Redis message carries ``doc_ids`` for downstream consumers
# (TaskHandler._run_raptor reads it). Identical to the fake-doc
# path's convention so we don't have to special-case the executor.
task["doc_ids"] = [doc["id"]]
assert REDIS_CONN.queue_product(
settings.get_svr_queue_name(priority, "raptor"),
message=task,
), "Can't access Redis. Please check the Redis' status."
return task["id"]
def get_queue_length(priority, suffix="common"):
group_info = REDIS_CONN.queue_info(settings.get_svr_queue_name(priority, suffix), SVR_CONSUMER_GROUP_NAME)
if not group_info:

523
api/db/services/file_commit_service.py
View File

@@ -15,11 +15,13 @@
#
import datetime
import difflib
import hashlib
import json
import logging
from typing import Optional
from api.db.db_models import DB, FileCommit, FileCommitItem, File
from api.db.db_models import DB, FileCommit, FileCommitItem, File, User
from api.db.services.common_service import CommonService
from api.db.services.file_service import FileService
from common import settings
@@ -29,6 +31,149 @@ from common.time_utils import current_timestamp, datetime_format
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------
# Artifact-commit extension
# ---------------------------------------------------------------------
# Artifact-page saves used to land in the retired ``ArtifactCommit`` table.
# They now flow through :class:`FileCommitService.record_page_edit`, which
# writes one FileCommit + one FileCommitItem per save with the artifact
# columns populated (title/comments on FileCommit; diff/content_after_*/
# slug_kwd/page_type_kwd on FileCommitItem).
#
# ``file_id`` for these commits is a stable content-hash of ``(kb_id, slug)``
# so per-page history queries can filter on it without a real File row —
# no pseudo-File / virtual-folder machinery is created, so the workspace
# UI stays free of ghost entries.
#
# ``folder_id`` is set to ``kb_id`` directly. The datasets URL prefix
# (``/datasets/<kb_id>/commits``) resolves the entity id to itself for
# this scope; workspace file-commit browsing still uses ``/folders/*`` or
# ``/workspace/*`` with the real folder id.
#
# Content storage for ``content_after`` is switched by a module-level
# constant so ops can move blobs between MinIO and the doc-store index
# without touching the schema.
ARTIFACT_CONTENT_STORAGE = "minio" # one of {"minio", "es"}
_ARTIFACT_COMMIT_BUCKET_PREFIX = ".artifact_commits"
_ARTIFACT_ES_KWD = "artifact_commit_content"
def _artifact_file_id(kb_id: str, slug: str) -> str:
"""Deterministic 32-char id for the artifact-page 'file' identity.
Not a real File row — just an index key that groups all commits for
the same page. Hashed so slugs longer than 32 chars still fit.
"""
return hashlib.md5(f"{kb_id}:{slug}".encode("utf-8")).hexdigest()
def _unified_diff(before: str, after: str, slug: str) -> str:
"""Return a unified diff between two markdown strings, or '' if equal."""
if (before or "") == (after or ""):
return ""
return "".join(
difflib.unified_diff(
(before or "").splitlines(keepends=True),
(after or "").splitlines(keepends=True),
fromfile=f"a/{slug}",
tofile=f"b/{slug}",
n=3,
)
)
def _store_content_after(kb_id: str, content: str) -> tuple[str, str]:
"""Persist ``content`` per :data:`ARTIFACT_CONTENT_STORAGE`. Returns
``(storage_kind, location)`` for the row's persistence columns.
Content-addressed by SHA-256 so re-saves with identical bodies share
the same blob.
"""
content_bytes = (content or "").encode("utf-8")
content_hash = hashlib.sha256(content_bytes).hexdigest()
if ARTIFACT_CONTENT_STORAGE == "minio":
location = f"{_ARTIFACT_COMMIT_BUCKET_PREFIX}/{content_hash}"
try:
storage = settings.STORAGE_IMPL
if storage is not None:
storage.put(kb_id, location, content_bytes)
except Exception:
logging.exception(
"record_page_edit: MinIO put failed for kb=%s hash=%s",
kb_id,
content_hash,
)
return "minio", location
if ARTIFACT_CONTENT_STORAGE == "es":
# Store as a single doc-store row so the same connector serves
# reads. The row is not retrievable (available_int=0).
from rag.nlp import search as _rag_search
index = _rag_search.index_name(kb_id) # kb-scoped index namespace
payload = {
"id": content_hash,
"kb_id": kb_id,
"doc_id": kb_id,
"compile_kwd": _ARTIFACT_ES_KWD,
"content_with_weight": content or "",
"available_int": 0,
}
try:
settings.docStoreConn.insert([payload], index, kb_id)
except Exception:
logging.exception(
"record_page_edit: ES insert failed for kb=%s hash=%s",
kb_id,
content_hash,
)
return "es", content_hash
# Unknown storage kind — fall through with empty location; the
# detail path treats missing location as "content not recoverable".
logging.warning(
"record_page_edit: unknown ARTIFACT_CONTENT_STORAGE=%r; content not persisted",
ARTIFACT_CONTENT_STORAGE,
)
return "", ""
def _read_content_after(kb_id: str, storage_kind: str, location: str) -> str:
"""Fetch the previously-stored artifact ``content_after`` blob.
Returns ``""`` when the location is empty (workspace commits) or the
blob is missing.
"""
if not location:
return ""
try:
if storage_kind == "minio":
storage = settings.STORAGE_IMPL
if storage is None:
return ""
raw = storage.get(kb_id, location)
if isinstance(raw, (bytes, bytearray)):
return raw.decode("utf-8", errors="replace")
return str(raw or "")
if storage_kind == "es":
from rag.nlp import search as _rag_search
index = _rag_search.index_name(kb_id)
row = settings.docStoreConn.get(location, index, [kb_id])
if isinstance(row, dict):
return row.get("content_with_weight") or ""
return ""
except Exception:
logging.exception(
"get_page_commit: content read failed kb=%s storage=%s loc=%s",
kb_id,
storage_kind,
location,
)
return ""
def _get_file_parent_id(file_id):
"""Look up a file's parent_id from the File table."""
try:
@@ -156,11 +301,13 @@ class FileCommitService(CommonService):
item["new_location"] = obj_key
# Update file record in DB
File.update({
"location": obj_key,
"size": len(content_bytes),
"update_time": current_timestamp(),
}).where(File.id == file_id).execute()
File.update(
{
"location": obj_key,
"size": len(content_bytes),
"update_time": current_timestamp(),
}
).where(File.id == file_id).execute()
# Update tree state
file_parent = _get_file_parent_id(file_id)
@@ -196,11 +343,13 @@ class FileCommitService(CommonService):
item["new_location"] = obj_key
# Update file record
File.update({
"location": obj_key,
"size": len(content_bytes),
"update_time": current_timestamp(),
}).where(File.id == file_id).execute()
File.update(
{
"location": obj_key,
"size": len(content_bytes),
"update_time": current_timestamp(),
}
).where(File.id == file_id).execute()
# Update tree state
file_parent = _get_file_parent_id(file_id)
@@ -222,9 +371,7 @@ class FileCommitService(CommonService):
item["old_location"] = old_location
# Soft-delete the file record
File.update(status="0", update_time=current_timestamp()).where(
File.id == file_id
).execute()
File.update(status="0", update_time=current_timestamp()).where(File.id == file_id).execute()
# Remove from tree state (mark deleted)
if file_id in tree_state:
@@ -237,9 +384,7 @@ class FileCommitService(CommonService):
item["new_name"] = new_name
# Update the file record name
File.update(name=new_name, update_time=current_timestamp()).where(
File.id == file_id
).execute()
File.update(name=new_name, update_time=current_timestamp()).where(File.id == file_id).execute()
# Update tree state
if file_id in tree_state:
@@ -259,9 +404,7 @@ class FileCommitService(CommonService):
def _get_latest_commit(cls, folder_id):
"""Get the latest (chain head) commit for a folder."""
try:
return cls.model.select().where(
cls.model.folder_id == folder_id
).order_by(cls.model.create_time.desc()).first()
return cls.model.select().where(cls.model.folder_id == folder_id).order_by(cls.model.create_time.desc()).first()
except Exception:
return None
@@ -359,27 +502,31 @@ class FileCommitService(CommonService):
if from_entry is not None and to_entry is None:
# Present in from, absent in to → deleted
diff.append({
"file_id": fid,
"file_name": from_name,
"operation": "delete",
"old_hash": from_hash or (from_item.new_hash if from_item else None),
"old_location": from_entry.get("location", "") if isinstance(from_entry, dict) else None,
"new_hash": None,
"new_location": None,
})
diff.append(
{
"file_id": fid,
"file_name": from_name,
"operation": "delete",
"old_hash": from_hash or (from_item.new_hash if from_item else None),
"old_location": from_entry.get("location", "") if isinstance(from_entry, dict) else None,
"new_hash": None,
"new_location": None,
}
)
elif from_entry is None and to_entry is not None:
# Present in to, absent in from → added
diff.append({
"file_id": fid,
"file_name": to_name,
"operation": "add",
"old_hash": None,
"old_location": None,
"new_hash": to_hash or (to_item.new_hash if to_item else None),
"new_location": to_entry.get("location", "") if isinstance(to_entry, dict) else None,
})
diff.append(
{
"file_id": fid,
"file_name": to_name,
"operation": "add",
"old_hash": None,
"old_location": None,
"new_hash": to_hash or (to_item.new_hash if to_item else None),
"new_location": to_entry.get("location", "") if isinstance(to_entry, dict) else None,
}
)
else:
# Both exist — check for changes
@@ -403,15 +550,17 @@ class FileCommitService(CommonService):
if changed:
old_loc = from_entry.get("location", "") if isinstance(from_entry, dict) else None
new_loc = to_entry.get("location", "") if isinstance(to_entry, dict) else None
diff.append({
"file_id": fid,
"file_name": to_name or from_name,
"operation": operation,
"old_hash": from_hash or (from_item.new_hash if from_item else None),
"old_location": old_loc or (from_item.new_location if from_item else None),
"new_hash": to_hash or (to_item.new_hash if to_item else None),
"new_location": new_loc or (to_item.new_location if to_item else None),
})
diff.append(
{
"file_id": fid,
"file_name": to_name or from_name,
"operation": operation,
"old_hash": from_hash or (from_item.new_hash if from_item else None),
"old_location": old_loc or (from_item.new_location if from_item else None),
"new_hash": to_hash or (to_item.new_hash if to_item else None),
"new_location": new_loc or (to_item.new_location if to_item else None),
}
)
return diff
@@ -449,27 +598,33 @@ class FileCommitService(CommonService):
live_hash = _compute_file_hash(folder_id, fid)
committed_hash = committed_entry.get("hash", "")
if live_hash and live_hash != committed_hash:
changes.append({
"file_id": fid,
"file_name": committed_entry.get("name", ""),
"operation": "modify",
})
changes.append(
{
"file_id": fid,
"file_name": committed_entry.get("name", ""),
"operation": "modify",
}
)
else:
if FileService.get_or_none(id=fid) is None:
changes.append({
"file_id": fid,
"file_name": committed_entry.get("name", ""),
"operation": "delete",
})
changes.append(
{
"file_id": fid,
"file_name": committed_entry.get("name", ""),
"operation": "delete",
}
)
# Check for newly added files
for fid, live_file in current_files.items():
if fid not in processed:
changes.append({
"file_id": fid,
"file_name": live_file.name,
"operation": "add",
})
changes.append(
{
"file_id": fid,
"file_name": live_file.name,
"operation": "add",
}
)
return changes
@@ -520,10 +675,14 @@ class FileCommitService(CommonService):
visited = set()
while current_id and current_id not in visited:
visited.add(current_id)
item = FileCommitItem.select().where(
FileCommitItem.commit_id == current_id,
FileCommitItem.file_id == file_id,
).first()
item = (
FileCommitItem.select()
.where(
FileCommitItem.commit_id == current_id,
FileCommitItem.file_id == file_id,
)
.first()
)
if item and item.new_hash:
obj_path = f".objects/{item.new_hash}"
storage_impl = settings.STORAGE_IMPL
@@ -538,6 +697,215 @@ class FileCommitService(CommonService):
return None
# ------------------------------------------------------------------
# Artifact-page commit surface
# ------------------------------------------------------------------
@classmethod
def record_page_edit(
cls,
*,
tenant_id: str,
kb_id: str,
page_type: str,
slug: str,
content_before: str,
content_after: str,
title: Optional[str] = None,
comments: Optional[str] = None,
user_id: Optional[str] = None,
) -> Optional[str]:
"""Persist one artifact-page edit as a FileCommit + FileCommitItem.
Returns the new commit id, or ``None`` when the diff is empty
(no-op save — skipped per the documented v1 contract).
Bypasses :func:`create_commit` because artifact commits have no
real ``File`` row backing them and don't participate in the
workspace ``tree_state`` snapshot chain.
"""
diff_text = _unified_diff(content_before or "", content_after or "", slug)
if not diff_text:
return None
title_ts = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
final_title = f"{(title or '').strip() or f'{title_ts} {slug}'} "
commit_id = get_uuid()
item_id = get_uuid()
file_id = _artifact_file_id(kb_id, slug)
now_ts = current_timestamp()
now_dt = datetime_format(date_time=datetime.datetime.now())
# Persist the post-save markdown per the configured storage.
# A failure here logs but doesn't block the commit row — the diff
# is still meaningful without content_after.
storage_kind, location = _store_content_after(kb_id, content_after or "")
# Chain to the previous commit for this page so the history stays
# ordered even under concurrent writes (auto-regen + user edit).
parent = (
FileCommit.select(FileCommit.id)
.join(
FileCommitItem,
on=(FileCommitItem.commit_id == FileCommit.id),
)
.where((FileCommit.folder_id == kb_id) & (FileCommitItem.file_id == file_id))
.order_by(FileCommit.create_time.desc())
.first()
)
parent_id = parent.id if parent else None
try:
with DB.atomic():
FileCommit(
id=commit_id,
folder_id=kb_id,
parent_id=parent_id,
# ``message`` stays populated with the same string as
# ``title`` so any generic file-commit consumer still
# renders something sensible.
message=final_title[:512],
author_id=user_id or "",
file_count=1,
tree_state=None,
title=final_title[:255],
comments=comments or "",
create_time=now_ts,
create_date=now_dt,
update_time=now_ts,
update_date=now_dt,
).save(force_insert=True)
FileCommitItem(
id=item_id,
commit_id=commit_id,
file_id=file_id,
operation="modify" if content_before else "add",
diff=diff_text,
content_after_storage=storage_kind or None,
content_after_location=location or None,
slug_kwd=slug,
page_type_kwd=page_type,
create_time=now_ts,
create_date=now_dt,
update_time=now_ts,
update_date=now_dt,
).save(force_insert=True)
except Exception:
logging.exception(
"record_page_edit: insert failed for kb=%s slug=%s",
kb_id,
slug,
)
return None
return commit_id
@classmethod
@DB.connection_context()
def list_page_commits(
cls,
tenant_id: str,
kb_id: str,
slug: str,
page: int = 1,
page_size: int = 50,
) -> tuple[int, list[dict]]:
"""Return (total, items) for one artifact page's history.
Filters by ``FileCommitItem.slug_kwd``; joins User for nickname.
Heavy columns (``diff``, ``content_after``) are excluded — the
detail path fetches them lazily.
"""
page = max(int(page or 1), 1)
page_size = max(min(int(page_size or 50), 200), 1)
file_id = _artifact_file_id(kb_id, slug)
base = (
FileCommit.select(
FileCommit.id,
FileCommit.title,
FileCommit.comments,
FileCommit.author_id,
FileCommit.create_time,
FileCommit.create_date,
)
.join(FileCommitItem, on=(FileCommitItem.commit_id == FileCommit.id))
.where((FileCommit.folder_id == kb_id) & (FileCommitItem.file_id == file_id) & (FileCommitItem.slug_kwd == slug))
)
total = base.count()
rows = list(base.order_by(FileCommit.create_time.desc()).paginate(page, page_size).dicts())
# Preserve the previous response key so callers only re-key once.
for r in rows:
r["user_id"] = r.pop("author_id", None)
user_ids = {r["user_id"] for r in rows if r.get("user_id")}
nickname_by_id: dict[str, str] = {}
if user_ids:
try:
for u in User.select(User.id, User.nickname).where(User.id.in_(list(user_ids))).dicts():
nickname_by_id[u["id"]] = u.get("nickname") or ""
except Exception:
logging.exception(
"list_page_commits: nickname lookup failed",
)
for r in rows:
r["user_nickname"] = nickname_by_id.get(r.get("user_id") or "", "")
return total, rows
@classmethod
@DB.connection_context()
def get_page_commit_detail(
cls,
tenant_id: str,
kb_id: str,
commit_id: str,
) -> Optional[dict]:
"""Return one artifact commit including ``diff`` +
``content_after`` (resolved from storage), or ``None`` when not
found. Scoped by ``folder_id == kb_id`` so a leaked commit id
can't be read cross-tenant.
"""
commit = FileCommit.get_or_none(
(FileCommit.id == commit_id) & (FileCommit.folder_id == kb_id),
)
if commit is None:
return None
item = FileCommitItem.get_or_none(FileCommitItem.commit_id == commit_id)
if item is None:
return None
content_after = _read_content_after(
kb_id,
item.content_after_storage or "",
item.content_after_location or "",
)
nickname = ""
if commit.author_id:
try:
u = User.get_or_none(User.id == commit.author_id)
if u is not None:
nickname = u.nickname or ""
except Exception:
pass
return {
"id": commit.id,
"tenant_id": tenant_id,
"kb_id": kb_id,
"page_type_kwd": item.page_type_kwd,
"slug": item.slug_kwd,
"user_id": commit.author_id or None,
"user_nickname": nickname,
"title": commit.title,
"comments": commit.comments,
"diff": item.diff,
"content_after": content_after,
"create_time": commit.create_time,
"create_date": commit.create_date,
}
@classmethod
@DB.connection_context()
def get_file_version_history(cls, file_id):
@@ -545,20 +913,21 @@ class FileCommitService(CommonService):
Returns list of dicts: [{"commit_id", "operation", "hash", "create_time", "message"}]
"""
items = FileCommitItem.select().where(FileCommitItem.file_id == file_id).order_by(
FileCommitItem.create_time.desc())
items = FileCommitItem.select().where(FileCommitItem.file_id == file_id).order_by(FileCommitItem.create_time.desc())
versions = []
for item in items:
commit = cls.get_commit(item.commit_id)
if commit:
versions.append({
"commit_id": item.commit_id,
"operation": item.operation,
"hash": item.new_hash or item.old_hash or "",
"create_time": item.create_time,
"message": commit.message,
})
versions.append(
{
"commit_id": item.commit_id,
"operation": item.operation,
"hash": item.new_hash or item.old_hash or "",
"create_time": item.create_time,
"message": commit.message,
}
)
return versions
@@ -620,9 +989,7 @@ def _build_hierarchical_tree(tree_state, root_folder_id):
}
# File children
for fid, entry in files_by_parent.get(node_id, []):
fn = {"id": fid, "name": entry.get("name", fid), "type": "file",
"hash": entry.get("hash", ""), "size": entry.get("size", 0),
"status": entry.get("status", "1")}
fn = {"id": fid, "name": entry.get("name", fid), "type": "file", "hash": entry.get("hash", ""), "size": entry.get("size", 0), "status": entry.get("status", "1")}
if entry.get("location"):
fn["location"] = entry["location"]
node["children"].append(fn)

110
api/db/services/knowledgebase_service.py
View File

@@ -46,6 +46,7 @@ class KnowledgebaseService(CommonService):
Attributes:
model: The Knowledgebase model class for database operations.
"""
model = Knowledgebase
@classmethod
@@ -59,13 +60,7 @@ class KnowledgebaseService(CommonService):
- KBs owned by the current user (`tenant_id == user_id`)
Always constrained to `StatusEnum.VALID`.
"""
return (
(
(cls.model.tenant_id.in_(joined_tenant_ids) & (cls.model.permission == TenantPermission.TEAM.value))
| (cls.model.tenant_id == user_id)
)
& (cls.model.status == StatusEnum.VALID.value)
)
return ((cls.model.tenant_id.in_(joined_tenant_ids) & (cls.model.permission == TenantPermission.TEAM.value)) | (cls.model.tenant_id == user_id)) & (cls.model.status == StatusEnum.VALID.value)
@classmethod
@DB.connection_context()
@@ -94,8 +89,7 @@ class KnowledgebaseService(CommonService):
2. The user is not the creator of the dataset
"""
# Check if a dataset can be deleted by a user
docs = cls.model.select(
cls.model.id).where(cls.model.id == kb_id, cls.model.created_by == user_id).paginate(0, 1)
docs = cls.model.select(cls.model.id).where(cls.model.id == kb_id, cls.model.created_by == user_id).paginate(0, 1)
docs = docs.dicts()
if not docs:
return False
@@ -127,10 +121,10 @@ class KnowledgebaseService(CommonService):
# Check parsing status of each document
for doc in docs:
# If document is being parsed, don't allow chat creation
if doc['run'] == TaskStatus.RUNNING.value or doc['run'] == TaskStatus.CANCEL.value or doc['run'] == TaskStatus.FAIL.value:
if doc["run"] == TaskStatus.RUNNING.value or doc["run"] == TaskStatus.CANCEL.value or doc["run"] == TaskStatus.FAIL.value:
return False, f"Document '{doc['name']}' in dataset '{kb.name}' is still being parsed. Please wait until all documents are parsed before starting a chat."
# If document is not yet parsed and has no chunks, don't allow chat creation
if doc['run'] == TaskStatus.UNSTART.value and doc['chunk_num'] == 0:
if doc["run"] == TaskStatus.UNSTART.value and doc["chunk_num"] == 0:
return False, f"Document '{doc['name']}' in dataset '{kb.name}' has not been parsed yet. Please parse all documents before starting a chat."
return True, None
@@ -143,20 +137,14 @@ class KnowledgebaseService(CommonService):
# kb_ids: List of dataset IDs
# Returns:
# List of document IDs
doc_ids = cls.model.select(Document.id.alias("document_id")).join(Document, on=(cls.model.id == Document.kb_id)).where(
cls.model.id.in_(kb_ids)
)
doc_ids = cls.model.select(Document.id.alias("document_id")).join(Document, on=(cls.model.id == Document.kb_id)).where(cls.model.id.in_(kb_ids))
doc_ids = list(doc_ids.dicts())
doc_ids = [doc["document_id"] for doc in doc_ids]
return doc_ids
@classmethod
@DB.connection_context()
def get_by_tenant_ids(cls, joined_tenant_ids, user_id,
page_number, items_per_page,
orderby, desc, keywords,
parser_id=None
):
def get_by_tenant_ids(cls, joined_tenant_ids, user_id, page_number, items_per_page, orderby, desc, keywords, parser_id=None):
# Get knowledge bases by tenant IDs with pagination and filtering
# Args:
# joined_tenant_ids: List of tenant IDs
@@ -183,17 +171,25 @@ class KnowledgebaseService(CommonService):
cls.model.parser_id,
cls.model.embd_id,
User.nickname,
User.avatar.alias('tenant_avatar'),
cls.model.update_time
User.avatar.alias("tenant_avatar"),
cls.model.update_time,
]
if keywords:
kbs = cls.model.select(*fields).join(User, on=(cls.model.tenant_id == User.id)).where(
cls._visibility_and_status_filter(joined_tenant_ids, user_id),
fn.LOWER(cls.model.name).contains(keywords.lower()),
kbs = (
cls.model.select(*fields)
.join(User, on=(cls.model.tenant_id == User.id))
.where(
cls._visibility_and_status_filter(joined_tenant_ids, user_id),
fn.LOWER(cls.model.name).contains(keywords.lower()),
)
)
else:
kbs = cls.model.select(*fields).join(User, on=(cls.model.tenant_id == User.id)).where(
cls._visibility_and_status_filter(joined_tenant_ids, user_id),
kbs = (
cls.model.select(*fields)
.join(User, on=(cls.model.tenant_id == User.id))
.where(
cls._visibility_and_status_filter(joined_tenant_ids, user_id),
)
)
if parser_id:
kbs = kbs.where(cls.model.parser_id == parser_id)
@@ -223,7 +219,7 @@ class KnowledgebaseService(CommonService):
cls.model.chunk_num,
cls.model.status,
cls.model.create_date,
cls.model.update_date
cls.model.update_date,
]
# find team kb and owned kb
kbs = cls.model.select(*fields).where(cls._visibility_and_status_filter(tenant_ids, user_id))
@@ -287,15 +283,19 @@ class KnowledgebaseService(CommonService):
cls.model.raptor_task_finish_at,
cls.model.mindmap_task_id,
cls.model.mindmap_task_finish_at,
cls.model.artifact_task_id,
cls.model.artifact_task_finish_at,
cls.model.skill_task_id,
cls.model.skill_task_finish_at,
cls.model.create_time,
cls.model.update_time
]
kbs = cls.model.select(*fields)\
.join(UserCanvas, on=(cls.model.pipeline_id == UserCanvas.id), join_type=JOIN.LEFT_OUTER)\
.where(
(cls.model.id == kb_id),
(cls.model.status == StatusEnum.VALID.value)
).dicts()
cls.model.update_time,
]
kbs = (
cls.model.select(*fields)
.join(UserCanvas, on=(cls.model.pipeline_id == UserCanvas.id), join_type=JOIN.LEFT_OUTER)
.where((cls.model.id == kb_id), (cls.model.status == StatusEnum.VALID.value))
.dicts()
)
if not kbs:
return None
return kbs[0]
@@ -362,11 +362,7 @@ class KnowledgebaseService(CommonService):
# tenant_id: Tenant ID
# Returns:
# Tuple of (exists, knowledge_base)
kb = cls.model.select().where(
(cls.model.name == kb_name)
& (cls.model.tenant_id == tenant_id)
& (cls.model.status == StatusEnum.VALID.value)
)
kb = cls.model.select().where((cls.model.name == kb_name) & (cls.model.tenant_id == tenant_id) & (cls.model.status == StatusEnum.VALID.value))
if kb:
return True, kb[0]
return False, None
@@ -379,17 +375,9 @@ class KnowledgebaseService(CommonService):
# List of all dataset IDs
return [m["id"] for m in cls.model.select(cls.model.id).dicts()]
@classmethod
@DB.connection_context()
def create_with_name(
cls,
*,
name: str,
tenant_id: str,
parser_id: str | None = None,
**kwargs
):
def create_with_name(cls, *, name: str, tenant_id: str, parser_id: str | None = None, **kwargs):
"""Create a dataset (knowledgebase) by name with kb_app defaults.
This encapsulates the creation logic used in kb_app.create so other callers
@@ -429,7 +417,7 @@ class KnowledgebaseService(CommonService):
"tenant_id": tenant_id,
"created_by": tenant_id,
"parser_id": (parser_id or "naive"),
**kwargs # Includes optional fields such as description, language, permission, avatar, parser_config, etc.
**kwargs, # Includes optional fields such as description, language, permission, avatar, parser_config, etc.
}
# Update parser_config (always override with validated default/merged config)
@@ -438,11 +426,9 @@ class KnowledgebaseService(CommonService):
return True, payload
@classmethod
@DB.connection_context()
def get_list(cls, joined_tenant_ids, user_id,
page_number, items_per_page, orderby, desc, id, name, keywords, parser_id=None):
def get_list(cls, joined_tenant_ids, user_id, page_number, items_per_page, orderby, desc, id, name, keywords, parser_id=None):
# Get list of knowledge bases with filtering and pagination
# Args:
# joined_tenant_ids: List of tenant IDs
@@ -566,7 +552,7 @@ class KnowledgebaseService(CommonService):
kb.save(only=dirty_fields)
except ValueError as e:
if str(e) == "no data to save!":
pass # that's OK
pass # that's OK
else:
raise e
@@ -577,21 +563,17 @@ class KnowledgebaseService(CommonService):
if not kb_row:
raise RuntimeError(f"kb_id {kb_id} does not exist")
update_dict = {
'doc_num': kb_row.doc_num - doc_num_info['doc_num'],
'chunk_num': kb_row.chunk_num - doc_num_info['chunk_num'],
'token_num': kb_row.token_num - doc_num_info['token_num'],
'update_time': current_timestamp(),
'update_date': datetime_format(datetime.now())
"doc_num": kb_row.doc_num - doc_num_info["doc_num"],
"chunk_num": kb_row.chunk_num - doc_num_info["chunk_num"],
"token_num": kb_row.token_num - doc_num_info["token_num"],
"update_time": current_timestamp(),
"update_date": datetime_format(datetime.now()),
}
return cls.model.update(update_dict).where(cls.model.id == kb_id).execute()
@classmethod
@DB.connection_context()
def get_null_tenant_embd_id_row(cls):
fields = [
cls.model.id,
cls.model.tenant_id,
cls.model.embd_id
]
fields = [cls.model.id, cls.model.tenant_id, cls.model.embd_id]
objs = cls.model.select(*fields).where(cls.model.tenant_embd_id.is_null())
return list(objs)

16
api/db/services/pipeline_operation_log_service.py
View File

@@ -98,8 +98,8 @@ class PipelineOperationLogService(CommonService):
if document_id != GRAPH_RAPTOR_FAKE_DOC_ID:
referred_document_id = document_id
# no need to update document for graph rag, raptor mindmap task
if task_type not in [PipelineTaskType.GRAPH_RAG, PipelineTaskType.RAPTOR, PipelineTaskType.MINDMAP]:
# no need to update document for KB-level fan-out tasks
if task_type not in [PipelineTaskType.GRAPH_RAG, PipelineTaskType.RAPTOR, PipelineTaskType.MINDMAP, PipelineTaskType.ARTIFACT, PipelineTaskType.SKILL]:
ok, document = DocumentService.get_by_id(referred_document_id)
if not ok:
logging.warning(f"Document for referred_document_id {referred_document_id} not found")
@@ -137,7 +137,7 @@ class PipelineOperationLogService(CommonService):
if task_type not in VALID_PIPELINE_TASK_TYPES:
raise ValueError(f"Invalid task type: {task_type}")
if task_type in [PipelineTaskType.GRAPH_RAG, PipelineTaskType.RAPTOR, PipelineTaskType.MINDMAP]:
if task_type in [PipelineTaskType.GRAPH_RAG, PipelineTaskType.RAPTOR, PipelineTaskType.MINDMAP, PipelineTaskType.ARTIFACT, PipelineTaskType.SKILL]:
# query task to get progress information from task
ok, task = TaskService.get_by_id(task_id)
if not ok:
@@ -166,6 +166,16 @@ class PipelineOperationLogService(CommonService):
document.kb_id,
{"mindmap_task_finish_at": finish_at},
)
elif task_type == PipelineTaskType.ARTIFACT:
KnowledgebaseService.update_by_id(
document.kb_id,
{"artifact_task_finish_at": finish_at},
)
elif task_type == PipelineTaskType.SKILL:
KnowledgebaseService.update_by_id(
document.kb_id,
{"skill_task_finish_at": finish_at},
)
log = dict(
id=get_uuid(),

188
api/db/services/task_service.py
View File

@@ -37,7 +37,95 @@ from rag.nlp import search
CANVAS_DEBUG_DOC_ID = "dataflow_x"
GRAPH_RAPTOR_FAKE_DOC_ID = "graph_raptor_x"
TASK_MAX_LOG_LENGTH = int(os.environ.get("TASK_MAX_LOG_LENGTH", 3000)) # TEXT MAX is 64 KiB bytes!
TASK_MAX_LOG_LENGTH = int(os.environ.get("TASK_MAX_LOG_LENGTH", 3000)) # TEXT MAX is 64 KiB bytes!
DOC_CHUNKING_COUNTER_TTL_SECONDS = 7 * 24 * 3600
def _doc_chunking_pending_key(doc_id: str) -> str:
return f"doc:chunking_pending:{doc_id}"
def _doc_chunking_aborted_key(doc_id: str) -> str:
return f"doc:chunking_aborted:{doc_id}"
def _doc_chunking_done_key(task_id: str) -> str:
return f"doc:chunking_done:{task_id}"
def seed_doc_chunking_counter(doc_id: str, pending_count: int) -> bool:
if not doc_id or pending_count <= 0:
return False
try:
REDIS_CONN.delete(_doc_chunking_aborted_key(doc_id))
return REDIS_CONN.set(
_doc_chunking_pending_key(doc_id),
str(pending_count),
exp=DOC_CHUNKING_COUNTER_TTL_SECONDS,
)
except Exception:
logging.exception("Failed to seed chunking counter for doc %s", doc_id)
return False
def clear_doc_chunking_counter(doc_id: str) -> None:
if not doc_id:
return
try:
REDIS_CONN.delete(_doc_chunking_pending_key(doc_id))
except Exception:
logging.exception("Failed to clear chunking counter for doc %s", doc_id)
def abort_doc_chunking_counter(doc_id: str) -> None:
if not doc_id:
return
try:
REDIS_CONN.delete(_doc_chunking_pending_key(doc_id))
REDIS_CONN.set(
_doc_chunking_aborted_key(doc_id),
"1",
exp=DOC_CHUNKING_COUNTER_TTL_SECONDS,
)
except Exception:
logging.exception("Failed to abort chunking counter for doc %s", doc_id)
def is_doc_chunking_aborted(doc_id: str) -> bool:
if not doc_id:
return False
try:
return bool(REDIS_CONN.get(_doc_chunking_aborted_key(doc_id)))
except Exception:
logging.exception("Failed to read chunking abort marker for doc %s", doc_id)
return False
def credit_doc_chunking_task(doc_id: str, task_id: str) -> int | None:
"""Credit one completed standard chunking task.
Returns the post-decrement pending count when this task was credited for
the first time. Returns a positive value when this task was already
credited, so callers treat retries as not-last.
"""
if not doc_id or not task_id:
return None
try:
first_credit = REDIS_CONN.set_if_absent(
_doc_chunking_done_key(task_id),
"1",
exp=DOC_CHUNKING_COUNTER_TTL_SECONDS,
)
if not first_credit:
return 1
pending_key = _doc_chunking_pending_key(doc_id)
if REDIS_CONN.get(pending_key) is None:
return -1
return REDIS_CONN.decrby(pending_key, 1)
except Exception:
logging.exception("Failed to credit chunking task %s for doc %s", task_id, doc_id)
return None
def trim_header_by_lines(text: str, max_length) -> str:
# Trim header text to maximum length while preserving line breaks
@@ -50,8 +138,8 @@ def trim_header_by_lines(text: str, max_length) -> str:
if len_text <= max_length:
return text
for i in range(len_text):
if text[i] == '\n' and len_text - i <= max_length:
return text[i + 1:]
if text[i] == "\n" and len_text - i <= max_length:
return text[i + 1 :]
return text
@@ -69,6 +157,7 @@ class TaskService(CommonService):
Attributes:
model: The Task model class for database operations.
"""
model = Task
@classmethod
@@ -96,6 +185,7 @@ class TaskService(CommonService):
cls.model.doc_id,
cls.model.from_page,
cls.model.to_page,
cls.model.task_type,
cls.model.retry_count,
Document.kb_id,
Document.parser_id,
@@ -116,32 +206,34 @@ class TaskService(CommonService):
]
docs = (
cls.model.select(*fields)
.join(Document, on=(doc_id == Document.id))
.join(Knowledgebase, on=(Document.kb_id == Knowledgebase.id))
.join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))
.where(cls.model.id == task_id)
.join(Document, on=(doc_id == Document.id))
.join(Knowledgebase, on=(Document.kb_id == Knowledgebase.id))
.join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))
.where(cls.model.id == task_id)
)
docs = list(docs.dicts())
if not docs:
return None
doc = docs[0]
msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task has been received."
prog = random.random() / 10.0
if docs[0]["retry_count"] >= 3:
if doc["retry_count"] >= 3:
msg = "\nERROR: Task is abandoned after 3 times attempts."
prog = -1
cls.model.update(
progress_msg=cls.model.progress_msg + msg,
progress=prog,
retry_count=docs[0]["retry_count"] + 1,
).where(cls.model.id == docs[0]["id"]).execute()
retry_count=doc["retry_count"] + 1,
).where(cls.model.id == doc["id"]).execute()
if docs[0]["retry_count"] >= 3:
abort_doc_chunking_counter(docs[0]["doc_id"])
DocumentService.update_by_id(docs[0]["doc_id"], {"progress": -1, "run": TaskStatus.FAIL.value, "update_time": current_timestamp(), "update_date": get_format_time()})
return None
return docs[0]
return doc
@classmethod
@DB.connection_context()
@@ -165,10 +257,7 @@ class TaskService(CommonService):
cls.model.digest,
cls.model.chunk_ids,
]
tasks = (
cls.model.select(*fields).order_by(cls.model.from_page.asc(), cls.model.create_time.desc())
.where(cls.model.doc_id == doc_id)
)
tasks = cls.model.select(*fields).order_by(cls.model.from_page.asc(), cls.model.create_time.desc()).where(cls.model.doc_id == doc_id)
tasks = list(tasks.dicts())
if not tasks:
return None
@@ -189,20 +278,8 @@ class TaskService(CommonService):
list[dict]: List of task dictionaries containing task details.
Returns None if no tasks are found.
"""
fields = [
cls.model.id,
cls.model.doc_id,
cls.model.from_page,
cls.model.progress,
cls.model.progress_msg,
cls.model.digest,
cls.model.chunk_ids,
cls.model.create_time
]
tasks = (
cls.model.select(*fields).order_by(cls.model.create_time.desc())
.where(cls.model.doc_id.in_(doc_ids))
)
fields = [cls.model.id, cls.model.doc_id, cls.model.from_page, cls.model.progress, cls.model.progress_msg, cls.model.digest, cls.model.chunk_ids, cls.model.create_time]
tasks = cls.model.select(*fields).order_by(cls.model.create_time.desc()).where(cls.model.doc_id.in_(doc_ids))
tasks = list(tasks.dicts())
if not tasks:
return None
@@ -238,9 +315,7 @@ class TaskService(CommonService):
"""
with DB.lock("get_task", -1):
docs = (
cls.model.select(
*[Document.id, Document.kb_id, Document.location, File.parent_id]
)
cls.model.select(*[Document.id, Document.kb_id, Document.location, File.parent_id])
.join(Document, on=(cls.model.doc_id == Document.id))
.join(
File2Document,
@@ -326,11 +401,7 @@ class TaskService(CommonService):
cls.model.update(progress_msg=progress_msg).where(cls.model.id == id).execute()
if "progress" in info:
prog = info["progress"]
cls.model.update(progress=prog).where(
(cls.model.id == id) &
((prog >= 1) | ((cls.model.progress != -1) &
((prog == -1) | (prog > cls.model.progress))))
).execute()
cls.model.update(progress=prog).where((cls.model.id == id) & ((prog >= 1) | ((cls.model.progress != -1) & ((prog == -1) | (prog > cls.model.progress))))).execute()
else:
with DB.lock("update_progress", -1):
if info["progress_msg"]:
@@ -338,11 +409,7 @@ class TaskService(CommonService):
cls.model.update(progress_msg=progress_msg).where(cls.model.id == id).execute()
if "progress" in info:
prog = info["progress"]
cls.model.update(progress=prog).where(
(cls.model.id == id) &
((prog >= 1) | ((cls.model.progress != -1) &
((prog == -1) | (prog > cls.model.progress))))
).execute()
cls.model.update(progress=prog).where((cls.model.id == id) & ((prog >= 1) | ((cls.model.progress != -1) & ((prog == -1) | (prog > cls.model.progress))))).execute()
begin_at = task.begin_at
if begin_at is not None:
@@ -456,18 +523,22 @@ def queue_tasks(doc: dict, bucket: str, name: str, priority: int):
if pre_task["chunk_ids"]:
pre_chunk_ids.extend(pre_task["chunk_ids"].split())
if pre_chunk_ids:
settings.docStoreConn.delete({"id": pre_chunk_ids}, search.index_name(chunking_config["tenant_id"]),
chunking_config["kb_id"])
settings.docStoreConn.delete({"id": pre_chunk_ids}, search.index_name(chunking_config["tenant_id"]), chunking_config["kb_id"])
DocumentService.update_by_id(doc["id"], {"chunk_num": ck_num})
bulk_insert_into_db(Task, parse_task_array, True)
DocumentService.begin2parse(doc["id"])
unfinished_task_array = [task for task in parse_task_array if task["progress"] < 1.0]
for unfinished_task in unfinished_task_array:
assert REDIS_CONN.queue_product(
settings.get_svr_queue_name(priority, suffix), message=unfinished_task
), "Can't access Redis. Please check the Redis' status."
chunking_n = sum(1 for task in unfinished_task_array if not task.get("task_type"))
if chunking_n > 0:
assert seed_doc_chunking_counter(doc["id"], chunking_n), "Can't access Redis. Please check the Redis' status."
try:
for unfinished_task in unfinished_task_array:
assert REDIS_CONN.queue_product(settings.get_svr_queue_name(priority, suffix), message=unfinished_task), "Can't access Redis. Please check the Redis' status."
except Exception:
abort_doc_chunking_counter(doc["id"])
raise
def reuse_prev_task_chunks(task: dict, prev_tasks: list[dict], chunking_config: dict):
@@ -494,8 +565,7 @@ def reuse_prev_task_chunks(task: dict, prev_tasks: list[dict], chunking_config:
idx = 0
while idx < len(prev_tasks):
prev_task = prev_tasks[idx]
if prev_task.get("from_page", 0) == task.get("from_page", 0) \
and prev_task.get("digest", 0) == task.get("digest", ""):
if prev_task.get("from_page", 0) == task.get("from_page", 0) and prev_task.get("digest", 0) == task.get("digest", ""):
break
idx += 1
@@ -506,18 +576,22 @@ def reuse_prev_task_chunks(task: dict, prev_tasks: list[dict], chunking_config:
return 0
task["chunk_ids"] = prev_task["chunk_ids"]
task["progress"] = 1.0
if "from_page" in task and "to_page" in task and (int(task['to_page']) - int(task['from_page']) >= 10 ** 6 or (int(task['from_page']) == MAXIMUM_TASK_PAGE_NUMBER and int(task['to_page']) == MAXIMUM_TASK_PAGE_NUMBER)):
if (
"from_page" in task
and "to_page" in task
and (int(task["to_page"]) - int(task["from_page"]) >= 10**6 or (int(task["from_page"]) == MAXIMUM_TASK_PAGE_NUMBER and int(task["to_page"]) == MAXIMUM_TASK_PAGE_NUMBER))
):
task["progress_msg"] = f"Page({task['from_page']}~{task['to_page']}): "
else:
task["progress_msg"] = ""
task["progress_msg"] = " ".join(
[datetime.now().strftime("%H:%M:%S"), task["progress_msg"], "Reused previous task's chunks."])
task["progress_msg"] = " ".join([datetime.now().strftime("%H:%M:%S"), task["progress_msg"], "Reused previous task's chunks."])
prev_task["chunk_ids"] = ""
return len(task["chunk_ids"].split())
def cancel_all_task_of(doc_id):
abort_doc_chunking_counter(doc_id)
for t in TaskService.query(doc_id=doc_id):
try:
REDIS_CONN.set(f"{t.id}-cancel", "x")
@@ -535,7 +609,7 @@ def has_canceled(task_id):
return False
def queue_dataflow(tenant_id:str, flow_id:str, task_id:str, doc_id:str=CANVAS_DEBUG_DOC_ID, file:dict=None, priority: int=0, rerun:bool=False) -> tuple[bool, str]:
def queue_dataflow(tenant_id: str, flow_id: str, task_id: str, doc_id: str = CANVAS_DEBUG_DOC_ID, file: dict = None, priority: int = 0, rerun: bool = False) -> tuple[bool, str]:
task = dict(
id=task_id,
@@ -544,7 +618,7 @@ def queue_dataflow(tenant_id:str, flow_id:str, task_id:str, doc_id:str=CANVAS_DE
to_page=MAXIMUM_TASK_PAGE_NUMBER,
task_type="dataflow" if not rerun else "dataflow_rerun",
priority=priority,
begin_at= datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
begin_at=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
)
if doc_id not in [CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID]:
TaskService.model.delete().where(TaskService.model.doc_id == doc_id).execute()
@@ -556,9 +630,7 @@ def queue_dataflow(tenant_id:str, flow_id:str, task_id:str, doc_id:str=CANVAS_DE
task["dataflow_id"] = flow_id
task["file"] = file
if not REDIS_CONN.queue_product(
settings.get_svr_queue_name(priority, "common"), message=task
):
if not REDIS_CONN.queue_product(settings.get_svr_queue_name(priority, "common"), message=task):
return False, "Can't access Redis. Please check the Redis' status."
return True, ""