Você não pode selecionar mais de 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.

document_service.py 31KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817
  1. #
  2. # Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. #
  16. import json
  17. import logging
  18. import random
  19. import re
  20. from concurrent.futures import ThreadPoolExecutor
  21. from copy import deepcopy
  22. from datetime import datetime
  23. from io import BytesIO
  24. import trio
  25. import xxhash
  26. from peewee import fn
  27. from api import settings
  28. from api.constants import IMG_BASE64_PREFIX
  29. from api.db import FileType, LLMType, ParserType, StatusEnum, TaskStatus, UserTenantRole
  30. from api.db.db_models import DB, Document, Knowledgebase, Task, Tenant, UserTenant, File2Document, File
  31. from api.db.db_utils import bulk_insert_into_db
  32. from api.db.services.common_service import CommonService
  33. from api.db.services.knowledgebase_service import KnowledgebaseService
  34. from api.utils import current_timestamp, get_format_time, get_uuid
  35. from rag.nlp import rag_tokenizer, search
  36. from rag.settings import get_svr_queue_name, SVR_CONSUMER_GROUP_NAME
  37. from rag.utils.redis_conn import REDIS_CONN
  38. from rag.utils.storage_factory import STORAGE_IMPL
  39. from rag.utils.doc_store_conn import OrderByExpr
  40. class DocumentService(CommonService):
  41. model = Document
  42. @classmethod
  43. def get_cls_model_fields(cls):
  44. return [
  45. cls.model.id,
  46. cls.model.thumbnail,
  47. cls.model.kb_id,
  48. cls.model.parser_id,
  49. cls.model.parser_config,
  50. cls.model.source_type,
  51. cls.model.type,
  52. cls.model.created_by,
  53. cls.model.name,
  54. cls.model.location,
  55. cls.model.size,
  56. cls.model.token_num,
  57. cls.model.chunk_num,
  58. cls.model.progress,
  59. cls.model.progress_msg,
  60. cls.model.process_begin_at,
  61. cls.model.process_duration,
  62. cls.model.meta_fields,
  63. cls.model.suffix,
  64. cls.model.run,
  65. cls.model.status,
  66. cls.model.create_time,
  67. cls.model.create_date,
  68. cls.model.update_time,
  69. cls.model.update_date,
  70. ]
  71. @classmethod
  72. @DB.connection_context()
  73. def get_list(cls, kb_id, page_number, items_per_page,
  74. orderby, desc, keywords, id, name):
  75. fields = cls.get_cls_model_fields()
  76. docs = cls.model.select(*fields).join(File2Document, on = (File2Document.document_id == cls.model.id)).join(File, on = (File.id == File2Document.file_id)).where(cls.model.kb_id == kb_id)
  77. if id:
  78. docs = docs.where(
  79. cls.model.id == id)
  80. if name:
  81. docs = docs.where(
  82. cls.model.name == name
  83. )
  84. if keywords:
  85. docs = docs.where(
  86. fn.LOWER(cls.model.name).contains(keywords.lower())
  87. )
  88. if desc:
  89. docs = docs.order_by(cls.model.getter_by(orderby).desc())
  90. else:
  91. docs = docs.order_by(cls.model.getter_by(orderby).asc())
  92. count = docs.count()
  93. docs = docs.paginate(page_number, items_per_page)
  94. return list(docs.dicts()), count
  95. @classmethod
  96. @DB.connection_context()
  97. def get_by_kb_id(cls, kb_id, page_number, items_per_page,
  98. orderby, desc, keywords, run_status, types, suffix):
  99. fields = cls.get_cls_model_fields()
  100. if keywords:
  101. docs = cls.model.select(*fields).join(File2Document, on=(File2Document.document_id == cls.model.id)).join(File, on=(File.id == File2Document.file_id)).where(
  102. (cls.model.kb_id == kb_id),
  103. (fn.LOWER(cls.model.name).contains(keywords.lower()))
  104. )
  105. else:
  106. docs = cls.model.select(*fields).join(File2Document, on=(File2Document.document_id == cls.model.id)).join(File, on=(File.id == File2Document.file_id)).where(cls.model.kb_id == kb_id)
  107. if run_status:
  108. docs = docs.where(cls.model.run.in_(run_status))
  109. if types:
  110. docs = docs.where(cls.model.type.in_(types))
  111. if suffix:
  112. docs = docs.where(cls.model.suffix.in_(suffix))
  113. count = docs.count()
  114. if desc:
  115. docs = docs.order_by(cls.model.getter_by(orderby).desc())
  116. else:
  117. docs = docs.order_by(cls.model.getter_by(orderby).asc())
  118. if page_number and items_per_page:
  119. docs = docs.paginate(page_number, items_per_page)
  120. return list(docs.dicts()), count
  121. @classmethod
  122. @DB.connection_context()
  123. def get_filter_by_kb_id(cls, kb_id, keywords, run_status, types, suffix):
  124. """
  125. returns:
  126. {
  127. "suffix": {
  128. "ppt": 1,
  129. "doxc": 2
  130. },
  131. "run_status": {
  132. "1": 2,
  133. "2": 2
  134. }
  135. }, total
  136. where "1" => RUNNING, "2" => CANCEL
  137. """
  138. fields = cls.get_cls_model_fields()
  139. if keywords:
  140. query = cls.model.select(*fields).join(File2Document, on=(File2Document.document_id == cls.model.id)).join(File, on=(File.id == File2Document.file_id)).where(
  141. (cls.model.kb_id == kb_id),
  142. (fn.LOWER(cls.model.name).contains(keywords.lower()))
  143. )
  144. else:
  145. query = cls.model.select(*fields).join(File2Document, on=(File2Document.document_id == cls.model.id)).join(File, on=(File.id == File2Document.file_id)).where(cls.model.kb_id == kb_id)
  146. if run_status:
  147. query = query.where(cls.model.run.in_(run_status))
  148. if types:
  149. query = query.where(cls.model.type.in_(types))
  150. if suffix:
  151. query = query.where(cls.model.suffix.in_(suffix))
  152. rows = query.select(cls.model.run, cls.model.suffix)
  153. total = rows.count()
  154. suffix_counter = {}
  155. run_status_counter = {}
  156. for row in rows:
  157. suffix_counter[row.suffix] = suffix_counter.get(row.suffix, 0) + 1
  158. run_status_counter[str(row.run)] = run_status_counter.get(str(row.run), 0) + 1
  159. return {
  160. "suffix": suffix_counter,
  161. "run_status": run_status_counter
  162. }, total
  163. @classmethod
  164. @DB.connection_context()
  165. def count_by_kb_id(cls, kb_id, keywords, run_status, types):
  166. if keywords:
  167. docs = cls.model.select().where(
  168. (cls.model.kb_id == kb_id),
  169. (fn.LOWER(cls.model.name).contains(keywords.lower()))
  170. )
  171. else:
  172. docs = cls.model.select().where(cls.model.kb_id == kb_id)
  173. if run_status:
  174. docs = docs.where(cls.model.run.in_(run_status))
  175. if types:
  176. docs = docs.where(cls.model.type.in_(types))
  177. count = docs.count()
  178. return count
  179. @classmethod
  180. @DB.connection_context()
  181. def get_total_size_by_kb_id(cls, kb_id, keywords="", run_status=[], types=[]):
  182. query = cls.model.select(fn.COALESCE(fn.SUM(cls.model.size), 0)).where(
  183. cls.model.kb_id == kb_id
  184. )
  185. if keywords:
  186. query = query.where(fn.LOWER(cls.model.name).contains(keywords.lower()))
  187. if run_status:
  188. query = query.where(cls.model.run.in_(run_status))
  189. if types:
  190. query = query.where(cls.model.type.in_(types))
  191. return int(query.scalar()) or 0
  192. @classmethod
  193. @DB.connection_context()
  194. def insert(cls, doc):
  195. if not cls.save(**doc):
  196. raise RuntimeError("Database error (Document)!")
  197. if not KnowledgebaseService.atomic_increase_doc_num_by_id(doc["kb_id"]):
  198. raise RuntimeError("Database error (Knowledgebase)!")
  199. return Document(**doc)
  200. @classmethod
  201. @DB.connection_context()
  202. def remove_document(cls, doc, tenant_id):
  203. cls.clear_chunk_num(doc.id)
  204. try:
  205. page = 0
  206. page_size = 1000
  207. all_chunk_ids = []
  208. while True:
  209. chunks = settings.docStoreConn.search(["img_id"], [], {"doc_id": doc.id}, [], OrderByExpr(),
  210. page * page_size, page_size, search.index_name(tenant_id),
  211. [doc.kb_id])
  212. chunk_ids = settings.docStoreConn.getChunkIds(chunks)
  213. if not chunk_ids:
  214. break
  215. all_chunk_ids.extend(chunk_ids)
  216. page += 1
  217. for cid in all_chunk_ids:
  218. if STORAGE_IMPL.obj_exist(doc.kb_id, cid):
  219. STORAGE_IMPL.rm(doc.kb_id, cid)
  220. if doc.thumbnail and not doc.thumbnail.startswith(IMG_BASE64_PREFIX):
  221. if STORAGE_IMPL.obj_exist(doc.kb_id, doc.thumbnail):
  222. STORAGE_IMPL.rm(doc.kb_id, doc.thumbnail)
  223. settings.docStoreConn.delete({"doc_id": doc.id}, search.index_name(tenant_id), doc.kb_id)
  224. graph_source = settings.docStoreConn.getFields(
  225. settings.docStoreConn.search(["source_id"], [], {"kb_id": doc.kb_id, "knowledge_graph_kwd": ["graph"]}, [], OrderByExpr(), 0, 1, search.index_name(tenant_id), [doc.kb_id]), ["source_id"]
  226. )
  227. if len(graph_source) > 0 and doc.id in list(graph_source.values())[0]["source_id"]:
  228. settings.docStoreConn.update({"kb_id": doc.kb_id, "knowledge_graph_kwd": ["entity", "relation", "graph", "subgraph", "community_report"], "source_id": doc.id},
  229. {"remove": {"source_id": doc.id}},
  230. search.index_name(tenant_id), doc.kb_id)
  231. settings.docStoreConn.update({"kb_id": doc.kb_id, "knowledge_graph_kwd": ["graph"]},
  232. {"removed_kwd": "Y"},
  233. search.index_name(tenant_id), doc.kb_id)
  234. settings.docStoreConn.delete({"kb_id": doc.kb_id, "knowledge_graph_kwd": ["entity", "relation", "graph", "subgraph", "community_report"], "must_not": {"exists": "source_id"}},
  235. search.index_name(tenant_id), doc.kb_id)
  236. except Exception:
  237. pass
  238. return cls.delete_by_id(doc.id)
  239. @classmethod
  240. @DB.connection_context()
  241. def get_newly_uploaded(cls):
  242. fields = [
  243. cls.model.id,
  244. cls.model.kb_id,
  245. cls.model.parser_id,
  246. cls.model.parser_config,
  247. cls.model.name,
  248. cls.model.type,
  249. cls.model.location,
  250. cls.model.size,
  251. Knowledgebase.tenant_id,
  252. Tenant.embd_id,
  253. Tenant.img2txt_id,
  254. Tenant.asr_id,
  255. cls.model.update_time]
  256. docs = cls.model.select(*fields) \
  257. .join(Knowledgebase, on=(cls.model.kb_id == Knowledgebase.id)) \
  258. .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id)) \
  259. .where(
  260. cls.model.status == StatusEnum.VALID.value,
  261. ~(cls.model.type == FileType.VIRTUAL.value),
  262. cls.model.progress == 0,
  263. cls.model.update_time >= current_timestamp() - 1000 * 600,
  264. cls.model.run == TaskStatus.RUNNING.value) \
  265. .order_by(cls.model.update_time.asc())
  266. return list(docs.dicts())
  267. @classmethod
  268. @DB.connection_context()
  269. def get_unfinished_docs(cls):
  270. fields = [cls.model.id, cls.model.process_begin_at, cls.model.parser_config, cls.model.progress_msg,
  271. cls.model.run, cls.model.parser_id]
  272. docs = cls.model.select(*fields) \
  273. .where(
  274. cls.model.status == StatusEnum.VALID.value,
  275. ~(cls.model.type == FileType.VIRTUAL.value),
  276. cls.model.progress < 1,
  277. cls.model.progress > 0)
  278. return list(docs.dicts())
  279. @classmethod
  280. @DB.connection_context()
  281. def increment_chunk_num(cls, doc_id, kb_id, token_num, chunk_num, duration):
  282. num = cls.model.update(token_num=cls.model.token_num + token_num,
  283. chunk_num=cls.model.chunk_num + chunk_num,
  284. process_duration=cls.model.process_duration + duration).where(
  285. cls.model.id == doc_id).execute()
  286. if num == 0:
  287. raise LookupError(
  288. "Document not found which is supposed to be there")
  289. num = Knowledgebase.update(
  290. token_num=Knowledgebase.token_num +
  291. token_num,
  292. chunk_num=Knowledgebase.chunk_num +
  293. chunk_num).where(
  294. Knowledgebase.id == kb_id).execute()
  295. return num
  296. @classmethod
  297. @DB.connection_context()
  298. def decrement_chunk_num(cls, doc_id, kb_id, token_num, chunk_num, duration):
  299. num = cls.model.update(token_num=cls.model.token_num - token_num,
  300. chunk_num=cls.model.chunk_num - chunk_num,
  301. process_duration=cls.model.process_duration + duration).where(
  302. cls.model.id == doc_id).execute()
  303. if num == 0:
  304. raise LookupError(
  305. "Document not found which is supposed to be there")
  306. num = Knowledgebase.update(
  307. token_num=Knowledgebase.token_num -
  308. token_num,
  309. chunk_num=Knowledgebase.chunk_num -
  310. chunk_num
  311. ).where(
  312. Knowledgebase.id == kb_id).execute()
  313. return num
  314. @classmethod
  315. @DB.connection_context()
  316. def clear_chunk_num(cls, doc_id):
  317. doc = cls.model.get_by_id(doc_id)
  318. assert doc, "Can't fine document in database."
  319. num = Knowledgebase.update(
  320. token_num=Knowledgebase.token_num -
  321. doc.token_num,
  322. chunk_num=Knowledgebase.chunk_num -
  323. doc.chunk_num,
  324. doc_num=Knowledgebase.doc_num - 1
  325. ).where(
  326. Knowledgebase.id == doc.kb_id).execute()
  327. return num
  328. @classmethod
  329. @DB.connection_context()
  330. def clear_chunk_num_when_rerun(cls, doc_id):
  331. doc = cls.model.get_by_id(doc_id)
  332. assert doc, "Can't fine document in database."
  333. num = (
  334. Knowledgebase.update(
  335. token_num=Knowledgebase.token_num - doc.token_num,
  336. chunk_num=Knowledgebase.chunk_num - doc.chunk_num,
  337. )
  338. .where(Knowledgebase.id == doc.kb_id)
  339. .execute()
  340. )
  341. return num
  342. @classmethod
  343. @DB.connection_context()
  344. def get_tenant_id(cls, doc_id):
  345. docs = cls.model.select(
  346. Knowledgebase.tenant_id).join(
  347. Knowledgebase, on=(
  348. Knowledgebase.id == cls.model.kb_id)).where(
  349. cls.model.id == doc_id, Knowledgebase.status == StatusEnum.VALID.value)
  350. docs = docs.dicts()
  351. if not docs:
  352. return
  353. return docs[0]["tenant_id"]
  354. @classmethod
  355. @DB.connection_context()
  356. def get_knowledgebase_id(cls, doc_id):
  357. docs = cls.model.select(cls.model.kb_id).where(cls.model.id == doc_id)
  358. docs = docs.dicts()
  359. if not docs:
  360. return
  361. return docs[0]["kb_id"]
  362. @classmethod
  363. @DB.connection_context()
  364. def get_tenant_id_by_name(cls, name):
  365. docs = cls.model.select(
  366. Knowledgebase.tenant_id).join(
  367. Knowledgebase, on=(
  368. Knowledgebase.id == cls.model.kb_id)).where(
  369. cls.model.name == name, Knowledgebase.status == StatusEnum.VALID.value)
  370. docs = docs.dicts()
  371. if not docs:
  372. return
  373. return docs[0]["tenant_id"]
  374. @classmethod
  375. @DB.connection_context()
  376. def accessible(cls, doc_id, user_id):
  377. docs = cls.model.select(
  378. cls.model.id).join(
  379. Knowledgebase, on=(
  380. Knowledgebase.id == cls.model.kb_id)
  381. ).join(UserTenant, on=(UserTenant.tenant_id == Knowledgebase.tenant_id)
  382. ).where(cls.model.id == doc_id, UserTenant.user_id == user_id).paginate(0, 1)
  383. docs = docs.dicts()
  384. if not docs:
  385. return False
  386. return True
  387. @classmethod
  388. @DB.connection_context()
  389. def accessible4deletion(cls, doc_id, user_id):
  390. docs = cls.model.select(cls.model.id
  391. ).join(
  392. Knowledgebase, on=(
  393. Knowledgebase.id == cls.model.kb_id)
  394. ).join(
  395. UserTenant, on=(
  396. (UserTenant.tenant_id == Knowledgebase.created_by) & (UserTenant.user_id == user_id))
  397. ).where(
  398. cls.model.id == doc_id,
  399. UserTenant.status == StatusEnum.VALID.value,
  400. ((UserTenant.role == UserTenantRole.NORMAL) | (UserTenant.role == UserTenantRole.OWNER))
  401. ).paginate(0, 1)
  402. docs = docs.dicts()
  403. if not docs:
  404. return False
  405. return True
  406. @classmethod
  407. @DB.connection_context()
  408. def get_embd_id(cls, doc_id):
  409. docs = cls.model.select(
  410. Knowledgebase.embd_id).join(
  411. Knowledgebase, on=(
  412. Knowledgebase.id == cls.model.kb_id)).where(
  413. cls.model.id == doc_id, Knowledgebase.status == StatusEnum.VALID.value)
  414. docs = docs.dicts()
  415. if not docs:
  416. return
  417. return docs[0]["embd_id"]
  418. @classmethod
  419. @DB.connection_context()
  420. def get_chunking_config(cls, doc_id):
  421. configs = (
  422. cls.model.select(
  423. cls.model.id,
  424. cls.model.kb_id,
  425. cls.model.parser_id,
  426. cls.model.parser_config,
  427. Knowledgebase.language,
  428. Knowledgebase.embd_id,
  429. Tenant.id.alias("tenant_id"),
  430. Tenant.img2txt_id,
  431. Tenant.asr_id,
  432. Tenant.llm_id,
  433. )
  434. .join(Knowledgebase, on=(cls.model.kb_id == Knowledgebase.id))
  435. .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))
  436. .where(cls.model.id == doc_id)
  437. )
  438. configs = configs.dicts()
  439. if not configs:
  440. return None
  441. return configs[0]
  442. @classmethod
  443. @DB.connection_context()
  444. def get_doc_id_by_doc_name(cls, doc_name):
  445. fields = [cls.model.id]
  446. doc_id = cls.model.select(*fields) \
  447. .where(cls.model.name == doc_name)
  448. doc_id = doc_id.dicts()
  449. if not doc_id:
  450. return
  451. return doc_id[0]["id"]
  452. @classmethod
  453. @DB.connection_context()
  454. def get_doc_ids_by_doc_names(cls, doc_names):
  455. if not doc_names:
  456. return []
  457. query = cls.model.select(cls.model.id).where(cls.model.name.in_(doc_names))
  458. return list(query.scalars().iterator())
  459. @classmethod
  460. @DB.connection_context()
  461. def get_thumbnails(cls, docids):
  462. fields = [cls.model.id, cls.model.kb_id, cls.model.thumbnail]
  463. return list(cls.model.select(
  464. *fields).where(cls.model.id.in_(docids)).dicts())
  465. @classmethod
  466. @DB.connection_context()
  467. def update_parser_config(cls, id, config):
  468. if not config:
  469. return
  470. e, d = cls.get_by_id(id)
  471. if not e:
  472. raise LookupError(f"Document({id}) not found.")
  473. def dfs_update(old, new):
  474. for k, v in new.items():
  475. if k not in old:
  476. old[k] = v
  477. continue
  478. if isinstance(v, dict):
  479. assert isinstance(old[k], dict)
  480. dfs_update(old[k], v)
  481. else:
  482. old[k] = v
  483. dfs_update(d.parser_config, config)
  484. if not config.get("raptor") and d.parser_config.get("raptor"):
  485. del d.parser_config["raptor"]
  486. cls.update_by_id(id, {"parser_config": d.parser_config})
  487. @classmethod
  488. @DB.connection_context()
  489. def get_doc_count(cls, tenant_id):
  490. docs = cls.model.select(cls.model.id).join(Knowledgebase,
  491. on=(Knowledgebase.id == cls.model.kb_id)).where(
  492. Knowledgebase.tenant_id == tenant_id)
  493. return len(docs)
  494. @classmethod
  495. @DB.connection_context()
  496. def begin2parse(cls, docid):
  497. cls.update_by_id(
  498. docid, {"progress": random.random() * 1 / 100.,
  499. "progress_msg": "Task is queued...",
  500. "process_begin_at": get_format_time()
  501. })
  502. @classmethod
  503. @DB.connection_context()
  504. def update_meta_fields(cls, doc_id, meta_fields):
  505. return cls.update_by_id(doc_id, {"meta_fields": meta_fields})
  506. @classmethod
  507. @DB.connection_context()
  508. def update_progress(cls):
  509. docs = cls.get_unfinished_docs()
  510. for d in docs:
  511. try:
  512. tsks = Task.query(doc_id=d["id"], order_by=Task.create_time)
  513. if not tsks:
  514. continue
  515. msg = []
  516. prg = 0
  517. finished = True
  518. bad = 0
  519. has_raptor = False
  520. has_graphrag = False
  521. e, doc = DocumentService.get_by_id(d["id"])
  522. status = doc.run # TaskStatus.RUNNING.value
  523. priority = 0
  524. for t in tsks:
  525. if 0 <= t.progress < 1:
  526. finished = False
  527. if t.progress == -1:
  528. bad += 1
  529. prg += t.progress if t.progress >= 0 else 0
  530. if t.progress_msg.strip():
  531. msg.append(t.progress_msg)
  532. if t.task_type == "raptor":
  533. has_raptor = True
  534. elif t.task_type == "graphrag":
  535. has_graphrag = True
  536. priority = max(priority, t.priority)
  537. prg /= len(tsks)
  538. if finished and bad:
  539. prg = -1
  540. status = TaskStatus.FAIL.value
  541. elif finished:
  542. if (d["parser_config"].get("raptor") or {}).get("use_raptor") and not has_raptor:
  543. queue_raptor_o_graphrag_tasks(d, "raptor", priority)
  544. prg = 0.98 * len(tsks) / (len(tsks) + 1)
  545. elif (d["parser_config"].get("graphrag") or {}).get("use_graphrag") and not has_graphrag:
  546. queue_raptor_o_graphrag_tasks(d, "graphrag", priority)
  547. prg = 0.98 * len(tsks) / (len(tsks) + 1)
  548. else:
  549. status = TaskStatus.DONE.value
  550. msg = "\n".join(sorted(msg))
  551. info = {
  552. "process_duration": datetime.timestamp(
  553. datetime.now()) -
  554. d["process_begin_at"].timestamp(),
  555. "run": status}
  556. if prg != 0:
  557. info["progress"] = prg
  558. if msg:
  559. info["progress_msg"] = msg
  560. if msg.endswith("created task graphrag") or msg.endswith("created task raptor"):
  561. info["progress_msg"] += "\n%d tasks are ahead in the queue..."%get_queue_length(priority)
  562. else:
  563. info["progress_msg"] = "%d tasks are ahead in the queue..."%get_queue_length(priority)
  564. cls.update_by_id(d["id"], info)
  565. except Exception as e:
  566. if str(e).find("'0'") < 0:
  567. logging.exception("fetch task exception")
  568. @classmethod
  569. @DB.connection_context()
  570. def get_kb_doc_count(cls, kb_id):
  571. return len(cls.model.select(cls.model.id).where(
  572. cls.model.kb_id == kb_id).dicts())
  573. @classmethod
  574. @DB.connection_context()
  575. def do_cancel(cls, doc_id):
  576. try:
  577. _, doc = DocumentService.get_by_id(doc_id)
  578. return doc.run == TaskStatus.CANCEL.value or doc.progress < 0
  579. except Exception:
  580. pass
  581. return False
  582. def queue_raptor_o_graphrag_tasks(doc, ty, priority):
  583. chunking_config = DocumentService.get_chunking_config(doc["id"])
  584. hasher = xxhash.xxh64()
  585. for field in sorted(chunking_config.keys()):
  586. hasher.update(str(chunking_config[field]).encode("utf-8"))
  587. def new_task():
  588. nonlocal doc
  589. return {
  590. "id": get_uuid(),
  591. "doc_id": doc["id"],
  592. "from_page": 100000000,
  593. "to_page": 100000000,
  594. "task_type": ty,
  595. "progress_msg": datetime.now().strftime("%H:%M:%S") + " created task " + ty
  596. }
  597. task = new_task()
  598. for field in ["doc_id", "from_page", "to_page"]:
  599. hasher.update(str(task.get(field, "")).encode("utf-8"))
  600. hasher.update(ty.encode("utf-8"))
  601. task["digest"] = hasher.hexdigest()
  602. bulk_insert_into_db(Task, [task], True)
  603. assert REDIS_CONN.queue_product(get_svr_queue_name(priority), message=task), "Can't access Redis. Please check the Redis' status."
  604. def get_queue_length(priority):
  605. group_info = REDIS_CONN.queue_info(get_svr_queue_name(priority), SVR_CONSUMER_GROUP_NAME)
  606. return int(group_info.get("lag", 0))
  607. def doc_upload_and_parse(conversation_id, file_objs, user_id):
  608. from api.db.services.api_service import API4ConversationService
  609. from api.db.services.conversation_service import ConversationService
  610. from api.db.services.dialog_service import DialogService
  611. from api.db.services.file_service import FileService
  612. from api.db.services.llm_service import LLMBundle
  613. from api.db.services.user_service import TenantService
  614. from rag.app import audio, email, naive, picture, presentation
  615. e, conv = ConversationService.get_by_id(conversation_id)
  616. if not e:
  617. e, conv = API4ConversationService.get_by_id(conversation_id)
  618. assert e, "Conversation not found!"
  619. e, dia = DialogService.get_by_id(conv.dialog_id)
  620. if not dia.kb_ids:
  621. raise LookupError("No knowledge base associated with this conversation. "
  622. "Please add a knowledge base before uploading documents")
  623. kb_id = dia.kb_ids[0]
  624. e, kb = KnowledgebaseService.get_by_id(kb_id)
  625. if not e:
  626. raise LookupError("Can't find this knowledgebase!")
  627. embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING, llm_name=kb.embd_id, lang=kb.language)
  628. err, files = FileService.upload_document(kb, file_objs, user_id)
  629. assert not err, "\n".join(err)
  630. def dummy(prog=None, msg=""):
  631. pass
  632. FACTORY = {
  633. ParserType.PRESENTATION.value: presentation,
  634. ParserType.PICTURE.value: picture,
  635. ParserType.AUDIO.value: audio,
  636. ParserType.EMAIL.value: email
  637. }
  638. parser_config = {"chunk_token_num": 4096, "delimiter": "\n!?;。;!?", "layout_recognize": "Plain Text"}
  639. exe = ThreadPoolExecutor(max_workers=12)
  640. threads = []
  641. doc_nm = {}
  642. for d, blob in files:
  643. doc_nm[d["id"]] = d["name"]
  644. for d, blob in files:
  645. kwargs = {
  646. "callback": dummy,
  647. "parser_config": parser_config,
  648. "from_page": 0,
  649. "to_page": 100000,
  650. "tenant_id": kb.tenant_id,
  651. "lang": kb.language
  652. }
  653. threads.append(exe.submit(FACTORY.get(d["parser_id"], naive).chunk, d["name"], blob, **kwargs))
  654. for (docinfo, _), th in zip(files, threads):
  655. docs = []
  656. doc = {
  657. "doc_id": docinfo["id"],
  658. "kb_id": [kb.id]
  659. }
  660. for ck in th.result():
  661. d = deepcopy(doc)
  662. d.update(ck)
  663. d["id"] = xxhash.xxh64((ck["content_with_weight"] + str(d["doc_id"])).encode("utf-8")).hexdigest()
  664. d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
  665. d["create_timestamp_flt"] = datetime.now().timestamp()
  666. if not d.get("image"):
  667. docs.append(d)
  668. continue
  669. output_buffer = BytesIO()
  670. if isinstance(d["image"], bytes):
  671. output_buffer = BytesIO(d["image"])
  672. else:
  673. d["image"].save(output_buffer, format='JPEG')
  674. STORAGE_IMPL.put(kb.id, d["id"], output_buffer.getvalue())
  675. d["img_id"] = "{}-{}".format(kb.id, d["id"])
  676. d.pop("image", None)
  677. docs.append(d)
  678. parser_ids = {d["id"]: d["parser_id"] for d, _ in files}
  679. docids = [d["id"] for d, _ in files]
  680. chunk_counts = {id: 0 for id in docids}
  681. token_counts = {id: 0 for id in docids}
  682. es_bulk_size = 64
  683. def embedding(doc_id, cnts, batch_size=16):
  684. nonlocal embd_mdl, chunk_counts, token_counts
  685. vects = []
  686. for i in range(0, len(cnts), batch_size):
  687. vts, c = embd_mdl.encode(cnts[i: i + batch_size])
  688. vects.extend(vts.tolist())
  689. chunk_counts[doc_id] += len(cnts[i:i + batch_size])
  690. token_counts[doc_id] += c
  691. return vects
  692. idxnm = search.index_name(kb.tenant_id)
  693. try_create_idx = True
  694. _, tenant = TenantService.get_by_id(kb.tenant_id)
  695. llm_bdl = LLMBundle(kb.tenant_id, LLMType.CHAT, tenant.llm_id)
  696. for doc_id in docids:
  697. cks = [c for c in docs if c["doc_id"] == doc_id]
  698. if parser_ids[doc_id] != ParserType.PICTURE.value:
  699. from graphrag.general.mind_map_extractor import MindMapExtractor
  700. mindmap = MindMapExtractor(llm_bdl)
  701. try:
  702. mind_map = trio.run(mindmap, [c["content_with_weight"] for c in docs if c["doc_id"] == doc_id])
  703. mind_map = json.dumps(mind_map.output, ensure_ascii=False, indent=2)
  704. if len(mind_map) < 32:
  705. raise Exception("Few content: " + mind_map)
  706. cks.append({
  707. "id": get_uuid(),
  708. "doc_id": doc_id,
  709. "kb_id": [kb.id],
  710. "docnm_kwd": doc_nm[doc_id],
  711. "title_tks": rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", doc_nm[doc_id])),
  712. "content_ltks": rag_tokenizer.tokenize("summary summarize 总结 概况 file 文件 概括"),
  713. "content_with_weight": mind_map,
  714. "knowledge_graph_kwd": "mind_map"
  715. })
  716. except Exception as e:
  717. logging.exception("Mind map generation error")
  718. vects = embedding(doc_id, [c["content_with_weight"] for c in cks])
  719. assert len(cks) == len(vects)
  720. for i, d in enumerate(cks):
  721. v = vects[i]
  722. d["q_%d_vec" % len(v)] = v
  723. for b in range(0, len(cks), es_bulk_size):
  724. if try_create_idx:
  725. if not settings.docStoreConn.indexExist(idxnm, kb_id):
  726. settings.docStoreConn.createIdx(idxnm, kb_id, len(vects[0]))
  727. try_create_idx = False
  728. settings.docStoreConn.insert(cks[b:b + es_bulk_size], idxnm, kb_id)
  729. DocumentService.increment_chunk_num(
  730. doc_id, kb.id, token_counts[doc_id], chunk_counts[doc_id], 0)
  731. return [d["id"] for d, _ in files]