zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-17 05:07:23 +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

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, ""