Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.

document_service.py 31KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819
  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. from api.db.services.task_service import TaskService
  204. cls.clear_chunk_num(doc.id)
  205. try:
  206. TaskService.filter_delete(Task.doc_id == doc.id)
  207. page = 0
  208. page_size = 1000
  209. all_chunk_ids = []
  210. while True:
  211. chunks = settings.docStoreConn.search(["img_id"], [], {"doc_id": doc.id}, [], OrderByExpr(),
  212. page * page_size, page_size, search.index_name(tenant_id),
  213. [doc.kb_id])
  214. chunk_ids = settings.docStoreConn.getChunkIds(chunks)
  215. if not chunk_ids:
  216. break
  217. all_chunk_ids.extend(chunk_ids)
  218. page += 1
  219. for cid in all_chunk_ids:
  220. if STORAGE_IMPL.obj_exist(doc.kb_id, cid):
  221. STORAGE_IMPL.rm(doc.kb_id, cid)
  222. if doc.thumbnail and not doc.thumbnail.startswith(IMG_BASE64_PREFIX):
  223. if STORAGE_IMPL.obj_exist(doc.kb_id, doc.thumbnail):
  224. STORAGE_IMPL.rm(doc.kb_id, doc.thumbnail)
  225. settings.docStoreConn.delete({"doc_id": doc.id}, search.index_name(tenant_id), doc.kb_id)
  226. graph_source = settings.docStoreConn.getFields(
  227. 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"]
  228. )
  229. if len(graph_source) > 0 and doc.id in list(graph_source.values())[0]["source_id"]:
  230. settings.docStoreConn.update({"kb_id": doc.kb_id, "knowledge_graph_kwd": ["entity", "relation", "graph", "subgraph", "community_report"], "source_id": doc.id},
  231. {"remove": {"source_id": doc.id}},
  232. search.index_name(tenant_id), doc.kb_id)
  233. settings.docStoreConn.update({"kb_id": doc.kb_id, "knowledge_graph_kwd": ["graph"]},
  234. {"removed_kwd": "Y"},
  235. search.index_name(tenant_id), doc.kb_id)
  236. settings.docStoreConn.delete({"kb_id": doc.kb_id, "knowledge_graph_kwd": ["entity", "relation", "graph", "subgraph", "community_report"], "must_not": {"exists": "source_id"}},
  237. search.index_name(tenant_id), doc.kb_id)
  238. except Exception:
  239. pass
  240. return cls.delete_by_id(doc.id)
  241. @classmethod
  242. @DB.connection_context()
  243. def get_newly_uploaded(cls):
  244. fields = [
  245. cls.model.id,
  246. cls.model.kb_id,
  247. cls.model.parser_id,
  248. cls.model.parser_config,
  249. cls.model.name,
  250. cls.model.type,
  251. cls.model.location,
  252. cls.model.size,
  253. Knowledgebase.tenant_id,
  254. Tenant.embd_id,
  255. Tenant.img2txt_id,
  256. Tenant.asr_id,
  257. cls.model.update_time]
  258. docs = cls.model.select(*fields) \
  259. .join(Knowledgebase, on=(cls.model.kb_id == Knowledgebase.id)) \
  260. .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id)) \
  261. .where(
  262. cls.model.status == StatusEnum.VALID.value,
  263. ~(cls.model.type == FileType.VIRTUAL.value),
  264. cls.model.progress == 0,
  265. cls.model.update_time >= current_timestamp() - 1000 * 600,
  266. cls.model.run == TaskStatus.RUNNING.value) \
  267. .order_by(cls.model.update_time.asc())
  268. return list(docs.dicts())
  269. @classmethod
  270. @DB.connection_context()
  271. def get_unfinished_docs(cls):
  272. fields = [cls.model.id, cls.model.process_begin_at, cls.model.parser_config, cls.model.progress_msg,
  273. cls.model.run, cls.model.parser_id]
  274. docs = cls.model.select(*fields) \
  275. .where(
  276. cls.model.status == StatusEnum.VALID.value,
  277. ~(cls.model.type == FileType.VIRTUAL.value),
  278. cls.model.progress < 1,
  279. cls.model.progress > 0)
  280. return list(docs.dicts())
  281. @classmethod
  282. @DB.connection_context()
  283. def increment_chunk_num(cls, doc_id, kb_id, token_num, chunk_num, duration):
  284. num = cls.model.update(token_num=cls.model.token_num + token_num,
  285. chunk_num=cls.model.chunk_num + chunk_num,
  286. process_duration=cls.model.process_duration + duration).where(
  287. cls.model.id == doc_id).execute()
  288. if num == 0:
  289. raise LookupError(
  290. "Document not found which is supposed to be there")
  291. num = Knowledgebase.update(
  292. token_num=Knowledgebase.token_num +
  293. token_num,
  294. chunk_num=Knowledgebase.chunk_num +
  295. chunk_num).where(
  296. Knowledgebase.id == kb_id).execute()
  297. return num
  298. @classmethod
  299. @DB.connection_context()
  300. def decrement_chunk_num(cls, doc_id, kb_id, token_num, chunk_num, duration):
  301. num = cls.model.update(token_num=cls.model.token_num - token_num,
  302. chunk_num=cls.model.chunk_num - chunk_num,
  303. process_duration=cls.model.process_duration + duration).where(
  304. cls.model.id == doc_id).execute()
  305. if num == 0:
  306. raise LookupError(
  307. "Document not found which is supposed to be there")
  308. num = Knowledgebase.update(
  309. token_num=Knowledgebase.token_num -
  310. token_num,
  311. chunk_num=Knowledgebase.chunk_num -
  312. chunk_num
  313. ).where(
  314. Knowledgebase.id == kb_id).execute()
  315. return num
  316. @classmethod
  317. @DB.connection_context()
  318. def clear_chunk_num(cls, doc_id):
  319. doc = cls.model.get_by_id(doc_id)
  320. assert doc, "Can't fine document in database."
  321. num = Knowledgebase.update(
  322. token_num=Knowledgebase.token_num -
  323. doc.token_num,
  324. chunk_num=Knowledgebase.chunk_num -
  325. doc.chunk_num,
  326. doc_num=Knowledgebase.doc_num - 1
  327. ).where(
  328. Knowledgebase.id == doc.kb_id).execute()
  329. return num
  330. @classmethod
  331. @DB.connection_context()
  332. def clear_chunk_num_when_rerun(cls, doc_id):
  333. doc = cls.model.get_by_id(doc_id)
  334. assert doc, "Can't fine document in database."
  335. num = (
  336. Knowledgebase.update(
  337. token_num=Knowledgebase.token_num - doc.token_num,
  338. chunk_num=Knowledgebase.chunk_num - doc.chunk_num,
  339. )
  340. .where(Knowledgebase.id == doc.kb_id)
  341. .execute()
  342. )
  343. return num
  344. @classmethod
  345. @DB.connection_context()
  346. def get_tenant_id(cls, doc_id):
  347. docs = cls.model.select(
  348. Knowledgebase.tenant_id).join(
  349. Knowledgebase, on=(
  350. Knowledgebase.id == cls.model.kb_id)).where(
  351. cls.model.id == doc_id, Knowledgebase.status == StatusEnum.VALID.value)
  352. docs = docs.dicts()
  353. if not docs:
  354. return
  355. return docs[0]["tenant_id"]
  356. @classmethod
  357. @DB.connection_context()
  358. def get_knowledgebase_id(cls, doc_id):
  359. docs = cls.model.select(cls.model.kb_id).where(cls.model.id == doc_id)
  360. docs = docs.dicts()
  361. if not docs:
  362. return
  363. return docs[0]["kb_id"]
  364. @classmethod
  365. @DB.connection_context()
  366. def get_tenant_id_by_name(cls, name):
  367. docs = cls.model.select(
  368. Knowledgebase.tenant_id).join(
  369. Knowledgebase, on=(
  370. Knowledgebase.id == cls.model.kb_id)).where(
  371. cls.model.name == name, Knowledgebase.status == StatusEnum.VALID.value)
  372. docs = docs.dicts()
  373. if not docs:
  374. return
  375. return docs[0]["tenant_id"]
  376. @classmethod
  377. @DB.connection_context()
  378. def accessible(cls, doc_id, user_id):
  379. docs = cls.model.select(
  380. cls.model.id).join(
  381. Knowledgebase, on=(
  382. Knowledgebase.id == cls.model.kb_id)
  383. ).join(UserTenant, on=(UserTenant.tenant_id == Knowledgebase.tenant_id)
  384. ).where(cls.model.id == doc_id, UserTenant.user_id == user_id).paginate(0, 1)
  385. docs = docs.dicts()
  386. if not docs:
  387. return False
  388. return True
  389. @classmethod
  390. @DB.connection_context()
  391. def accessible4deletion(cls, doc_id, user_id):
  392. docs = cls.model.select(cls.model.id
  393. ).join(
  394. Knowledgebase, on=(
  395. Knowledgebase.id == cls.model.kb_id)
  396. ).join(
  397. UserTenant, on=(
  398. (UserTenant.tenant_id == Knowledgebase.created_by) & (UserTenant.user_id == user_id))
  399. ).where(
  400. cls.model.id == doc_id,
  401. UserTenant.status == StatusEnum.VALID.value,
  402. ((UserTenant.role == UserTenantRole.NORMAL) | (UserTenant.role == UserTenantRole.OWNER))
  403. ).paginate(0, 1)
  404. docs = docs.dicts()
  405. if not docs:
  406. return False
  407. return True
  408. @classmethod
  409. @DB.connection_context()
  410. def get_embd_id(cls, doc_id):
  411. docs = cls.model.select(
  412. Knowledgebase.embd_id).join(
  413. Knowledgebase, on=(
  414. Knowledgebase.id == cls.model.kb_id)).where(
  415. cls.model.id == doc_id, Knowledgebase.status == StatusEnum.VALID.value)
  416. docs = docs.dicts()
  417. if not docs:
  418. return
  419. return docs[0]["embd_id"]
  420. @classmethod
  421. @DB.connection_context()
  422. def get_chunking_config(cls, doc_id):
  423. configs = (
  424. cls.model.select(
  425. cls.model.id,
  426. cls.model.kb_id,
  427. cls.model.parser_id,
  428. cls.model.parser_config,
  429. Knowledgebase.language,
  430. Knowledgebase.embd_id,
  431. Tenant.id.alias("tenant_id"),
  432. Tenant.img2txt_id,
  433. Tenant.asr_id,
  434. Tenant.llm_id,
  435. )
  436. .join(Knowledgebase, on=(cls.model.kb_id == Knowledgebase.id))
  437. .join(Tenant, on=(Knowledgebase.tenant_id == Tenant.id))
  438. .where(cls.model.id == doc_id)
  439. )
  440. configs = configs.dicts()
  441. if not configs:
  442. return None
  443. return configs[0]
  444. @classmethod
  445. @DB.connection_context()
  446. def get_doc_id_by_doc_name(cls, doc_name):
  447. fields = [cls.model.id]
  448. doc_id = cls.model.select(*fields) \
  449. .where(cls.model.name == doc_name)
  450. doc_id = doc_id.dicts()
  451. if not doc_id:
  452. return
  453. return doc_id[0]["id"]
  454. @classmethod
  455. @DB.connection_context()
  456. def get_doc_ids_by_doc_names(cls, doc_names):
  457. if not doc_names:
  458. return []
  459. query = cls.model.select(cls.model.id).where(cls.model.name.in_(doc_names))
  460. return list(query.scalars().iterator())
  461. @classmethod
  462. @DB.connection_context()
  463. def get_thumbnails(cls, docids):
  464. fields = [cls.model.id, cls.model.kb_id, cls.model.thumbnail]
  465. return list(cls.model.select(
  466. *fields).where(cls.model.id.in_(docids)).dicts())
  467. @classmethod
  468. @DB.connection_context()
  469. def update_parser_config(cls, id, config):
  470. if not config:
  471. return
  472. e, d = cls.get_by_id(id)
  473. if not e:
  474. raise LookupError(f"Document({id}) not found.")
  475. def dfs_update(old, new):
  476. for k, v in new.items():
  477. if k not in old:
  478. old[k] = v
  479. continue
  480. if isinstance(v, dict):
  481. assert isinstance(old[k], dict)
  482. dfs_update(old[k], v)
  483. else:
  484. old[k] = v
  485. dfs_update(d.parser_config, config)
  486. if not config.get("raptor") and d.parser_config.get("raptor"):
  487. del d.parser_config["raptor"]
  488. cls.update_by_id(id, {"parser_config": d.parser_config})
  489. @classmethod
  490. @DB.connection_context()
  491. def get_doc_count(cls, tenant_id):
  492. docs = cls.model.select(cls.model.id).join(Knowledgebase,
  493. on=(Knowledgebase.id == cls.model.kb_id)).where(
  494. Knowledgebase.tenant_id == tenant_id)
  495. return len(docs)
  496. @classmethod
  497. @DB.connection_context()
  498. def begin2parse(cls, docid):
  499. cls.update_by_id(
  500. docid, {"progress": random.random() * 1 / 100.,
  501. "progress_msg": "Task is queued...",
  502. "process_begin_at": get_format_time()
  503. })
  504. @classmethod
  505. @DB.connection_context()
  506. def update_meta_fields(cls, doc_id, meta_fields):
  507. return cls.update_by_id(doc_id, {"meta_fields": meta_fields})
  508. @classmethod
  509. @DB.connection_context()
  510. def update_progress(cls):
  511. docs = cls.get_unfinished_docs()
  512. for d in docs:
  513. try:
  514. tsks = Task.query(doc_id=d["id"], order_by=Task.create_time)
  515. if not tsks:
  516. continue
  517. msg = []
  518. prg = 0
  519. finished = True
  520. bad = 0
  521. has_raptor = False
  522. has_graphrag = False
  523. e, doc = DocumentService.get_by_id(d["id"])
  524. status = doc.run # TaskStatus.RUNNING.value
  525. priority = 0
  526. for t in tsks:
  527. if 0 <= t.progress < 1:
  528. finished = False
  529. if t.progress == -1:
  530. bad += 1
  531. prg += t.progress if t.progress >= 0 else 0
  532. if t.progress_msg.strip():
  533. msg.append(t.progress_msg)
  534. if t.task_type == "raptor":
  535. has_raptor = True
  536. elif t.task_type == "graphrag":
  537. has_graphrag = True
  538. priority = max(priority, t.priority)
  539. prg /= len(tsks)
  540. if finished and bad:
  541. prg = -1
  542. status = TaskStatus.FAIL.value
  543. elif finished:
  544. if (d["parser_config"].get("raptor") or {}).get("use_raptor") and not has_raptor:
  545. queue_raptor_o_graphrag_tasks(d, "raptor", priority)
  546. prg = 0.98 * len(tsks) / (len(tsks) + 1)
  547. elif (d["parser_config"].get("graphrag") or {}).get("use_graphrag") and not has_graphrag:
  548. queue_raptor_o_graphrag_tasks(d, "graphrag", priority)
  549. prg = 0.98 * len(tsks) / (len(tsks) + 1)
  550. else:
  551. status = TaskStatus.DONE.value
  552. msg = "\n".join(sorted(msg))
  553. info = {
  554. "process_duration": datetime.timestamp(
  555. datetime.now()) -
  556. d["process_begin_at"].timestamp(),
  557. "run": status}
  558. if prg != 0:
  559. info["progress"] = prg
  560. if msg:
  561. info["progress_msg"] = msg
  562. if msg.endswith("created task graphrag") or msg.endswith("created task raptor"):
  563. info["progress_msg"] += "\n%d tasks are ahead in the queue..."%get_queue_length(priority)
  564. else:
  565. info["progress_msg"] = "%d tasks are ahead in the queue..."%get_queue_length(priority)
  566. cls.update_by_id(d["id"], info)
  567. except Exception as e:
  568. if str(e).find("'0'") < 0:
  569. logging.exception("fetch task exception")
  570. @classmethod
  571. @DB.connection_context()
  572. def get_kb_doc_count(cls, kb_id):
  573. return len(cls.model.select(cls.model.id).where(
  574. cls.model.kb_id == kb_id).dicts())
  575. @classmethod
  576. @DB.connection_context()
  577. def do_cancel(cls, doc_id):
  578. try:
  579. _, doc = DocumentService.get_by_id(doc_id)
  580. return doc.run == TaskStatus.CANCEL.value or doc.progress < 0
  581. except Exception:
  582. pass
  583. return False
  584. def queue_raptor_o_graphrag_tasks(doc, ty, priority):
  585. chunking_config = DocumentService.get_chunking_config(doc["id"])
  586. hasher = xxhash.xxh64()
  587. for field in sorted(chunking_config.keys()):
  588. hasher.update(str(chunking_config[field]).encode("utf-8"))
  589. def new_task():
  590. nonlocal doc
  591. return {
  592. "id": get_uuid(),
  593. "doc_id": doc["id"],
  594. "from_page": 100000000,
  595. "to_page": 100000000,
  596. "task_type": ty,
  597. "progress_msg": datetime.now().strftime("%H:%M:%S") + " created task " + ty
  598. }
  599. task = new_task()
  600. for field in ["doc_id", "from_page", "to_page"]:
  601. hasher.update(str(task.get(field, "")).encode("utf-8"))
  602. hasher.update(ty.encode("utf-8"))
  603. task["digest"] = hasher.hexdigest()
  604. bulk_insert_into_db(Task, [task], True)
  605. assert REDIS_CONN.queue_product(get_svr_queue_name(priority), message=task), "Can't access Redis. Please check the Redis' status."
  606. def get_queue_length(priority):
  607. group_info = REDIS_CONN.queue_info(get_svr_queue_name(priority), SVR_CONSUMER_GROUP_NAME)
  608. return int(group_info.get("lag", 0))
  609. def doc_upload_and_parse(conversation_id, file_objs, user_id):
  610. from api.db.services.api_service import API4ConversationService
  611. from api.db.services.conversation_service import ConversationService
  612. from api.db.services.dialog_service import DialogService
  613. from api.db.services.file_service import FileService
  614. from api.db.services.llm_service import LLMBundle
  615. from api.db.services.user_service import TenantService
  616. from rag.app import audio, email, naive, picture, presentation
  617. e, conv = ConversationService.get_by_id(conversation_id)
  618. if not e:
  619. e, conv = API4ConversationService.get_by_id(conversation_id)
  620. assert e, "Conversation not found!"
  621. e, dia = DialogService.get_by_id(conv.dialog_id)
  622. if not dia.kb_ids:
  623. raise LookupError("No knowledge base associated with this conversation. "
  624. "Please add a knowledge base before uploading documents")
  625. kb_id = dia.kb_ids[0]
  626. e, kb = KnowledgebaseService.get_by_id(kb_id)
  627. if not e:
  628. raise LookupError("Can't find this knowledgebase!")
  629. embd_mdl = LLMBundle(kb.tenant_id, LLMType.EMBEDDING, llm_name=kb.embd_id, lang=kb.language)
  630. err, files = FileService.upload_document(kb, file_objs, user_id)
  631. assert not err, "\n".join(err)
  632. def dummy(prog=None, msg=""):
  633. pass
  634. FACTORY = {
  635. ParserType.PRESENTATION.value: presentation,
  636. ParserType.PICTURE.value: picture,
  637. ParserType.AUDIO.value: audio,
  638. ParserType.EMAIL.value: email
  639. }
  640. parser_config = {"chunk_token_num": 4096, "delimiter": "\n!?;。;!?", "layout_recognize": "Plain Text"}
  641. exe = ThreadPoolExecutor(max_workers=12)
  642. threads = []
  643. doc_nm = {}
  644. for d, blob in files:
  645. doc_nm[d["id"]] = d["name"]
  646. for d, blob in files:
  647. kwargs = {
  648. "callback": dummy,
  649. "parser_config": parser_config,
  650. "from_page": 0,
  651. "to_page": 100000,
  652. "tenant_id": kb.tenant_id,
  653. "lang": kb.language
  654. }
  655. threads.append(exe.submit(FACTORY.get(d["parser_id"], naive).chunk, d["name"], blob, **kwargs))
  656. for (docinfo, _), th in zip(files, threads):
  657. docs = []
  658. doc = {
  659. "doc_id": docinfo["id"],
  660. "kb_id": [kb.id]
  661. }
  662. for ck in th.result():
  663. d = deepcopy(doc)
  664. d.update(ck)
  665. d["id"] = xxhash.xxh64((ck["content_with_weight"] + str(d["doc_id"])).encode("utf-8")).hexdigest()
  666. d["create_time"] = str(datetime.now()).replace("T", " ")[:19]
  667. d["create_timestamp_flt"] = datetime.now().timestamp()
  668. if not d.get("image"):
  669. docs.append(d)
  670. continue
  671. output_buffer = BytesIO()
  672. if isinstance(d["image"], bytes):
  673. output_buffer = BytesIO(d["image"])
  674. else:
  675. d["image"].save(output_buffer, format='JPEG')
  676. STORAGE_IMPL.put(kb.id, d["id"], output_buffer.getvalue())
  677. d["img_id"] = "{}-{}".format(kb.id, d["id"])
  678. d.pop("image", None)
  679. docs.append(d)
  680. parser_ids = {d["id"]: d["parser_id"] for d, _ in files}
  681. docids = [d["id"] for d, _ in files]
  682. chunk_counts = {id: 0 for id in docids}
  683. token_counts = {id: 0 for id in docids}
  684. es_bulk_size = 64
  685. def embedding(doc_id, cnts, batch_size=16):
  686. nonlocal embd_mdl, chunk_counts, token_counts
  687. vects = []
  688. for i in range(0, len(cnts), batch_size):
  689. vts, c = embd_mdl.encode(cnts[i: i + batch_size])
  690. vects.extend(vts.tolist())
  691. chunk_counts[doc_id] += len(cnts[i:i + batch_size])
  692. token_counts[doc_id] += c
  693. return vects
  694. idxnm = search.index_name(kb.tenant_id)
  695. try_create_idx = True
  696. _, tenant = TenantService.get_by_id(kb.tenant_id)
  697. llm_bdl = LLMBundle(kb.tenant_id, LLMType.CHAT, tenant.llm_id)
  698. for doc_id in docids:
  699. cks = [c for c in docs if c["doc_id"] == doc_id]
  700. if parser_ids[doc_id] != ParserType.PICTURE.value:
  701. from graphrag.general.mind_map_extractor import MindMapExtractor
  702. mindmap = MindMapExtractor(llm_bdl)
  703. try:
  704. mind_map = trio.run(mindmap, [c["content_with_weight"] for c in docs if c["doc_id"] == doc_id])
  705. mind_map = json.dumps(mind_map.output, ensure_ascii=False, indent=2)
  706. if len(mind_map) < 32:
  707. raise Exception("Few content: " + mind_map)
  708. cks.append({
  709. "id": get_uuid(),
  710. "doc_id": doc_id,
  711. "kb_id": [kb.id],
  712. "docnm_kwd": doc_nm[doc_id],
  713. "title_tks": rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", doc_nm[doc_id])),
  714. "content_ltks": rag_tokenizer.tokenize("summary summarize 总结 概况 file 文件 概括"),
  715. "content_with_weight": mind_map,
  716. "knowledge_graph_kwd": "mind_map"
  717. })
  718. except Exception as e:
  719. logging.exception("Mind map generation error")
  720. vects = embedding(doc_id, [c["content_with_weight"] for c in cks])
  721. assert len(cks) == len(vects)
  722. for i, d in enumerate(cks):
  723. v = vects[i]
  724. d["q_%d_vec" % len(v)] = v
  725. for b in range(0, len(cks), es_bulk_size):
  726. if try_create_idx:
  727. if not settings.docStoreConn.indexExist(idxnm, kb_id):
  728. settings.docStoreConn.createIdx(idxnm, kb_id, len(vects[0]))
  729. try_create_idx = False
  730. settings.docStoreConn.insert(cks[b:b + es_bulk_size], idxnm, kb_id)
  731. DocumentService.increment_chunk_num(
  732. doc_id, kb.id, token_counts[doc_id], chunk_counts[doc_id], 0)
  733. return [d["id"] for d, _ in files]