mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 20:42:30 +08:00
Add task moduel, and pipline the task and every parser (#49)
This commit is contained in:
KevinHuSh
@ -22,6 +22,8 @@ from elasticsearch_dsl import Q
|
||||
from flask import request
|
||||
from flask_login import login_required, current_user
|
||||
|
||||
from api.db.db_models import Task
|
||||
from api.db.services.task_service import TaskService
|
||||
from rag.nlp import search
|
||||
from rag.utils import ELASTICSEARCH
|
||||
from api.db.services import duplicate_name
|
||||
@ -205,6 +207,26 @@ def rm():
|
||||
return server_error_response(e)
|
||||
|
||||
|
||||
@manager.route('/run', methods=['POST'])
|
||||
@login_required
|
||||
@validate_request("doc_ids", "run")
|
||||
def rm():
|
||||
req = request.json
|
||||
try:
|
||||
for id in req["doc_ids"]:
|
||||
DocumentService.update_by_id(id, {"run": str(req["run"])})
|
||||
if req["run"] == "2":
|
||||
TaskService.filter_delete([Task.doc_id == id])
|
||||
tenant_id = DocumentService.get_tenant_id(id)
|
||||
if not tenant_id:
|
||||
return get_data_error_result(retmsg="Tenant not found!")
|
||||
ELASTICSEARCH.deleteByQuery(Q("match", doc_id=id), idxnm=search.index_name(tenant_id))
|
||||
|
||||
return get_json_result(data=True)
|
||||
except Exception as e:
|
||||
return server_error_response(e)
|
||||
|
||||
|
||||
@manager.route('/rename', methods=['POST'])
|
||||
@login_required
|
||||
@validate_request("doc_id", "name", "old_name")
|
||||
@ -262,7 +284,7 @@ def change_parser():
|
||||
if doc.parser_id.lower() == req["parser_id"].lower():
|
||||
return get_json_result(data=True)
|
||||
|
||||
e = DocumentService.update_by_id(doc.id, {"parser_id": req["parser_id"], "progress":0, "progress_msg": ""})
|
||||
e = DocumentService.update_by_id(doc.id, {"parser_id": req["parser_id"], "progress":0, "progress_msg": "", "run": 1})
|
||||
if not e:
|
||||
return get_data_error_result(retmsg="Document not found!")
|
||||
e = DocumentService.increment_chunk_num(doc.id, doc.kb_id, doc.token_num*-1, doc.chunk_num*-1, doc.process_duation*-1)
|
||||
|
||||
@ -59,3 +59,14 @@ class ChatStyle(StrEnum):
|
||||
PRECISE = 'Precise'
|
||||
EVENLY = 'Evenly'
|
||||
CUSTOM = 'Custom'
|
||||
|
||||
|
||||
class ParserType(StrEnum):
|
||||
GENERAL = "general"
|
||||
PRESENTATION = "presentation"
|
||||
LAWS = "laws"
|
||||
MANUAL = "manual"
|
||||
PAPER = "paper"
|
||||
RESUME = ""
|
||||
BOOK = ""
|
||||
QA = ""
|
||||
|
||||
@ -496,15 +496,27 @@ class Document(DataBaseModel):
|
||||
token_num = IntegerField(default=0)
|
||||
chunk_num = IntegerField(default=0)
|
||||
progress = FloatField(default=0)
|
||||
progress_msg = CharField(max_length=255, null=True, help_text="process message", default="")
|
||||
progress_msg = CharField(max_length=512, null=True, help_text="process message", default="")
|
||||
process_begin_at = DateTimeField(null=True)
|
||||
process_duation = FloatField(default=0)
|
||||
run = CharField(max_length=1, null=True, help_text="start to run processing or cancel.(1: run it; 2: cancel)", default="0")
|
||||
status = CharField(max_length=1, null=True, help_text="is it validate(0: wasted，1: validate)", default="1")
|
||||
|
||||
class Meta:
|
||||
db_table = "document"
|
||||
|
||||
|
||||
class Task(DataBaseModel):
|
||||
id = CharField(max_length=32, primary_key=True)
|
||||
doc_id = CharField(max_length=32, null=False, index=True)
|
||||
from_page = IntegerField(default=0)
|
||||
to_page = IntegerField(default=-1)
|
||||
begin_at = DateTimeField(null=True)
|
||||
process_duation = FloatField(default=0)
|
||||
progress = FloatField(default=0)
|
||||
progress_msg = CharField(max_length=255, null=True, help_text="process message", default="")
|
||||
|
||||
|
||||
class Dialog(DataBaseModel):
|
||||
id = CharField(max_length=32, primary_key=True)
|
||||
tenant_id = CharField(max_length=32, null=False)
|
||||
@ -553,72 +565,6 @@ class Conversation(DataBaseModel):
|
||||
|
||||
|
||||
"""
|
||||
class Job(DataBaseModel):
|
||||
# multi-party common configuration
|
||||
f_user_id = CharField(max_length=25, null=True)
|
||||
f_job_id = CharField(max_length=25, index=True)
|
||||
f_name = CharField(max_length=500, null=True, default='')
|
||||
f_description = TextField(null=True, default='')
|
||||
f_tag = CharField(max_length=50, null=True, default='')
|
||||
f_dsl = JSONField()
|
||||
f_runtime_conf = JSONField()
|
||||
f_runtime_conf_on_party = JSONField()
|
||||
f_train_runtime_conf = JSONField(null=True)
|
||||
f_roles = JSONField()
|
||||
f_initiator_role = CharField(max_length=50)
|
||||
f_initiator_party_id = CharField(max_length=50)
|
||||
f_status = CharField(max_length=50)
|
||||
f_status_code = IntegerField(null=True)
|
||||
f_user = JSONField()
|
||||
# this party configuration
|
||||
f_role = CharField(max_length=50, index=True)
|
||||
f_party_id = CharField(max_length=10, index=True)
|
||||
f_is_initiator = BooleanField(null=True, default=False)
|
||||
f_progress = IntegerField(null=True, default=0)
|
||||
f_ready_signal = BooleanField(default=False)
|
||||
f_ready_time = BigIntegerField(null=True)
|
||||
f_cancel_signal = BooleanField(default=False)
|
||||
f_cancel_time = BigIntegerField(null=True)
|
||||
f_rerun_signal = BooleanField(default=False)
|
||||
f_end_scheduling_updates = IntegerField(null=True, default=0)
|
||||
|
||||
f_engine_name = CharField(max_length=50, null=True)
|
||||
f_engine_type = CharField(max_length=10, null=True)
|
||||
f_cores = IntegerField(default=0)
|
||||
f_memory = IntegerField(default=0) # MB
|
||||
f_remaining_cores = IntegerField(default=0)
|
||||
f_remaining_memory = IntegerField(default=0) # MB
|
||||
f_resource_in_use = BooleanField(default=False)
|
||||
f_apply_resource_time = BigIntegerField(null=True)
|
||||
f_return_resource_time = BigIntegerField(null=True)
|
||||
|
||||
f_inheritance_info = JSONField(null=True)
|
||||
f_inheritance_status = CharField(max_length=50, null=True)
|
||||
|
||||
f_start_time = BigIntegerField(null=True)
|
||||
f_start_date = DateTimeField(null=True)
|
||||
f_end_time = BigIntegerField(null=True)
|
||||
f_end_date = DateTimeField(null=True)
|
||||
f_elapsed = BigIntegerField(null=True)
|
||||
|
||||
class Meta:
|
||||
db_table = "t_job"
|
||||
primary_key = CompositeKey('f_job_id', 'f_role', 'f_party_id')
|
||||
|
||||
|
||||
|
||||
class PipelineComponentMeta(DataBaseModel):
|
||||
f_model_id = CharField(max_length=100, index=True)
|
||||
f_model_version = CharField(max_length=100, index=True)
|
||||
f_role = CharField(max_length=50, index=True)
|
||||
f_party_id = CharField(max_length=10, index=True)
|
||||
f_component_name = CharField(max_length=100, index=True)
|
||||
f_component_module_name = CharField(max_length=100)
|
||||
f_model_alias = CharField(max_length=100, index=True)
|
||||
f_model_proto_index = JSONField(null=True)
|
||||
f_run_parameters = JSONField(null=True)
|
||||
f_archive_sha256 = CharField(max_length=100, null=True)
|
||||
f_archive_from_ip = CharField(max_length=100, null=True)
|
||||
|
||||
class Meta:
|
||||
db_table = 't_pipeline_component_meta'
|
||||
|
||||
@ -32,19 +32,19 @@ LOGGER = getLogger()
|
||||
def bulk_insert_into_db(model, data_source, replace_on_conflict=False):
|
||||
DB.create_tables([model])
|
||||
|
||||
current_time = current_timestamp()
|
||||
current_date = timestamp_to_date(current_time)
|
||||
|
||||
for data in data_source:
|
||||
if 'f_create_time' not in data:
|
||||
data['f_create_time'] = current_time
|
||||
data['f_create_date'] = timestamp_to_date(data['f_create_time'])
|
||||
data['f_update_time'] = current_time
|
||||
data['f_update_date'] = current_date
|
||||
current_time = current_timestamp()
|
||||
current_date = timestamp_to_date(current_time)
|
||||
if 'create_time' not in data:
|
||||
data['create_time'] = current_time
|
||||
data['create_date'] = timestamp_to_date(data['create_time'])
|
||||
data['update_time'] = current_time
|
||||
data['update_date'] = current_date
|
||||
|
||||
preserve = tuple(data_source[0].keys() - {'f_create_time', 'f_create_date'})
|
||||
preserve = tuple(data_source[0].keys() - {'create_time', 'create_date'})
|
||||
|
||||
batch_size = 50 if RuntimeConfig.USE_LOCAL_DATABASE else 1000
|
||||
batch_size = 1000
|
||||
|
||||
for i in range(0, len(data_source), batch_size):
|
||||
with DB.atomic():
|
||||
|
||||
@ -70,6 +70,7 @@ class CommonService:
|
||||
@DB.connection_context()
|
||||
def insert_many(cls, data_list, batch_size=100):
|
||||
with DB.atomic():
|
||||
for d in data_list: d["create_time"] = datetime_format(datetime.now())
|
||||
for i in range(0, len(data_list), batch_size):
|
||||
cls.model.insert_many(data_list[i:i + batch_size]).execute()
|
||||
|
||||
|
||||
@ -61,8 +61,8 @@ class DocumentService(CommonService):
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_newly_uploaded(cls, tm, mod, comm, items_per_page=64):
|
||||
fields = [cls.model.id, cls.model.kb_id, cls.model.parser_id, cls.model.name, cls.model.location, cls.model.size, Knowledgebase.tenant_id, Tenant.embd_id, Tenant.img2txt_id, cls.model.update_time]
|
||||
def get_newly_uploaded(cls, tm, mod=0, comm=1, items_per_page=64):
|
||||
fields = [cls.model.id, cls.model.kb_id, cls.model.parser_id, cls.model.name, cls.model.type, cls.model.location, cls.model.size, Knowledgebase.tenant_id, Tenant.embd_id, Tenant.img2txt_id, Tenant.asr_id, cls.model.update_time]
|
||||
docs = cls.model.select(*fields) \
|
||||
.join(Knowledgebase, on=(cls.model.kb_id == Knowledgebase.id)) \
|
||||
.join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))\
|
||||
@ -76,6 +76,18 @@ class DocumentService(CommonService):
|
||||
.paginate(1, items_per_page)
|
||||
return list(docs.dicts())
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_unfinished_docs(cls):
|
||||
fields = [cls.model.id, cls.model.process_begin_at]
|
||||
docs = cls.model.select(*fields) \
|
||||
.where(
|
||||
cls.model.status == StatusEnum.VALID.value,
|
||||
~(cls.model.type == FileType.VIRTUAL.value),
|
||||
cls.model.progress < 1,
|
||||
cls.model.progress > 0)
|
||||
return list(docs.dicts())
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def increment_chunk_num(cls, doc_id, kb_id, token_num, chunk_num, duation):
|
||||
|
||||
53
api/db/services/task_service.py
Normal file
53
api/db/services/task_service.py
Normal file
@ -0,0 +1,53 @@
|
||||
#
|
||||
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
from peewee import Expression
|
||||
from api.db.db_models import DB
|
||||
from api.db import StatusEnum, FileType
|
||||
from api.db.db_models import Task, Document, Knowledgebase, Tenant
|
||||
from api.db.services.common_service import CommonService
|
||||
|
||||
|
||||
class TaskService(CommonService):
|
||||
model = Task
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def get_tasks(cls, tm, mod=0, comm=1, items_per_page=64):
|
||||
fields = [cls.model.id, cls.model.doc_id, cls.model.from_page,cls.model.to_page, Document.kb_id, Document.parser_id, Document.name, Document.type, Document.location, Document.size, Knowledgebase.tenant_id, Tenant.embd_id, Tenant.img2txt_id, Tenant.asr_id, cls.model.update_time]
|
||||
docs = cls.model.select(*fields) \
|
||||
.join(Document, on=(cls.model.doc_id == Document.id)) \
|
||||
.join(Knowledgebase, on=(Document.kb_id == Knowledgebase.id)) \
|
||||
.join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))\
|
||||
.where(
|
||||
Document.status == StatusEnum.VALID.value,
|
||||
~(Document.type == FileType.VIRTUAL.value),
|
||||
cls.model.progress == 0,
|
||||
cls.model.update_time >= tm,
|
||||
(Expression(cls.model.create_time, "%%", comm) == mod))\
|
||||
.order_by(cls.model.update_time.asc())\
|
||||
.paginate(1, items_per_page)
|
||||
return list(docs.dicts())
|
||||
|
||||
|
||||
@classmethod
|
||||
@DB.connection_context()
|
||||
def do_cancel(cls, id):
|
||||
try:
|
||||
cls.model.get_by_id(id)
|
||||
return False
|
||||
except Exception as e:
|
||||
pass
|
||||
return True
|
||||
Reference in New Issue
Block a user