lightrag.py 174 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799380038013802380338043805380638073808380938103811381238133814381538163817381838193820382138223823382438253826382738283829383038313832383338343835383638373838383938403841384238433844384538463847384838493850385138523853385438553856385738583859386038613862386338643865386638673868386938703871387238733874387538763877387838793880388138823883388438853886388738883889389038913892389338943895389638973898389939003901390239033904390539063907390839093910391139123913391439153916391739183919392039213922392339243925392639273928392939303931393239333934393539363937393839393940394139423943394439453946394739483949395039513952395339543955395639573958395939603961396239633964396539663967396839693970397139723973397439753976397739783979398039813982398339843985398639873988398939903991399239933994399539963997399839994000400140024003400440054006400740084009401040114012401340144015401640174018401940204021402240234024402540264027402840294030403140324033403440354036403740384039404040414042404340444045404640474048404940504051405240534054405540564057405840594060406140624063406440654066406740684069407040714072407340744075407640774078407940804081408240834084408540864087408840894090
  1. from __future__ import annotations
  2. import traceback
  3. import asyncio
  4. import os
  5. import time
  6. import warnings
  7. from copy import deepcopy
  8. try:
  9. import httpx
  10. except Exception: # pragma: no cover - optional dependency
  11. httpx = None
  12. from dataclasses import InitVar, asdict, dataclass, field, replace
  13. from datetime import datetime, timezone
  14. from functools import partial
  15. from typing import (
  16. Any,
  17. AsyncIterator,
  18. Awaitable,
  19. Callable,
  20. Iterator,
  21. cast,
  22. final,
  23. Literal,
  24. Mapping,
  25. Optional,
  26. List,
  27. Dict,
  28. Union,
  29. )
  30. from lightrag.prompt import (
  31. PROMPTS,
  32. get_default_entity_extraction_prompt_profile,
  33. resolve_entity_extraction_prompt_profile,
  34. validate_entity_extraction_prompt_profile_for_mode,
  35. )
  36. from lightrag.constants import (
  37. DEFAULT_CHUNK_P_SIZE,
  38. DEFAULT_MAX_GLEANING,
  39. DEFAULT_MAX_EXTRACTION_RECORDS,
  40. DEFAULT_MAX_EXTRACTION_ENTITIES,
  41. DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE,
  42. DEFAULT_TOP_K,
  43. DEFAULT_CHUNK_TOP_K,
  44. DEFAULT_MAX_ENTITY_TOKENS,
  45. DEFAULT_MAX_RELATION_TOKENS,
  46. DEFAULT_MAX_TOTAL_TOKENS,
  47. DEFAULT_COSINE_THRESHOLD,
  48. DEFAULT_RELATED_CHUNK_NUMBER,
  49. DEFAULT_KG_CHUNK_PICK_METHOD,
  50. DEFAULT_MIN_RERANK_SCORE,
  51. DEFAULT_SUMMARY_MAX_TOKENS,
  52. DEFAULT_SUMMARY_CONTEXT_SIZE,
  53. DEFAULT_SUMMARY_LENGTH_RECOMMENDED,
  54. DEFAULT_MAX_ASYNC,
  55. DEFAULT_MAX_PARALLEL_INSERT,
  56. DEFAULT_MAX_GRAPH_NODES,
  57. DEFAULT_MAX_SOURCE_IDS_PER_ENTITY,
  58. DEFAULT_MAX_SOURCE_IDS_PER_RELATION,
  59. DEFAULT_SUMMARY_LANGUAGE,
  60. DEFAULT_LLM_TIMEOUT,
  61. DEFAULT_EMBEDDING_TIMEOUT,
  62. DEFAULT_RERANK_TIMEOUT,
  63. DEFAULT_SOURCE_IDS_LIMIT_METHOD,
  64. DEFAULT_MAX_FILE_PATHS,
  65. DEFAULT_MAX_PARALLEL_ANALYZE,
  66. DEFAULT_MAX_PARALLEL_PARSE_NATIVE,
  67. DEFAULT_MAX_PARALLEL_PARSE_MINERU,
  68. DEFAULT_MAX_PARALLEL_PARSE_DOCLING,
  69. DEFAULT_QUEUE_SIZE_DEFAULT,
  70. DEFAULT_QUEUE_SIZE_INSERT,
  71. DEFAULT_FILE_PATH_MORE_PLACEHOLDER,
  72. )
  73. from lightrag.utils import get_env_value
  74. from lightrag.kg import (
  75. verify_storage_implementation,
  76. )
  77. from lightrag.kg.shared_storage import (
  78. get_namespace_data,
  79. get_default_workspace,
  80. set_default_workspace,
  81. get_namespace_lock,
  82. get_storage_keyed_lock,
  83. )
  84. from lightrag.base import (
  85. BaseGraphStorage,
  86. BaseKVStorage,
  87. BaseVectorStorage,
  88. DocProcessingStatus,
  89. DocStatus,
  90. DocStatusStorage,
  91. QueryParam,
  92. StorageNameSpace,
  93. StoragesStatus,
  94. DeletionResult,
  95. OllamaServerInfos,
  96. QueryResult,
  97. )
  98. from lightrag.namespace import NameSpace
  99. from lightrag.chunker import chunking_by_token_size
  100. from lightrag.operate import (
  101. extract_entities,
  102. kg_query,
  103. naive_query,
  104. rebuild_knowledge_from_chunks,
  105. )
  106. from lightrag.utils_pipeline import normalize_document_file_path
  107. from lightrag.constants import GRAPH_FIELD_SEP
  108. from lightrag.utils import (
  109. Tokenizer,
  110. TiktokenTokenizer,
  111. EmbeddingFunc,
  112. always_get_an_event_loop,
  113. compute_mdhash_id,
  114. priority_limit_async_func_call,
  115. sanitize_text_for_encoding,
  116. check_storage_env_vars,
  117. generate_track_id,
  118. convert_to_user_format,
  119. logger,
  120. make_relation_vdb_ids,
  121. subtract_source_ids,
  122. make_relation_chunk_key,
  123. normalize_source_ids_limit_method,
  124. normalize_string_list,
  125. )
  126. from lightrag.types import KnowledgeGraph
  127. from dotenv import load_dotenv
  128. from lightrag.pipeline import _PipelineMixin
  129. from lightrag.kg.factory import get_storage_class
  130. from lightrag.addon_params import (
  131. ObservableAddonParams,
  132. normalize_addon_params,
  133. )
  134. from lightrag.llm_roles import (
  135. ROLE_NAMES,
  136. ROLES,
  137. RoleLLMConfig,
  138. RoleSpec, # noqa: F401 # re-exported via lightrag/__init__.py
  139. _optional_env_int,
  140. _RoleLLMMixin,
  141. _RoleLLMState,
  142. )
  143. from lightrag.storage_migrations import _StorageMigrationMixin
  144. # use the .env that is inside the current folder
  145. # allows to use different .env file for each lightrag instance
  146. # the OS environment variables take precedence over the .env file
  147. load_dotenv(dotenv_path=".env", override=False)
  148. @final
  149. @dataclass
  150. class LightRAG(_RoleLLMMixin, _StorageMigrationMixin, _PipelineMixin):
  151. """LightRAG: Simple and Fast Retrieval-Augmented Generation."""
  152. # Directory
  153. # ---
  154. working_dir: str = field(default="./rag_storage")
  155. """Directory where cache and temporary files are stored."""
  156. # Storage
  157. # ---
  158. kv_storage: str = field(default="JsonKVStorage")
  159. """Storage backend for key-value data."""
  160. vector_storage: str = field(default="NanoVectorDBStorage")
  161. """Storage backend for vector embeddings."""
  162. graph_storage: str = field(default="NetworkXStorage")
  163. """Storage backend for knowledge graphs."""
  164. doc_status_storage: str = field(default="JsonDocStatusStorage")
  165. """Storage type for tracking document processing statuses."""
  166. # Workspace
  167. # ---
  168. workspace: str = field(default_factory=lambda: os.getenv("WORKSPACE", ""))
  169. """Workspace for data isolation. Defaults to empty string if WORKSPACE environment variable is not set."""
  170. # ---
  171. # TODO: Deprecated, use setup_logger in utils.py instead
  172. log_level: int | None = field(default=None)
  173. log_file_path: str | None = field(default=None)
  174. # Query parameters
  175. # ---
  176. top_k: int = field(default=get_env_value("TOP_K", DEFAULT_TOP_K, int))
  177. """Number of entities/relations to retrieve for each query."""
  178. chunk_top_k: int = field(
  179. default=get_env_value("CHUNK_TOP_K", DEFAULT_CHUNK_TOP_K, int)
  180. )
  181. """Maximum number of chunks in context."""
  182. max_entity_tokens: int = field(
  183. default=get_env_value("MAX_ENTITY_TOKENS", DEFAULT_MAX_ENTITY_TOKENS, int)
  184. )
  185. """Maximum number of tokens for entity in context."""
  186. max_relation_tokens: int = field(
  187. default=get_env_value("MAX_RELATION_TOKENS", DEFAULT_MAX_RELATION_TOKENS, int)
  188. )
  189. """Maximum number of tokens for relation in context."""
  190. max_total_tokens: int = field(
  191. default=get_env_value("MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS, int)
  192. )
  193. """Maximum total tokens in context (including system prompt, entities, relations and chunks)."""
  194. cosine_threshold: int = field(
  195. default=get_env_value("COSINE_THRESHOLD", DEFAULT_COSINE_THRESHOLD, int)
  196. )
  197. """Cosine threshold of vector DB retrieval for entities, relations and chunks."""
  198. related_chunk_number: int = field(
  199. default=get_env_value("RELATED_CHUNK_NUMBER", DEFAULT_RELATED_CHUNK_NUMBER, int)
  200. )
  201. """Number of related chunks to grab from single entity or relation."""
  202. kg_chunk_pick_method: str = field(
  203. default=get_env_value("KG_CHUNK_PICK_METHOD", DEFAULT_KG_CHUNK_PICK_METHOD, str)
  204. )
  205. """Method for selecting text chunks: 'WEIGHT' for weight-based selection, 'VECTOR' for embedding similarity-based selection."""
  206. # Entity extraction
  207. # ---
  208. entity_extract_max_gleaning: int = field(
  209. default=get_env_value("MAX_GLEANING", DEFAULT_MAX_GLEANING, int)
  210. )
  211. """Maximum number of entity extraction attempts for ambiguous content."""
  212. entity_extract_max_records: int = field(
  213. default=get_env_value(
  214. "MAX_EXTRACTION_RECORDS", DEFAULT_MAX_EXTRACTION_RECORDS, int
  215. )
  216. )
  217. """Per-response cap on total entity+relationship rows/records."""
  218. entity_extract_max_entities: int = field(
  219. default=get_env_value(
  220. "MAX_EXTRACTION_ENTITIES", DEFAULT_MAX_EXTRACTION_ENTITIES, int
  221. )
  222. )
  223. """Per-response cap on entity rows/objects."""
  224. force_llm_summary_on_merge: int = field(
  225. default=get_env_value(
  226. "FORCE_LLM_SUMMARY_ON_MERGE", DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE, int
  227. )
  228. )
  229. # Text chunking
  230. # ---
  231. chunk_token_size: int | None = field(default=None)
  232. """Maximum number of tokens per text chunk when splitting documents.
  233. ``None`` means "use ``addon_params['chunker']['chunk_token_size']``"
  234. (env-driven via ``CHUNK_SIZE``). When the constructor is given a
  235. non-None value it overlays onto ``addon_params['chunker']`` in
  236. ``__post_init__`` so the per-document ``chunk_options`` snapshot
  237. actually picks it up. Always an ``int`` after construction —
  238. back-filled from the resolved chunker config so legacy readers
  239. (``self.chunk_token_size``) keep working."""
  240. chunk_overlap_token_size: int | None = field(default=None)
  241. """Number of overlapping tokens between consecutive text chunks (F-strategy semantics).
  242. ``None`` means "use the per-strategy default in
  243. ``addon_params['chunker']``" (env-driven via
  244. ``CHUNK_F_OVERLAP_SIZE`` / ``CHUNK_R_OVERLAP_SIZE`` falling back to
  245. ``CHUNK_OVERLAP_SIZE``). When non-None at construction time, the
  246. value overlays onto every strategy sub-dict that natively takes
  247. ``chunk_overlap_token_size`` (``fixed_token``, ``recursive_character``)
  248. so the per-doc snapshot reflects the constructor choice. Per-strategy
  249. chunker parameters (R / V separators, thresholds, overlap overrides,
  250. etc.) live in ``addon_params['chunker']`` and are documented in
  251. :func:`lightrag.parser.routing.default_chunker_config`. Per-doc
  252. snapshots are persisted to ``full_docs[doc_id]['chunk_options']``
  253. at enqueue time."""
  254. tokenizer: Optional[Tokenizer] = field(default=None)
  255. """
  256. A function that returns a Tokenizer instance.
  257. If None, and a `tiktoken_model_name` is provided, a TiktokenTokenizer will be created.
  258. If both are None, the default TiktokenTokenizer is used.
  259. """
  260. tiktoken_model_name: str = field(default="gpt-4o-mini")
  261. """Model name used for tokenization when chunking text with tiktoken. Defaults to `gpt-4o-mini`."""
  262. chunking_func: Callable[
  263. [
  264. Tokenizer,
  265. str,
  266. Optional[str],
  267. bool,
  268. int,
  269. int,
  270. ],
  271. Union[List[Dict[str, Any]], Awaitable[List[Dict[str, Any]]]],
  272. ] = field(default_factory=lambda: chunking_by_token_size)
  273. """
  274. Legacy chunking-function customization point. Synchronous or async.
  275. **When this function is actually invoked.** The chunker dispatch in
  276. ``_PipelineMixin.process_single_document`` is driven by the
  277. document's ``process_options``:
  278. - If ``process_options`` explicitly contains a chunking selector
  279. char (``F``/``R``/``V``/``P``), the dispatcher routes to a
  280. chunker that follows the new file-chunker contract — see
  281. :mod:`lightrag.chunker` (``chunking_by_fixed_token`` for ``F``,
  282. ``chunking_by_paragraph_semantic`` for ``P``; ``R``/``V`` are
  283. not yet implemented and fall back to ``F``). **This
  284. ``chunking_func`` is NOT called in that case** — it is a
  285. legacy escape hatch and is intentionally bypassed when the user
  286. opted into a specific strategy.
  287. - If ``process_options`` does **not** name a chunking strategy
  288. (empty string, or only non-chunking flags such as ``i`` / ``t``
  289. / ``e`` / ``!``), the dispatcher invokes this ``chunking_func``
  290. with the legacy 6-arg signature below. This is the path taken
  291. by direct ``ainsert(text)`` calls and by any document whose
  292. ``process_options`` simply does not select a chunker.
  293. The presence/absence of the selector is exposed by
  294. :attr:`lightrag.parser.routing.ProcessOptions.chunking_explicit`.
  295. **Signature** — preserved unchanged from earlier LightRAG releases
  296. so externally-supplied chunkers continue to drop in without edits:
  297. - `tokenizer`: A Tokenizer instance to use for tokenization.
  298. - `content`: The text to be split into chunks.
  299. - `split_by_character`: The character to split the text on. If
  300. None, the text is split into chunks of `chunk_token_size`
  301. tokens.
  302. - `split_by_character_only`: If True, the text is split only on
  303. the specified character.
  304. - `chunk_overlap_token_size`: The number of overlapping tokens
  305. between consecutive chunks.
  306. - `chunk_token_size`: The maximum number of tokens per chunk.
  307. The function should return a list of dictionaries (or an awaitable
  308. that resolves to one), each containing:
  309. - `tokens` (int): The number of tokens in the chunk.
  310. - `content` (str): The text content of the chunk.
  311. - `chunk_order_index` (int): Zero-based index indicating the
  312. chunk's order in the document.
  313. Defaults to :func:`lightrag.chunker.chunking_by_token_size`.
  314. """
  315. # Embedding
  316. # ---
  317. embedding_func: EmbeddingFunc | None = field(default=None)
  318. """Function for computing text embeddings. Must be set before use."""
  319. embedding_token_limit: int | None = field(default=None, init=False)
  320. """Token limit for embedding model. Set automatically from embedding_func.max_token_size in __post_init__."""
  321. embedding_batch_num: int = field(default=int(os.getenv("EMBEDDING_BATCH_NUM", 10)))
  322. """Batch size for embedding computations."""
  323. embedding_func_max_async: int = field(
  324. default=int(os.getenv("EMBEDDING_FUNC_MAX_ASYNC", 8))
  325. )
  326. """Maximum number of concurrent embedding function calls."""
  327. embedding_cache_config: dict[str, Any] = field(
  328. default_factory=lambda: {
  329. "enabled": False,
  330. "similarity_threshold": 0.95,
  331. "use_llm_check": False,
  332. }
  333. )
  334. """Configuration for embedding cache.
  335. - enabled: If True, enables caching to avoid redundant computations.
  336. - similarity_threshold: Minimum similarity score to use cached embeddings.
  337. - use_llm_check: If True, validates cached embeddings using an LLM.
  338. """
  339. default_embedding_timeout: int = field(
  340. default=int(os.getenv("EMBEDDING_TIMEOUT", DEFAULT_EMBEDDING_TIMEOUT))
  341. )
  342. # LLM Configuration
  343. # ---
  344. llm_model_func: Callable[..., object] | None = field(default=None)
  345. """Function for interacting with the large language model (LLM). Must be set before use."""
  346. role_llm_configs: dict[str, RoleLLMConfig | dict[str, Any]] | None = field(
  347. default=None
  348. )
  349. """Per-role LLM overrides keyed by role name (see :data:`ROLES`).
  350. Each entry is a :class:`RoleLLMConfig` (or a plain dict with the same
  351. keys ``func`` / ``kwargs`` / ``max_async`` / ``timeout``). Any field left
  352. as ``None`` falls back to the corresponding base LLM setting. Roles not
  353. present in the dict are wrapped from the base ``llm_model_func`` and
  354. pick up ``{ROLE_PREFIX}_MAX_ASYNC_LLM`` env defaults."""
  355. llm_model_name: str = field(default="gpt-4o-mini")
  356. """Name of the LLM model used for generating responses."""
  357. summary_max_tokens: int = field(
  358. default=int(os.getenv("SUMMARY_MAX_TOKENS", DEFAULT_SUMMARY_MAX_TOKENS))
  359. )
  360. """Maximum tokens allowed for entity/relation description."""
  361. summary_context_size: int = field(
  362. default=int(os.getenv("SUMMARY_CONTEXT_SIZE", DEFAULT_SUMMARY_CONTEXT_SIZE))
  363. )
  364. """Maximum number of tokens allowed per LLM response."""
  365. summary_length_recommended: int = field(
  366. default=int(
  367. os.getenv("SUMMARY_LENGTH_RECOMMENDED", DEFAULT_SUMMARY_LENGTH_RECOMMENDED)
  368. )
  369. )
  370. """Recommended length of LLM summary output."""
  371. llm_model_max_async: int = field(
  372. default=int(os.getenv("MAX_ASYNC", DEFAULT_MAX_ASYNC))
  373. )
  374. """Maximum number of concurrent LLM calls."""
  375. llm_model_kwargs: dict[str, Any] = field(default_factory=dict)
  376. """Additional keyword arguments passed to the LLM model function."""
  377. default_llm_timeout: int = field(
  378. default=int(os.getenv("LLM_TIMEOUT", DEFAULT_LLM_TIMEOUT))
  379. )
  380. entity_extraction_use_json: bool = field(
  381. default=os.getenv("ENTITY_EXTRACTION_USE_JSON", "false").lower() == "true"
  382. )
  383. """When True, entity extraction uses JSON structured output instead of delimiter-based text.
  384. JSON mode is slower but significantly improves extraction quality and compatibility with smaller models.
  385. Providers with native structured output support (OpenAI, Ollama, Gemini) will use their
  386. native capabilities. Other providers rely on JSON-formatted prompts with json_repair parsing.
  387. Default: False. Set ENTITY_EXTRACTION_USE_JSON=true in .env to enable."""
  388. # Rerank Configuration
  389. # ---
  390. rerank_model_func: Callable[..., object] | None = field(default=None)
  391. """Function for reranking retrieved documents. All rerank configurations (model name, API keys, top_k, etc.) should be included in this function. Optional."""
  392. rerank_model_max_async: int = field(
  393. default=int(
  394. os.getenv(
  395. "MAX_ASYNC_RERANK",
  396. os.getenv("MAX_ASYNC", DEFAULT_MAX_ASYNC),
  397. )
  398. )
  399. )
  400. """Maximum number of concurrent rerank calls.
  401. Falls back to MAX_ASYNC when MAX_ASYNC_RERANK is unset."""
  402. default_rerank_timeout: int = field(
  403. default=int(os.getenv("RERANK_TIMEOUT", DEFAULT_RERANK_TIMEOUT))
  404. )
  405. """Rerank request timeout in seconds.
  406. Independent from LLM_TIMEOUT since reranker calls are much shorter
  407. than full LLM generation."""
  408. min_rerank_score: float = field(
  409. default=get_env_value("MIN_RERANK_SCORE", DEFAULT_MIN_RERANK_SCORE, float)
  410. )
  411. """Minimum rerank score threshold for filtering chunks after reranking."""
  412. # Storage
  413. # ---
  414. vector_db_storage_cls_kwargs: dict[str, Any] = field(default_factory=dict)
  415. """Additional parameters for vector database storage."""
  416. enable_llm_cache: bool = field(default=True)
  417. """Enables caching for LLM responses to avoid redundant computations."""
  418. enable_llm_cache_for_entity_extract: bool = field(default=True)
  419. """If True, enables caching for entity extraction steps to reduce LLM costs."""
  420. vlm_process_enable: bool = field(
  421. default_factory=lambda: get_env_value("VLM_PROCESS_ENABLE", False, bool)
  422. )
  423. """Master switch for VLM multimodal analysis (i/t/e items).
  424. When False, the pipeline emits a warning and skips every multimodal item
  425. without invoking the VLM. When True, the configured VLM binding must
  426. support image inputs.
  427. """
  428. # Extensions
  429. # ---
  430. max_parallel_insert: int = field(
  431. default=int(os.getenv("MAX_PARALLEL_INSERT", DEFAULT_MAX_PARALLEL_INSERT))
  432. )
  433. """Maximum number of parallel insert operations."""
  434. max_parallel_parse_native: int = field(
  435. default=int(
  436. os.getenv(
  437. "MAX_PARALLEL_PARSE_NATIVE", str(DEFAULT_MAX_PARALLEL_PARSE_NATIVE)
  438. )
  439. )
  440. )
  441. max_parallel_parse_mineru: int = field(
  442. default=int(
  443. os.getenv(
  444. "MAX_PARALLEL_PARSE_MINERU", str(DEFAULT_MAX_PARALLEL_PARSE_MINERU)
  445. )
  446. )
  447. )
  448. max_parallel_parse_docling: int = field(
  449. default=int(
  450. os.getenv(
  451. "MAX_PARALLEL_PARSE_DOCLING", str(DEFAULT_MAX_PARALLEL_PARSE_DOCLING)
  452. )
  453. )
  454. )
  455. max_parallel_analyze: int = field(
  456. default=int(
  457. os.getenv("MAX_PARALLEL_ANALYZE", str(DEFAULT_MAX_PARALLEL_ANALYZE))
  458. )
  459. )
  460. queue_size_default: int = field(
  461. default=int(os.getenv("QUEUE_SIZE_DEFAULT", str(DEFAULT_QUEUE_SIZE_DEFAULT)))
  462. )
  463. queue_size_insert: int = field(
  464. default=int(os.getenv("QUEUE_SIZE_INSERT", str(DEFAULT_QUEUE_SIZE_INSERT)))
  465. )
  466. max_graph_nodes: int = field(
  467. default=get_env_value("MAX_GRAPH_NODES", DEFAULT_MAX_GRAPH_NODES, int)
  468. )
  469. """Maximum number of graph nodes to return in knowledge graph queries."""
  470. max_source_ids_per_entity: int = field(
  471. default=get_env_value(
  472. "MAX_SOURCE_IDS_PER_ENTITY", DEFAULT_MAX_SOURCE_IDS_PER_ENTITY, int
  473. )
  474. )
  475. """Maximum number of source (chunk) ids in entity Grpah + VDB."""
  476. max_source_ids_per_relation: int = field(
  477. default=get_env_value(
  478. "MAX_SOURCE_IDS_PER_RELATION",
  479. DEFAULT_MAX_SOURCE_IDS_PER_RELATION,
  480. int,
  481. )
  482. )
  483. """Maximum number of source (chunk) ids in relation Graph + VDB."""
  484. source_ids_limit_method: str = field(
  485. default_factory=lambda: normalize_source_ids_limit_method(
  486. get_env_value(
  487. "SOURCE_IDS_LIMIT_METHOD",
  488. DEFAULT_SOURCE_IDS_LIMIT_METHOD,
  489. str,
  490. )
  491. )
  492. )
  493. """Strategy for enforcing source_id limits: IGNORE_NEW or FIFO."""
  494. max_file_paths: int = field(
  495. default=get_env_value("MAX_FILE_PATHS", DEFAULT_MAX_FILE_PATHS, int)
  496. )
  497. """Maximum number of file paths to store in entity/relation file_path field."""
  498. file_path_more_placeholder: str = field(default=DEFAULT_FILE_PATH_MORE_PLACEHOLDER)
  499. """Placeholder text when file paths exceed max_file_paths limit."""
  500. addon_params: InitVar[dict[str, Any] | None] = None
  501. _addon_params: ObservableAddonParams = field(
  502. default_factory=ObservableAddonParams,
  503. init=False,
  504. repr=False,
  505. )
  506. _addon_params_dirty: bool = field(default=True, init=False, repr=False)
  507. _entity_extraction_prompt_profile: dict[str, Any] = field(
  508. default_factory=get_default_entity_extraction_prompt_profile,
  509. init=False,
  510. repr=False,
  511. )
  512. _cached_entity_extraction_use_json: bool | None = field(
  513. default=None,
  514. init=False,
  515. repr=False,
  516. )
  517. _resolved_summary_language: str = field(
  518. default=DEFAULT_SUMMARY_LANGUAGE,
  519. init=False,
  520. repr=False,
  521. )
  522. # Storages Management
  523. # ---
  524. # TODO: Deprecated (LightRAG will never initialize storage automatically on creation,and finalize should be call before destroying)
  525. auto_manage_storages_states: bool = field(default=False)
  526. """If True, lightrag will automatically calls initialize_storages and finalize_storages at the appropriate times."""
  527. cosine_better_than_threshold: float = field(
  528. default=float(os.getenv("COSINE_THRESHOLD", 0.2))
  529. )
  530. ollama_server_infos: Optional[OllamaServerInfos] = field(default=None)
  531. """Configuration for Ollama server information."""
  532. _storages_status: StoragesStatus = field(default=StoragesStatus.NOT_CREATED)
  533. def _mark_addon_params_dirty(self) -> None:
  534. self._addon_params_dirty = True
  535. def _replace_addon_params(
  536. self, addon_params: Mapping[str, Any] | None, *, mark_dirty: bool
  537. ) -> None:
  538. wrapped = ObservableAddonParams(
  539. normalize_addon_params(addon_params),
  540. on_change=self._mark_addon_params_dirty,
  541. )
  542. self._addon_params = wrapped
  543. if mark_dirty:
  544. self._mark_addon_params_dirty()
  545. def _get_addon_params(self) -> ObservableAddonParams:
  546. """Return the live addon_params store.
  547. Mutations on the returned instance trigger a cache refresh on the next
  548. _build_global_config() call. If the whole mapping is replaced via the
  549. setter, previously captured references point at the old instance and
  550. will no longer propagate changes — always re-read `rag.addon_params`
  551. after replacement rather than caching references.
  552. """
  553. return self._addon_params
  554. def _set_runtime_addon_params(self, addon_params: Mapping[str, Any] | None) -> None:
  555. self._replace_addon_params(addon_params, mark_dirty=True)
  556. self._apply_chunk_size_overlay()
  557. def _apply_chunk_size_overlay(self) -> None:
  558. """Reconcile chunk-size config across all four configuration tiers.
  559. Specificity-ordered precedence (high → low) per slot:
  560. 1. ``addon_params['chunker']`` explicit (user-supplied dict that
  561. already carries the key).
  562. 2. Strategy-specific env (``CHUNK_F_SIZE`` / ``CHUNK_R_SIZE`` /
  563. ``CHUNK_V_SIZE`` for per-strategy ``chunk_token_size``;
  564. ``CHUNK_F_OVERLAP_SIZE`` / ``CHUNK_R_OVERLAP_SIZE`` /
  565. ``CHUNK_P_OVERLAP_SIZE`` for overlap). These are pre-filled into
  566. the strategy sub-dict by
  567. :func:`lightrag.parser.routing.default_chunker_config` when it
  568. builds the dict from scratch; for a *partial*
  569. ``addon_params['chunker']`` (which skips that builder) this overlay
  570. mirrors the size-env reads below so the env still applies. Either
  571. way the slot is filled *only* when the env var is set. No strategy
  572. env feeds the *top-level* ``chunk_token_size`` slot; that chain
  573. stays addon_params > legacy ctor > ``CHUNK_SIZE``.
  574. 3. Legacy constructor field
  575. (``LightRAG(chunk_token_size=…, chunk_overlap_token_size=…)``).
  576. Strategy-agnostic; only fills slots that were not already set
  577. by tiers 1–2.
  578. 4. Legacy env (``CHUNK_SIZE`` / ``CHUNK_OVERLAP_SIZE``) — final
  579. fallback.
  580. After this runs, ``self._addon_params['chunker']`` carries fully
  581. resolved values for every slot the pipeline needs, and the
  582. legacy ``self.chunk_token_size`` / ``self.chunk_overlap_token_size``
  583. instance fields are back-filled to ``int`` so downstream readers
  584. (e.g. ``process_single_document``'s
  585. ``chunk_opts.get("chunk_token_size") or self.chunk_token_size``
  586. fallback) keep working.
  587. """
  588. chunker_cfg = self._addon_params.get("chunker")
  589. if not isinstance(chunker_cfg, dict):
  590. chunker_cfg = {}
  591. self._addon_params["chunker"] = chunker_cfg
  592. # Top-level chunk_token_size — no strategy-specific env exists,
  593. # so the chain is: addon_params > legacy ctor > CHUNK_SIZE env.
  594. if "chunk_token_size" not in chunker_cfg:
  595. if self.chunk_token_size is not None:
  596. chunker_cfg["chunk_token_size"] = self.chunk_token_size
  597. else:
  598. chunker_cfg["chunk_token_size"] = int(os.getenv("CHUNK_SIZE", 1200))
  599. # Per-strategy chunk_overlap_token_size — strategy env (if set)
  600. # already lives in the sub-dict. Slots still missing fall back
  601. # to the legacy ctor field, then CHUNK_OVERLAP_SIZE env.
  602. if self.chunk_overlap_token_size is not None:
  603. legacy_overlap_default = self.chunk_overlap_token_size
  604. else:
  605. legacy_overlap_default = int(os.getenv("CHUNK_OVERLAP_SIZE", 100))
  606. for strategy_key in (
  607. "fixed_token",
  608. "recursive_character",
  609. "paragraph_semantic",
  610. ):
  611. sub = chunker_cfg.get(strategy_key)
  612. if not isinstance(sub, dict):
  613. sub = {}
  614. chunker_cfg[strategy_key] = sub
  615. if "chunk_overlap_token_size" not in sub:
  616. sub["chunk_overlap_token_size"] = legacy_overlap_default
  617. # P-specific chunk_token_size backfill — P does NOT inherit the
  618. # top-level chunk_token_size (CHUNK_SIZE / legacy ctor) when
  619. # nothing more specific was set; paragraph-semantic merging
  620. # needs more headroom than the global default to keep related
  621. # paragraphs together. ``default_chunker_config`` already
  622. # pre-fills this slot for the default-built chunker dict, but
  623. # when the caller hands us a partial ``addon_params['chunker']``
  624. # that lacks the slot (e.g. ``{"paragraph_semantic": {}}``)
  625. # ``normalize_addon_params`` does not re-run the defaults
  626. # builder — so this overlay is the last guard that ensures the
  627. # slot is always populated. Precedence (high → low):
  628. # explicit ``addon_params`` > ``CHUNK_P_SIZE`` env >
  629. # ``DEFAULT_CHUNK_P_SIZE``. ``setdefault`` preserves any
  630. # explicit value the caller did provide; the env read here
  631. # mirrors ``default_chunker_config`` so partial-addon-params
  632. # callers still pick up env overrides.
  633. p_size_raw = os.getenv("CHUNK_P_SIZE")
  634. chunker_cfg["paragraph_semantic"].setdefault(
  635. "chunk_token_size",
  636. int(p_size_raw) if p_size_raw is not None else DEFAULT_CHUNK_P_SIZE,
  637. )
  638. # Per-strategy F/R/V chunk_token_size from strategy env
  639. # (CHUNK_F_SIZE / CHUNK_R_SIZE / CHUNK_V_SIZE). Same rationale as the
  640. # P backfill above: ``default_chunker_config`` seeds these when it
  641. # builds the chunker dict from scratch, but a partial
  642. # ``addon_params['chunker']`` skips that builder
  643. # (``normalize_addon_params`` only defaults the whole ``chunker`` key
  644. # when it is absent), so this overlay is the last guard. Unlike P,
  645. # the slot is filled ONLY when the env is actually set — leaving it
  646. # absent otherwise so F/R/V inherit the top-level ``chunk_token_size``
  647. # at consumption time. ``setdefault`` preserves an explicit
  648. # caller-supplied value (tier 1 wins over the env tier 2).
  649. for strategy_key, size_env in (
  650. ("fixed_token", "CHUNK_F_SIZE"),
  651. ("recursive_character", "CHUNK_R_SIZE"),
  652. ("semantic_vector", "CHUNK_V_SIZE"),
  653. ):
  654. size_raw = os.getenv(size_env)
  655. if size_raw is None:
  656. continue
  657. sub = chunker_cfg.get(strategy_key)
  658. if not isinstance(sub, dict):
  659. sub = {}
  660. chunker_cfg[strategy_key] = sub
  661. sub.setdefault("chunk_token_size", int(size_raw))
  662. # Back-fill legacy instance fields → always int afterwards.
  663. # Overlap mirrors the F-strategy resolved value, matching the
  664. # F-flavoured legacy ``self.chunk_overlap_token_size`` semantics
  665. # used by the legacy 6-arg ``chunking_func`` path.
  666. self.chunk_token_size = chunker_cfg["chunk_token_size"]
  667. self.chunk_overlap_token_size = chunker_cfg["fixed_token"][
  668. "chunk_overlap_token_size"
  669. ]
  670. def _refresh_addon_params_cache(self) -> None:
  671. summary_language = self._addon_params.get("language", DEFAULT_SUMMARY_LANGUAGE)
  672. if not isinstance(summary_language, str) or not summary_language.strip():
  673. summary_language = DEFAULT_SUMMARY_LANGUAGE
  674. self._resolved_summary_language = summary_language
  675. resolved_prompt_profile = resolve_entity_extraction_prompt_profile(
  676. self._addon_params,
  677. self.entity_extraction_use_json,
  678. )
  679. self._entity_extraction_prompt_profile = (
  680. validate_entity_extraction_prompt_profile_for_mode(
  681. resolved_prompt_profile,
  682. self.entity_extraction_use_json,
  683. self._addon_params.get("entity_type_prompt_file"),
  684. )
  685. )
  686. self._cached_entity_extraction_use_json = self.entity_extraction_use_json
  687. self._addon_params_dirty = False
  688. def _ensure_addon_params_cache(self) -> None:
  689. if (
  690. not self._addon_params_dirty
  691. and self._cached_entity_extraction_use_json
  692. == self.entity_extraction_use_json
  693. ):
  694. return
  695. self._refresh_addon_params_cache()
  696. def _build_global_config(self) -> dict[str, Any]:
  697. self._ensure_addon_params_cache()
  698. global_config = asdict(self)
  699. global_config.pop("_addon_params", None)
  700. global_config.pop("_addon_params_dirty", None)
  701. global_config.pop("_cached_entity_extraction_use_json", None)
  702. global_config["addon_params"] = dict(self._addon_params)
  703. # Inject runtime per-role wrapped LLM funcs (callable; not part of asdict
  704. # because they live in the private _role_llm_states map). The first
  705. # _build_global_config() call from __post_init__ runs before the role
  706. # state is built, so fall back to an empty dict in that case.
  707. states = getattr(self, "_role_llm_states", None) or {}
  708. global_config["role_llm_funcs"] = {
  709. spec.name: states[spec.name].wrapped if spec.name in states else None
  710. for spec in ROLES
  711. }
  712. global_config["llm_cache_identities"] = {
  713. spec.name: self._build_role_llm_cache_identity(
  714. spec.name, states.get(spec.name)
  715. )
  716. for spec in ROLES
  717. }
  718. return global_config
  719. def _build_role_llm_cache_identity(
  720. self, role: str, state: _RoleLLMState | None
  721. ) -> dict[str, Any]:
  722. # `state` is None during the first _build_global_config() call from
  723. # __post_init__ — role builders have not run yet, so metadata is empty
  724. # and we fall back to self.llm_model_name. Once roles are initialized
  725. # or aupdate_llm_role_config() runs, metadata always carries `model`.
  726. metadata = state.metadata if state is not None else {}
  727. return {
  728. "role": role,
  729. "binding": metadata.get("binding"),
  730. "model": metadata.get("model") or self.llm_model_name,
  731. "host": metadata.get("host"),
  732. }
  733. def __post_init__(self, addon_params: dict[str, Any] | None):
  734. from lightrag.kg.shared_storage import (
  735. initialize_share_data,
  736. )
  737. # Fail fast if deprecated ENTITY_TYPES env var is set
  738. if os.getenv("ENTITY_TYPES") is not None:
  739. raise SystemExit(
  740. "ERROR: ENTITY_TYPES environment variable is no longer supported. "
  741. "Please customize entity type guidance through the prompt template instead. "
  742. "Set addon_params={'entity_types_guidance': '...'} or replace the prompt template."
  743. )
  744. self._replace_addon_params(addon_params, mark_dirty=False)
  745. self._apply_chunk_size_overlay()
  746. self._refresh_addon_params_cache()
  747. # Handle deprecated parameters
  748. if self.log_level is not None:
  749. warnings.warn(
  750. "WARNING: log_level parameter is deprecated, use setup_logger in utils.py instead",
  751. UserWarning,
  752. stacklevel=2,
  753. )
  754. if self.log_file_path is not None:
  755. warnings.warn(
  756. "WARNING: log_file_path parameter is deprecated, use setup_logger in utils.py instead",
  757. UserWarning,
  758. stacklevel=2,
  759. )
  760. # Remove these attributes to prevent their use
  761. if hasattr(self, "log_level"):
  762. delattr(self, "log_level")
  763. if hasattr(self, "log_file_path"):
  764. delattr(self, "log_file_path")
  765. initialize_share_data()
  766. if not os.path.exists(self.working_dir):
  767. logger.info(f"Creating working directory {self.working_dir}")
  768. os.makedirs(self.working_dir)
  769. # Verify storage implementation compatibility and environment variables
  770. storage_configs = [
  771. ("KV_STORAGE", self.kv_storage),
  772. ("VECTOR_STORAGE", self.vector_storage),
  773. ("GRAPH_STORAGE", self.graph_storage),
  774. ("DOC_STATUS_STORAGE", self.doc_status_storage),
  775. ]
  776. for storage_type, storage_name in storage_configs:
  777. # Verify storage implementation compatibility
  778. verify_storage_implementation(storage_type, storage_name)
  779. # Check environment variables
  780. check_storage_env_vars(storage_name)
  781. # Ensure vector_db_storage_cls_kwargs has required fields
  782. self.vector_db_storage_cls_kwargs = {
  783. "cosine_better_than_threshold": self.cosine_better_than_threshold,
  784. **self.vector_db_storage_cls_kwargs,
  785. }
  786. # Init Tokenizer
  787. # Post-initialization hook to handle backward compatabile tokenizer initialization based on provided parameters
  788. if self.tokenizer is None:
  789. if self.tiktoken_model_name:
  790. self.tokenizer = TiktokenTokenizer(self.tiktoken_model_name)
  791. else:
  792. self.tokenizer = TiktokenTokenizer()
  793. # Initialize ollama_server_infos if not provided
  794. if self.ollama_server_infos is None:
  795. self.ollama_server_infos = OllamaServerInfos()
  796. # Validate config
  797. if self.force_llm_summary_on_merge < 3:
  798. logger.warning(
  799. f"force_llm_summary_on_merge should be at least 3, got {self.force_llm_summary_on_merge}"
  800. )
  801. if self.summary_context_size > self.max_total_tokens:
  802. logger.warning(
  803. f"summary_context_size({self.summary_context_size}) should no greater than max_total_tokens({self.max_total_tokens})"
  804. )
  805. if self.summary_length_recommended > self.summary_max_tokens:
  806. logger.warning(
  807. f"max_total_tokens({self.summary_max_tokens}) should greater than summary_length_recommended({self.summary_length_recommended})"
  808. )
  809. if self.rerank_model_func is not None:
  810. self.rerank_model_func = priority_limit_async_func_call(
  811. self.rerank_model_max_async,
  812. llm_timeout=self.default_rerank_timeout,
  813. queue_name="Rerank func",
  814. )(self.rerank_model_func)
  815. # Init Embedding
  816. # Step 1: Capture embedding_func and max_token_size before applying rate_limit decorator
  817. original_embedding_func = self.embedding_func
  818. embedding_max_token_size = None
  819. if self.embedding_func and hasattr(self.embedding_func, "max_token_size"):
  820. embedding_max_token_size = self.embedding_func.max_token_size
  821. logger.debug(
  822. f"Captured embedding max_token_size: {embedding_max_token_size}"
  823. )
  824. self.embedding_token_limit = embedding_max_token_size
  825. # Fix global_config now
  826. global_config = self._build_global_config()
  827. # Restore original EmbeddingFunc object (asdict converts it to dict)
  828. global_config["embedding_func"] = original_embedding_func
  829. _print_config = ",\n ".join([f"{k} = {v}" for k, v in global_config.items()])
  830. logger.debug(f"LightRAG init with param:\n {_print_config}\n")
  831. # Step 2: Apply priority wrapper decorator to EmbeddingFunc's inner func
  832. # Create a NEW EmbeddingFunc instance with the wrapped func to avoid mutating the caller's object
  833. # This ensures _generate_collection_suffix can still access attributes (model_name, embedding_dim)
  834. # while preventing side effects when the same EmbeddingFunc is reused across multiple LightRAG instances
  835. if self.embedding_func is not None:
  836. wrapped_func = priority_limit_async_func_call(
  837. self.embedding_func_max_async,
  838. llm_timeout=self.default_embedding_timeout,
  839. queue_name="Embedding func",
  840. )(self.embedding_func.func)
  841. # Use dataclasses.replace() to create a new instance, leaving the original unchanged
  842. self.embedding_func = replace(self.embedding_func, func=wrapped_func)
  843. # Initialize all storages
  844. self.key_string_value_json_storage_cls: type[BaseKVStorage] = get_storage_class(
  845. self.kv_storage
  846. ) # type: ignore
  847. self.vector_db_storage_cls: type[BaseVectorStorage] = get_storage_class(
  848. self.vector_storage
  849. ) # type: ignore
  850. self.graph_storage_cls: type[BaseGraphStorage] = get_storage_class(
  851. self.graph_storage
  852. ) # type: ignore
  853. self.key_string_value_json_storage_cls = partial( # type: ignore
  854. self.key_string_value_json_storage_cls, global_config=global_config
  855. )
  856. self.vector_db_storage_cls = partial( # type: ignore
  857. self.vector_db_storage_cls, global_config=global_config
  858. )
  859. self.graph_storage_cls = partial( # type: ignore
  860. self.graph_storage_cls, global_config=global_config
  861. )
  862. # Initialize document status storage
  863. self.doc_status_storage_cls = get_storage_class(self.doc_status_storage)
  864. self.llm_response_cache: BaseKVStorage = self.key_string_value_json_storage_cls( # type: ignore
  865. namespace=NameSpace.KV_STORE_LLM_RESPONSE_CACHE,
  866. workspace=self.workspace,
  867. global_config=global_config,
  868. embedding_func=self.embedding_func,
  869. )
  870. self.text_chunks: BaseKVStorage = self.key_string_value_json_storage_cls( # type: ignore
  871. namespace=NameSpace.KV_STORE_TEXT_CHUNKS,
  872. workspace=self.workspace,
  873. embedding_func=self.embedding_func,
  874. )
  875. self.full_docs: BaseKVStorage = self.key_string_value_json_storage_cls( # type: ignore
  876. namespace=NameSpace.KV_STORE_FULL_DOCS,
  877. workspace=self.workspace,
  878. embedding_func=self.embedding_func,
  879. )
  880. self.full_entities: BaseKVStorage = self.key_string_value_json_storage_cls( # type: ignore
  881. namespace=NameSpace.KV_STORE_FULL_ENTITIES,
  882. workspace=self.workspace,
  883. embedding_func=self.embedding_func,
  884. )
  885. self.full_relations: BaseKVStorage = self.key_string_value_json_storage_cls( # type: ignore
  886. namespace=NameSpace.KV_STORE_FULL_RELATIONS,
  887. workspace=self.workspace,
  888. embedding_func=self.embedding_func,
  889. )
  890. self.entity_chunks: BaseKVStorage = self.key_string_value_json_storage_cls( # type: ignore
  891. namespace=NameSpace.KV_STORE_ENTITY_CHUNKS,
  892. workspace=self.workspace,
  893. embedding_func=self.embedding_func,
  894. )
  895. self.relation_chunks: BaseKVStorage = self.key_string_value_json_storage_cls( # type: ignore
  896. namespace=NameSpace.KV_STORE_RELATION_CHUNKS,
  897. workspace=self.workspace,
  898. embedding_func=self.embedding_func,
  899. )
  900. self.chunk_entity_relation_graph: BaseGraphStorage = self.graph_storage_cls( # type: ignore
  901. namespace=NameSpace.GRAPH_STORE_CHUNK_ENTITY_RELATION,
  902. workspace=self.workspace,
  903. embedding_func=self.embedding_func,
  904. )
  905. self.entities_vdb: BaseVectorStorage = self.vector_db_storage_cls( # type: ignore
  906. namespace=NameSpace.VECTOR_STORE_ENTITIES,
  907. workspace=self.workspace,
  908. embedding_func=self.embedding_func,
  909. meta_fields={"entity_name", "source_id", "content", "file_path"},
  910. )
  911. self.relationships_vdb: BaseVectorStorage = self.vector_db_storage_cls( # type: ignore
  912. namespace=NameSpace.VECTOR_STORE_RELATIONSHIPS,
  913. workspace=self.workspace,
  914. embedding_func=self.embedding_func,
  915. meta_fields={"src_id", "tgt_id", "source_id", "content", "file_path"},
  916. )
  917. self.chunks_vdb: BaseVectorStorage = self.vector_db_storage_cls( # type: ignore
  918. namespace=NameSpace.VECTOR_STORE_CHUNKS,
  919. workspace=self.workspace,
  920. embedding_func=self.embedding_func,
  921. meta_fields={"full_doc_id", "content", "file_path"},
  922. )
  923. # Initialize document status storage
  924. self.doc_status: DocStatusStorage = self.doc_status_storage_cls(
  925. namespace=NameSpace.DOC_STATUS,
  926. workspace=self.workspace,
  927. global_config=global_config,
  928. embedding_func=None,
  929. )
  930. # Per-role isolated LLM wrappers (independent queues per role).
  931. # The base ``self.llm_model_func`` is intentionally NOT queue-wrapped:
  932. # every code path that calls an LLM goes through one of the role
  933. # wrappers built below, so concurrency is enforced at the role layer.
  934. base_llm_func = self.llm_model_func
  935. if base_llm_func is None:
  936. raise ValueError("llm_model_func must be provided")
  937. self._llm_role_builder = None
  938. self._retired_llm_queue_cleanup_tasks: set[asyncio.Task] = set()
  939. user_role_configs = self.role_llm_configs or {}
  940. if not isinstance(user_role_configs, Mapping):
  941. raise TypeError(
  942. "role_llm_configs must be a Mapping or None, got "
  943. f"{type(user_role_configs).__name__}"
  944. )
  945. unknown_roles = [role for role in user_role_configs if role not in ROLE_NAMES]
  946. if unknown_roles:
  947. valid_roles = ", ".join(sorted(ROLE_NAMES))
  948. unknown = ", ".join(repr(role) for role in unknown_roles)
  949. raise ValueError(
  950. f"Unknown role_llm_configs key(s): {unknown}. "
  951. f"Valid roles are: {valid_roles}"
  952. )
  953. self._role_llm_states: dict[str, _RoleLLMState] = {}
  954. for spec in ROLES:
  955. override = user_role_configs.get(spec.name)
  956. if override is None:
  957. cfg = RoleLLMConfig()
  958. elif isinstance(override, RoleLLMConfig):
  959. cfg = override
  960. elif isinstance(override, Mapping):
  961. cfg = RoleLLMConfig(**dict(override))
  962. else:
  963. raise TypeError(
  964. f"role_llm_configs[{spec.name!r}] must be RoleLLMConfig or "
  965. f"a dict, got {type(override).__name__}"
  966. )
  967. max_async = cfg.max_async
  968. if max_async is None:
  969. max_async = _optional_env_int(f"{spec.env_prefix}_MAX_ASYNC_LLM")
  970. metadata = {}
  971. if cfg.metadata is not None:
  972. if not isinstance(cfg.metadata, Mapping):
  973. raise TypeError(
  974. f"role_llm_configs[{spec.name!r}].metadata must be a "
  975. f"Mapping or None, got {type(cfg.metadata).__name__}"
  976. )
  977. metadata = deepcopy(dict(cfg.metadata))
  978. self._role_llm_states[spec.name] = _RoleLLMState(
  979. raw_func=cfg.func or base_llm_func,
  980. kwargs=cfg.kwargs,
  981. max_async=max_async,
  982. timeout=cfg.timeout,
  983. metadata=metadata,
  984. )
  985. self._rebuild_role_llm_funcs()
  986. self._log_llm_role_config("initialized")
  987. self._storages_status = StoragesStatus.CREATED
  988. async def initialize_storages(self):
  989. """Storage initialization must be called one by one to prevent deadlock"""
  990. if self._storages_status == StoragesStatus.CREATED:
  991. # Set the first initialized workspace will set the default workspace
  992. # Allows namespace operation without specifying workspace for backward compatibility
  993. default_workspace = get_default_workspace()
  994. if default_workspace is None:
  995. set_default_workspace(self.workspace)
  996. elif default_workspace != self.workspace:
  997. logger.info(
  998. f"Creating LightRAG instance with workspace='{self.workspace}' "
  999. f"while default workspace is set to '{default_workspace}'"
  1000. )
  1001. # Auto-initialize pipeline_status for this workspace
  1002. from lightrag.kg.shared_storage import initialize_pipeline_status
  1003. await initialize_pipeline_status(workspace=self.workspace)
  1004. for storage in (
  1005. self.full_docs,
  1006. self.text_chunks,
  1007. self.full_entities,
  1008. self.full_relations,
  1009. self.entity_chunks,
  1010. self.relation_chunks,
  1011. self.entities_vdb,
  1012. self.relationships_vdb,
  1013. self.chunks_vdb,
  1014. self.chunk_entity_relation_graph,
  1015. self.llm_response_cache,
  1016. self.doc_status,
  1017. ):
  1018. if storage:
  1019. # logger.debug(f"Initializing storage: {storage}")
  1020. await storage.initialize()
  1021. self._storages_status = StoragesStatus.INITIALIZED
  1022. logger.debug("All storage types initialized")
  1023. async def finalize_storages(self):
  1024. """Asynchronously finalize the storages with improved error handling"""
  1025. if self._storages_status == StoragesStatus.INITIALIZED:
  1026. storages = [
  1027. ("full_docs", self.full_docs),
  1028. ("text_chunks", self.text_chunks),
  1029. ("full_entities", self.full_entities),
  1030. ("full_relations", self.full_relations),
  1031. ("entity_chunks", self.entity_chunks),
  1032. ("relation_chunks", self.relation_chunks),
  1033. ("entities_vdb", self.entities_vdb),
  1034. ("relationships_vdb", self.relationships_vdb),
  1035. ("chunks_vdb", self.chunks_vdb),
  1036. ("chunk_entity_relation_graph", self.chunk_entity_relation_graph),
  1037. ("llm_response_cache", self.llm_response_cache),
  1038. ("doc_status", self.doc_status),
  1039. ]
  1040. # Finalize each storage individually to ensure one failure doesn't prevent others from closing
  1041. successful_finalizations = []
  1042. failed_finalizations = []
  1043. for storage_name, storage in storages:
  1044. if storage:
  1045. try:
  1046. await storage.finalize()
  1047. successful_finalizations.append(storage_name)
  1048. logger.debug(f"Successfully finalized {storage_name}")
  1049. except Exception as e:
  1050. error_msg = f"Failed to finalize {storage_name}: {e}"
  1051. logger.error(error_msg)
  1052. failed_finalizations.append(storage_name)
  1053. # Log summary of finalization results
  1054. if successful_finalizations:
  1055. logger.info(
  1056. f"Successfully finalized {len(successful_finalizations)} storages"
  1057. )
  1058. if failed_finalizations:
  1059. logger.error(
  1060. f"Failed to finalize {len(failed_finalizations)} storages: {', '.join(failed_finalizations)}"
  1061. )
  1062. else:
  1063. logger.debug("All storages finalized successfully")
  1064. self._storages_status = StoragesStatus.FINALIZED
  1065. async def get_graph_labels(self):
  1066. text = await self.chunk_entity_relation_graph.get_all_labels()
  1067. return text
  1068. async def get_knowledge_graph(
  1069. self,
  1070. node_label: str,
  1071. max_depth: int = 3,
  1072. max_nodes: int = None,
  1073. ) -> KnowledgeGraph:
  1074. """Get knowledge graph for a given label
  1075. Args:
  1076. node_label (str): Label to get knowledge graph for
  1077. max_depth (int): Maximum depth of graph
  1078. max_nodes (int, optional): Maximum number of nodes to return. Defaults to self.max_graph_nodes.
  1079. Returns:
  1080. KnowledgeGraph: Knowledge graph containing nodes and edges
  1081. """
  1082. # Use self.max_graph_nodes as default if max_nodes is None
  1083. if max_nodes is None:
  1084. max_nodes = self.max_graph_nodes
  1085. else:
  1086. # Limit max_nodes to not exceed self.max_graph_nodes
  1087. max_nodes = min(max_nodes, self.max_graph_nodes)
  1088. return await self.chunk_entity_relation_graph.get_knowledge_graph(
  1089. node_label, max_depth, max_nodes
  1090. )
  1091. def insert(
  1092. self,
  1093. input: str | list[str],
  1094. split_by_character: str | None = None,
  1095. split_by_character_only: bool = False,
  1096. ids: str | list[str] | None = None,
  1097. file_paths: str | list[str] | None = None,
  1098. track_id: str | None = None,
  1099. ) -> str:
  1100. """Sync Insert documents with checkpoint support
  1101. Args:
  1102. input: Single document string or list of document strings
  1103. split_by_character: if split_by_character is not None, split the string by character, if chunk longer than
  1104. chunk_token_size, it will be split again by token size.
  1105. split_by_character_only: if split_by_character_only is True, split the string by character only, when
  1106. split_by_character is None, this parameter is ignored.
  1107. ids: single string of the document ID or list of unique document IDs, if not provided, MD5 hash IDs will be generated
  1108. file_paths: single string of the file path or list of file paths, used for citation
  1109. track_id: tracking ID for monitoring processing status, if not provided, will be generated
  1110. Returns:
  1111. str: tracking ID for monitoring processing status
  1112. """
  1113. loop = always_get_an_event_loop()
  1114. return loop.run_until_complete(
  1115. self.ainsert(
  1116. input,
  1117. split_by_character,
  1118. split_by_character_only,
  1119. ids,
  1120. file_paths,
  1121. track_id,
  1122. )
  1123. )
  1124. async def ainsert(
  1125. self,
  1126. input: str | list[str],
  1127. split_by_character: str | None = None,
  1128. split_by_character_only: bool = False,
  1129. ids: str | list[str] | None = None,
  1130. file_paths: str | list[str] | None = None,
  1131. track_id: str | None = None,
  1132. ) -> str:
  1133. """Async insert documents with checkpoint support (fixed-token chunking only).
  1134. SDK convenience entry point. It **always** chunks with the fixed-token
  1135. (F) strategy: ``process_options`` is intentionally not passed, so the
  1136. document runs the F chunker. ``split_by_character`` /
  1137. ``split_by_character_only`` are F-strategy runtime args; the rest of
  1138. the F config (``chunk_token_size`` / ``chunk_overlap_token_size``,
  1139. seeded from ``CHUNK_F_SIZE`` / ``CHUNK_SIZE`` etc.) comes from
  1140. ``addon_params['chunker']['fixed_token']``. ``ainsert`` cannot select
  1141. the recursive-character (R), semantic-vector (V), or paragraph-semantic
  1142. (P) strategies.
  1143. The LightRAG **server / REST API does not call this method** — it
  1144. ingests via :meth:`apipeline_enqueue_documents` +
  1145. :meth:`apipeline_process_enqueue_documents` with a per-document
  1146. ``process_options`` selector, which is how F/R/V/P are chosen there.
  1147. To use R/V/P (or pass an explicit per-document ``chunk_options``) from
  1148. the SDK, call those two methods directly with ``process_options=…``
  1149. instead of ``ainsert``.
  1150. Args:
  1151. input: Single document string or list of document strings
  1152. split_by_character: if split_by_character is not None, split the string by character, if chunk longer than
  1153. chunk_token_size, it will be split again by token size.
  1154. split_by_character_only: if split_by_character_only is True, split the string by character only, when
  1155. split_by_character is None, this parameter is ignored.
  1156. ids: list of unique document IDs, if not provided, MD5 hash IDs will be generated
  1157. file_paths: list of file paths corresponding to each document, used for citation
  1158. track_id: tracking ID for monitoring processing status, if not provided, will be generated
  1159. Returns:
  1160. str: tracking ID for monitoring processing status
  1161. """
  1162. # Generate track_id if not provided
  1163. if track_id is None:
  1164. track_id = generate_track_id("insert")
  1165. # Capture the F-strategy runtime args into a chunk_options
  1166. # snapshot before enqueue so they become a per-document
  1167. # setting. ``apipeline_enqueue_documents`` itself doesn't take
  1168. # split args — chunk_options is the canonical chunker-config
  1169. # carrier; runtime split args are an ainsert-only concern.
  1170. from lightrag.parser.routing import resolve_chunk_options
  1171. chunk_opts = resolve_chunk_options(
  1172. self.addon_params,
  1173. split_by_character=split_by_character,
  1174. split_by_character_only=split_by_character_only,
  1175. )
  1176. await self.apipeline_enqueue_documents(
  1177. input,
  1178. ids,
  1179. file_paths,
  1180. track_id,
  1181. chunk_options=chunk_opts,
  1182. )
  1183. await self.apipeline_process_enqueue_documents()
  1184. return track_id
  1185. # TODO: deprecated, use insert instead
  1186. def insert_custom_chunks(
  1187. self,
  1188. full_text: str,
  1189. text_chunks: list[str],
  1190. doc_id: str | list[str] | None = None,
  1191. ) -> None:
  1192. loop = always_get_an_event_loop()
  1193. loop.run_until_complete(
  1194. self.ainsert_custom_chunks(full_text, text_chunks, doc_id)
  1195. )
  1196. # TODO: deprecated, use ainsert instead
  1197. async def ainsert_custom_chunks(
  1198. self, full_text: str, text_chunks: list[str], doc_id: str | None = None
  1199. ) -> None:
  1200. update_storage = False
  1201. try:
  1202. # Clean input texts
  1203. full_text = sanitize_text_for_encoding(full_text)
  1204. text_chunks = [sanitize_text_for_encoding(chunk) for chunk in text_chunks]
  1205. file_path = normalize_document_file_path("")
  1206. # Process cleaned texts
  1207. if doc_id is None:
  1208. doc_key = compute_mdhash_id(full_text, prefix="doc-")
  1209. else:
  1210. doc_key = doc_id
  1211. new_docs = {doc_key: {"content": full_text, "file_path": file_path}}
  1212. _add_doc_keys = await self.full_docs.filter_keys({doc_key})
  1213. new_docs = {k: v for k, v in new_docs.items() if k in _add_doc_keys}
  1214. if not len(new_docs):
  1215. logger.warning("This document is already in the storage.")
  1216. return
  1217. update_storage = True
  1218. logger.info(f"Inserting {len(new_docs)} docs")
  1219. inserting_chunks: dict[str, Any] = {}
  1220. for index, chunk_text in enumerate(text_chunks):
  1221. chunk_key = compute_mdhash_id(chunk_text, prefix="chunk-")
  1222. tokens = len(self.tokenizer.encode(chunk_text))
  1223. inserting_chunks[chunk_key] = {
  1224. "content": chunk_text,
  1225. "full_doc_id": doc_key,
  1226. "tokens": tokens,
  1227. "chunk_order_index": index,
  1228. "file_path": file_path,
  1229. }
  1230. doc_ids = set(inserting_chunks.keys())
  1231. add_chunk_keys = await self.text_chunks.filter_keys(doc_ids)
  1232. inserting_chunks = {
  1233. k: v for k, v in inserting_chunks.items() if k in add_chunk_keys
  1234. }
  1235. if not len(inserting_chunks):
  1236. logger.warning("All chunks are already in the storage.")
  1237. return
  1238. tasks = [
  1239. self.chunks_vdb.upsert(inserting_chunks),
  1240. self._process_extract_entities(inserting_chunks),
  1241. self.full_docs.upsert(new_docs),
  1242. self.text_chunks.upsert(inserting_chunks),
  1243. ]
  1244. await asyncio.gather(*tasks)
  1245. finally:
  1246. if update_storage:
  1247. await self._insert_done()
  1248. async def _process_extract_entities(
  1249. self, chunk: dict[str, Any], pipeline_status=None, pipeline_status_lock=None
  1250. ) -> list:
  1251. try:
  1252. chunk_results = await extract_entities(
  1253. chunk,
  1254. global_config=self._build_global_config(),
  1255. pipeline_status=pipeline_status,
  1256. pipeline_status_lock=pipeline_status_lock,
  1257. llm_response_cache=self.llm_response_cache,
  1258. text_chunks_storage=self.text_chunks,
  1259. )
  1260. return chunk_results
  1261. except Exception as e:
  1262. error_msg = f"Failed to extract entities and relationships: {str(e)}"
  1263. logger.error(error_msg)
  1264. async with pipeline_status_lock:
  1265. pipeline_status["latest_message"] = error_msg
  1266. pipeline_status["history_messages"].append(error_msg)
  1267. raise e
  1268. async def _insert_done(
  1269. self, pipeline_status=None, pipeline_status_lock=None
  1270. ) -> None:
  1271. tasks = [
  1272. cast(StorageNameSpace, storage_inst).index_done_callback()
  1273. for storage_inst in [ # type: ignore
  1274. self.full_docs,
  1275. self.doc_status,
  1276. self.text_chunks,
  1277. self.full_entities,
  1278. self.full_relations,
  1279. self.entity_chunks,
  1280. self.relation_chunks,
  1281. self.llm_response_cache,
  1282. self.entities_vdb,
  1283. self.relationships_vdb,
  1284. self.chunks_vdb,
  1285. self.chunk_entity_relation_graph,
  1286. ]
  1287. if storage_inst is not None
  1288. ]
  1289. await asyncio.gather(*tasks)
  1290. log_message = "In memory DB persist to disk"
  1291. logger.info(log_message)
  1292. if pipeline_status is not None and pipeline_status_lock is not None:
  1293. async with pipeline_status_lock:
  1294. pipeline_status["latest_message"] = log_message
  1295. pipeline_status["history_messages"].append(log_message)
  1296. def insert_custom_kg(
  1297. self, custom_kg: dict[str, Any], full_doc_id: str = None
  1298. ) -> None:
  1299. loop = always_get_an_event_loop()
  1300. loop.run_until_complete(self.ainsert_custom_kg(custom_kg, full_doc_id))
  1301. async def ainsert_custom_kg(
  1302. self,
  1303. custom_kg: dict[str, Any],
  1304. full_doc_id: str = None,
  1305. ) -> None:
  1306. update_storage = False
  1307. try:
  1308. # Insert chunks into vector storage
  1309. all_chunks_data: dict[str, dict[str, str]] = {}
  1310. chunk_to_source_map: dict[str, str] = {}
  1311. for chunk_data in custom_kg.get("chunks", []):
  1312. chunk_content = sanitize_text_for_encoding(chunk_data["content"])
  1313. source_id = chunk_data["source_id"]
  1314. file_path = normalize_document_file_path(
  1315. chunk_data.get("file_path", "custom_kg")
  1316. )
  1317. tokens = len(self.tokenizer.encode(chunk_content))
  1318. chunk_order_index = (
  1319. 0
  1320. if "chunk_order_index" not in chunk_data.keys()
  1321. else chunk_data["chunk_order_index"]
  1322. )
  1323. chunk_id = compute_mdhash_id(chunk_content, prefix="chunk-")
  1324. chunk_entry = {
  1325. "content": chunk_content,
  1326. "source_id": source_id,
  1327. "tokens": tokens,
  1328. "chunk_order_index": chunk_order_index,
  1329. "full_doc_id": full_doc_id
  1330. if full_doc_id is not None
  1331. else source_id,
  1332. "file_path": file_path,
  1333. "status": DocStatus.PROCESSED,
  1334. }
  1335. all_chunks_data[chunk_id] = chunk_entry
  1336. chunk_to_source_map[source_id] = chunk_id
  1337. update_storage = True
  1338. if all_chunks_data:
  1339. await asyncio.gather(
  1340. self.chunks_vdb.upsert(all_chunks_data),
  1341. self.text_chunks.upsert(all_chunks_data),
  1342. )
  1343. # Keep the last declaration for each entity_name so batch backends
  1344. # preserve the old serial upsert semantics deterministically.
  1345. deduped_entities: dict[str, dict[str, Any]] = {}
  1346. for entity_data in custom_kg.get("entities", []):
  1347. entity_name = entity_data["entity_name"]
  1348. deduped_entities.pop(entity_name, None)
  1349. deduped_entities[entity_name] = entity_data
  1350. # Insert entities into knowledge graph (batch for performance)
  1351. all_entities_data: list[dict[str, str]] = []
  1352. entity_nodes: list[tuple[str, dict[str, str]]] = []
  1353. for entity_data in deduped_entities.values():
  1354. entity_name = entity_data["entity_name"]
  1355. entity_type = entity_data.get("entity_type", "UNKNOWN")
  1356. description = entity_data.get("description", "No description provided")
  1357. source_chunk_id = entity_data.get("source_id", "UNKNOWN")
  1358. source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
  1359. file_path = normalize_document_file_path(
  1360. entity_data.get("file_path", "custom_kg")
  1361. )
  1362. if source_id == "UNKNOWN":
  1363. logger.warning(
  1364. f"Entity '{entity_name}' has an UNKNOWN source_id. Please check the source mapping."
  1365. )
  1366. node_data: dict[str, str] = {
  1367. "entity_id": entity_name,
  1368. "entity_type": entity_type,
  1369. "description": description,
  1370. "source_id": source_id,
  1371. "file_path": file_path,
  1372. "created_at": int(time.time()),
  1373. }
  1374. entity_nodes.append((entity_name, node_data))
  1375. node_data_copy = dict(node_data)
  1376. node_data_copy["entity_name"] = entity_name
  1377. all_entities_data.append(node_data_copy)
  1378. update_storage = True
  1379. # Relationship storage is undirected, so keep only the last update
  1380. # for each endpoint pair regardless of order.
  1381. deduped_relationships: dict[tuple[str, str], dict[str, Any]] = {}
  1382. for relationship_data in custom_kg.get("relationships", []):
  1383. src_id = relationship_data["src_id"]
  1384. tgt_id = relationship_data["tgt_id"]
  1385. relation_key = tuple(sorted((src_id, tgt_id)))
  1386. deduped_relationships.pop(relation_key, None)
  1387. deduped_relationships[relation_key] = relationship_data
  1388. # Coarse-grained keyed lock covering every entity name and every
  1389. # relationship endpoint this batch will write. Keys collide with
  1390. # the per-entity and sorted([src, tgt]) edge locks held by the
  1391. # doc-ingest pipeline (operate.py:_locked_process_entity_name and
  1392. # _locked_process_edges) in the same namespace, so a concurrent
  1393. # insert_custom_kg waits behind an in-flight document ingest
  1394. # rather than racing it. Two concurrent custom-KG inserts that
  1395. # touch overlapping entities likewise mutually exclude here.
  1396. # An empty batch skips the lock entirely — nothing to serialise on.
  1397. lock_key_set: set[str] = {entity_name for entity_name, _ in entity_nodes}
  1398. for relationship_data in deduped_relationships.values():
  1399. lock_key_set.add(relationship_data["src_id"])
  1400. lock_key_set.add(relationship_data["tgt_id"])
  1401. workspace = self.workspace or ""
  1402. namespace = f"{workspace}:GraphDB" if workspace else "GraphDB"
  1403. async def _do_graph_and_vdb_writes() -> None:
  1404. # Batch insert entities (reduces N serial awaits to 1)
  1405. if entity_nodes:
  1406. await self.chunk_entity_relation_graph.upsert_nodes_batch(
  1407. entity_nodes
  1408. )
  1409. # Insert relationships into knowledge graph (batch for performance)
  1410. all_relationships_data: list[dict[str, str]] = []
  1411. edge_list: list[tuple[str, str, dict[str, str]]] = []
  1412. # Batch check which relationship endpoints exist (1 await instead of 2M)
  1413. needed_node_ids: set[str] = set()
  1414. for relationship_data in deduped_relationships.values():
  1415. needed_node_ids.add(relationship_data["src_id"])
  1416. needed_node_ids.add(relationship_data["tgt_id"])
  1417. existing_nodes = await self.chunk_entity_relation_graph.has_nodes_batch(
  1418. list(needed_node_ids)
  1419. )
  1420. # Create missing nodes in batch
  1421. missing_nodes: list[tuple[str, dict[str, str]]] = []
  1422. for relationship_data in deduped_relationships.values():
  1423. src_id = relationship_data["src_id"]
  1424. tgt_id = relationship_data["tgt_id"]
  1425. source_chunk_id = relationship_data.get("source_id", "UNKNOWN")
  1426. source_id = chunk_to_source_map.get(source_chunk_id, "UNKNOWN")
  1427. file_path = normalize_document_file_path(
  1428. relationship_data.get("file_path", "custom_kg")
  1429. )
  1430. if source_id == "UNKNOWN":
  1431. logger.warning(
  1432. f"Relationship from '{src_id}' to '{tgt_id}' has an UNKNOWN source_id. Please check the source mapping."
  1433. )
  1434. for need_insert_id in [src_id, tgt_id]:
  1435. if need_insert_id not in existing_nodes:
  1436. missing_nodes.append(
  1437. (
  1438. need_insert_id,
  1439. {
  1440. "entity_id": need_insert_id,
  1441. "source_id": source_id,
  1442. "description": "UNKNOWN",
  1443. "entity_type": "UNKNOWN",
  1444. "file_path": file_path,
  1445. "created_at": int(time.time()),
  1446. },
  1447. )
  1448. )
  1449. existing_nodes.add(need_insert_id)
  1450. normalized_src_id, normalized_tgt_id = sorted((src_id, tgt_id))
  1451. edge_data = {
  1452. "weight": relationship_data.get("weight", 1.0),
  1453. "description": relationship_data["description"],
  1454. "keywords": relationship_data["keywords"],
  1455. "source_id": source_id,
  1456. "file_path": file_path,
  1457. "created_at": int(time.time()),
  1458. }
  1459. edge_list.append((src_id, tgt_id, edge_data))
  1460. all_relationships_data.append(
  1461. {
  1462. "src_id": normalized_src_id,
  1463. "tgt_id": normalized_tgt_id,
  1464. "description": relationship_data["description"],
  1465. "keywords": relationship_data["keywords"],
  1466. "source_id": source_id,
  1467. "weight": relationship_data.get("weight", 1.0),
  1468. "file_path": file_path,
  1469. "created_at": int(time.time()),
  1470. }
  1471. )
  1472. # Batch insert missing placeholder nodes
  1473. if missing_nodes:
  1474. await self.chunk_entity_relation_graph.upsert_nodes_batch(
  1475. missing_nodes
  1476. )
  1477. # Batch insert edges
  1478. if edge_list:
  1479. await self.chunk_entity_relation_graph.upsert_edges_batch(edge_list)
  1480. # Insert entities and relationships into vector storage (parallel)
  1481. data_for_entities_vdb = {
  1482. compute_mdhash_id(dp["entity_name"], prefix="ent-"): {
  1483. "content": dp["entity_name"] + "\n" + dp["description"],
  1484. "entity_name": dp["entity_name"],
  1485. "source_id": dp["source_id"],
  1486. "description": dp["description"],
  1487. "entity_type": dp["entity_type"],
  1488. "file_path": dp.get("file_path", "custom_kg"),
  1489. }
  1490. for dp in all_entities_data
  1491. }
  1492. data_for_rels_vdb = {
  1493. compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): {
  1494. "src_id": dp["src_id"],
  1495. "tgt_id": dp["tgt_id"],
  1496. "source_id": dp["source_id"],
  1497. "content": f"{dp['keywords']}\t{dp['src_id']}\n{dp['tgt_id']}\n{dp['description']}",
  1498. "keywords": dp["keywords"],
  1499. "description": dp["description"],
  1500. "weight": dp["weight"],
  1501. "file_path": dp.get("file_path", "custom_kg"),
  1502. }
  1503. for dp in all_relationships_data
  1504. }
  1505. legacy_rel_ids_to_delete = sorted(
  1506. {
  1507. rel_id
  1508. for dp in all_relationships_data
  1509. for rel_id in make_relation_vdb_ids(dp["src_id"], dp["tgt_id"])[
  1510. 1:
  1511. ]
  1512. }
  1513. )
  1514. # Parallel VDB upserts (was serial in original)
  1515. await asyncio.gather(
  1516. self.entities_vdb.upsert(data_for_entities_vdb),
  1517. self.relationships_vdb.upsert(data_for_rels_vdb),
  1518. )
  1519. if legacy_rel_ids_to_delete:
  1520. await self.relationships_vdb.delete(legacy_rel_ids_to_delete)
  1521. if lock_key_set:
  1522. if entity_nodes or deduped_relationships:
  1523. update_storage = True
  1524. async with get_storage_keyed_lock(
  1525. sorted(lock_key_set),
  1526. namespace=namespace,
  1527. enable_logging=False,
  1528. ):
  1529. await _do_graph_and_vdb_writes()
  1530. else:
  1531. # No entities, no relationships — nothing to serialise on.
  1532. await _do_graph_and_vdb_writes()
  1533. except Exception as e:
  1534. logger.error(f"Error in ainsert_custom_kg: {e}")
  1535. raise
  1536. finally:
  1537. if update_storage:
  1538. await self._insert_done()
  1539. def query(
  1540. self,
  1541. query: str,
  1542. param: QueryParam = QueryParam(),
  1543. system_prompt: str | None = None,
  1544. ) -> str | Iterator[str]:
  1545. """
  1546. Perform a sync query.
  1547. Args:
  1548. query (str): The query to be executed.
  1549. param (QueryParam): Configuration parameters for query execution.
  1550. prompt (Optional[str]): Custom prompts for fine-tuned control over the system's behavior. Defaults to None, which uses PROMPTS["rag_response"].
  1551. Returns:
  1552. str: The result of the query execution.
  1553. """
  1554. loop = always_get_an_event_loop()
  1555. return loop.run_until_complete(self.aquery(query, param, system_prompt)) # type: ignore
  1556. async def aquery(
  1557. self,
  1558. query: str,
  1559. param: QueryParam = QueryParam(),
  1560. system_prompt: str | None = None,
  1561. ) -> str | AsyncIterator[str]:
  1562. """
  1563. Perform a async query (backward compatibility wrapper).
  1564. This function is now a wrapper around aquery_llm that maintains backward compatibility
  1565. by returning only the LLM response content in the original format.
  1566. Args:
  1567. query (str): The query to be executed.
  1568. param (QueryParam): Configuration parameters for query execution.
  1569. system_prompt (Optional[str]): Custom prompts for fine-tuned control over the system's behavior. Defaults to None, which uses PROMPTS["rag_response"].
  1570. Returns:
  1571. str | AsyncIterator[str]: The LLM response content.
  1572. - Non-streaming: Returns str
  1573. - Streaming: Returns AsyncIterator[str]
  1574. """
  1575. # Call the new aquery_llm function to get complete results
  1576. result = await self.aquery_llm(query, param, system_prompt)
  1577. # Extract and return only the LLM response for backward compatibility
  1578. llm_response = result.get("llm_response", {})
  1579. if llm_response.get("is_streaming"):
  1580. return llm_response.get("response_iterator")
  1581. else:
  1582. return llm_response.get("content", "")
  1583. def query_data(
  1584. self,
  1585. query: str,
  1586. param: QueryParam = QueryParam(),
  1587. ) -> dict[str, Any]:
  1588. """
  1589. Synchronous data retrieval API: returns structured retrieval results without LLM generation.
  1590. This function is the synchronous version of aquery_data, providing the same functionality
  1591. for users who prefer synchronous interfaces.
  1592. Args:
  1593. query: Query text for retrieval.
  1594. param: Query parameters controlling retrieval behavior (same as aquery).
  1595. Returns:
  1596. dict[str, Any]: Same structured data result as aquery_data.
  1597. """
  1598. loop = always_get_an_event_loop()
  1599. return loop.run_until_complete(self.aquery_data(query, param))
  1600. async def aquery_data(
  1601. self,
  1602. query: str,
  1603. param: QueryParam = QueryParam(),
  1604. ) -> dict[str, Any]:
  1605. """
  1606. Asynchronous data retrieval API: returns structured retrieval results without LLM generation.
  1607. This function reuses the same logic as aquery but stops before LLM generation,
  1608. returning the final processed entities, relationships, and chunks data that would be sent to LLM.
  1609. Args:
  1610. query: Query text for retrieval.
  1611. param: Query parameters controlling retrieval behavior (same as aquery).
  1612. Returns:
  1613. dict[str, Any]: Structured data result in the following format:
  1614. **Success Response:**
  1615. ```python
  1616. {
  1617. "status": "success",
  1618. "message": "Query executed successfully",
  1619. "data": {
  1620. "entities": [
  1621. {
  1622. "entity_name": str, # Entity identifier
  1623. "entity_type": str, # Entity category/type
  1624. "description": str, # Entity description
  1625. "source_id": str, # Source chunk references
  1626. "file_path": str, # Origin file path
  1627. "created_at": str, # Creation timestamp
  1628. "reference_id": str # Reference identifier for citations
  1629. }
  1630. ],
  1631. "relationships": [
  1632. {
  1633. "src_id": str, # Source entity name
  1634. "tgt_id": str, # Target entity name
  1635. "description": str, # Relationship description
  1636. "keywords": str, # Relationship keywords
  1637. "weight": float, # Relationship strength
  1638. "source_id": str, # Source chunk references
  1639. "file_path": str, # Origin file path
  1640. "created_at": str, # Creation timestamp
  1641. "reference_id": str # Reference identifier for citations
  1642. }
  1643. ],
  1644. "chunks": [
  1645. {
  1646. "content": str, # Document chunk content
  1647. "file_path": str, # Origin file path
  1648. "chunk_id": str, # Unique chunk identifier
  1649. "reference_id": str # Reference identifier for citations
  1650. }
  1651. ],
  1652. "references": [
  1653. {
  1654. "reference_id": str, # Reference identifier
  1655. "file_path": str # Corresponding file path
  1656. }
  1657. ]
  1658. },
  1659. "metadata": {
  1660. "query_mode": str, # Query mode used ("local", "global", "hybrid", "mix", "naive", "bypass")
  1661. "keywords": {
  1662. "high_level": List[str], # High-level keywords extracted
  1663. "low_level": List[str] # Low-level keywords extracted
  1664. },
  1665. "processing_info": {
  1666. "total_entities_found": int, # Total entities before truncation
  1667. "total_relations_found": int, # Total relations before truncation
  1668. "entities_after_truncation": int, # Entities after token truncation
  1669. "relations_after_truncation": int, # Relations after token truncation
  1670. "merged_chunks_count": int, # Chunks before final processing
  1671. "final_chunks_count": int # Final chunks in result
  1672. }
  1673. }
  1674. }
  1675. ```
  1676. **Query Mode Differences:**
  1677. - **local**: Focuses on entities and their related chunks based on low-level keywords
  1678. - **global**: Focuses on relationships and their connected entities based on high-level keywords
  1679. - **hybrid**: Combines local and global results using round-robin merging
  1680. - **mix**: Includes knowledge graph data plus vector-retrieved document chunks
  1681. - **naive**: Only vector-retrieved chunks, entities and relationships arrays are empty
  1682. - **bypass**: All data arrays are empty, used for direct LLM queries
  1683. ** processing_info is optional and may not be present in all responses, especially when query result is empty**
  1684. **Failure Response:**
  1685. ```python
  1686. {
  1687. "status": "failure",
  1688. "message": str, # Error description
  1689. "data": {} # Empty data object
  1690. }
  1691. ```
  1692. **Common Failure Cases:**
  1693. - Empty query string
  1694. - Both high-level and low-level keywords are empty
  1695. - Query returns empty dataset
  1696. - Missing tokenizer or system configuration errors
  1697. Note:
  1698. The function adapts to the new data format from convert_to_user_format where
  1699. actual data is nested under the 'data' field, with 'status' and 'message'
  1700. fields at the top level.
  1701. """
  1702. global_config = self._build_global_config()
  1703. # Create a copy of param to avoid modifying the original
  1704. data_param = QueryParam(
  1705. mode=param.mode,
  1706. only_need_context=True, # Skip LLM generation, only get context and data
  1707. only_need_prompt=False,
  1708. response_type=param.response_type,
  1709. stream=False, # Data retrieval doesn't need streaming
  1710. top_k=param.top_k,
  1711. chunk_top_k=param.chunk_top_k,
  1712. max_entity_tokens=param.max_entity_tokens,
  1713. max_relation_tokens=param.max_relation_tokens,
  1714. max_total_tokens=param.max_total_tokens,
  1715. hl_keywords=param.hl_keywords,
  1716. ll_keywords=param.ll_keywords,
  1717. conversation_history=param.conversation_history,
  1718. user_prompt=param.user_prompt,
  1719. enable_rerank=param.enable_rerank,
  1720. )
  1721. query_result = None
  1722. if data_param.mode in ["local", "global", "hybrid", "mix"]:
  1723. logger.debug(f"[aquery_data] Using kg_query for mode: {data_param.mode}")
  1724. query_result = await kg_query(
  1725. query.strip(),
  1726. self.chunk_entity_relation_graph,
  1727. self.entities_vdb,
  1728. self.relationships_vdb,
  1729. self.text_chunks,
  1730. data_param, # Use data_param with only_need_context=True
  1731. global_config,
  1732. hashing_kv=self.llm_response_cache,
  1733. system_prompt=None,
  1734. chunks_vdb=self.chunks_vdb,
  1735. )
  1736. elif data_param.mode == "naive":
  1737. logger.debug(f"[aquery_data] Using naive_query for mode: {data_param.mode}")
  1738. query_result = await naive_query(
  1739. query.strip(),
  1740. self.chunks_vdb,
  1741. data_param, # Use data_param with only_need_context=True
  1742. global_config,
  1743. hashing_kv=self.llm_response_cache,
  1744. system_prompt=None,
  1745. )
  1746. elif data_param.mode == "bypass":
  1747. logger.debug("[aquery_data] Using bypass mode")
  1748. # bypass mode returns empty data using convert_to_user_format
  1749. empty_raw_data = convert_to_user_format(
  1750. [], # no entities
  1751. [], # no relationships
  1752. [], # no chunks
  1753. [], # no references
  1754. "bypass",
  1755. )
  1756. query_result = QueryResult(content="", raw_data=empty_raw_data)
  1757. else:
  1758. raise ValueError(f"Unknown mode {data_param.mode}")
  1759. if query_result is None:
  1760. no_result_message = "Query returned no results"
  1761. if data_param.mode == "naive":
  1762. no_result_message = "No relevant document chunks found."
  1763. final_data: dict[str, Any] = {
  1764. "status": "failure",
  1765. "message": no_result_message,
  1766. "data": {},
  1767. "metadata": {
  1768. "failure_reason": "no_results",
  1769. "mode": data_param.mode,
  1770. },
  1771. }
  1772. logger.info("[aquery_data] Query returned no results.")
  1773. else:
  1774. # Extract raw_data from QueryResult
  1775. final_data = query_result.raw_data or {}
  1776. # Log final result counts - adapt to new data format from convert_to_user_format
  1777. if final_data and "data" in final_data:
  1778. data_section = final_data["data"]
  1779. entities_count = len(data_section.get("entities", []))
  1780. relationships_count = len(data_section.get("relationships", []))
  1781. chunks_count = len(data_section.get("chunks", []))
  1782. logger.debug(
  1783. f"[aquery_data] Final result: {entities_count} entities, {relationships_count} relationships, {chunks_count} chunks"
  1784. )
  1785. else:
  1786. logger.warning("[aquery_data] No data section found in query result")
  1787. await self._query_done()
  1788. return final_data
  1789. async def aquery_llm(
  1790. self,
  1791. query: str,
  1792. param: QueryParam = QueryParam(),
  1793. system_prompt: str | None = None,
  1794. ) -> dict[str, Any]:
  1795. """
  1796. Asynchronous complete query API: returns structured retrieval results with LLM generation.
  1797. This function performs a single query operation and returns both structured data and LLM response,
  1798. based on the original aquery logic to avoid duplicate calls.
  1799. Args:
  1800. query: Query text for retrieval and LLM generation.
  1801. param: Query parameters controlling retrieval and LLM behavior.
  1802. system_prompt: Optional custom system prompt for LLM generation.
  1803. Returns:
  1804. dict[str, Any]: Complete response with structured data and LLM response.
  1805. """
  1806. logger.debug(f"[aquery_llm] Query param: {param}")
  1807. global_config = self._build_global_config()
  1808. try:
  1809. query_result = None
  1810. if param.mode in ["local", "global", "hybrid", "mix"]:
  1811. query_result = await kg_query(
  1812. query.strip(),
  1813. self.chunk_entity_relation_graph,
  1814. self.entities_vdb,
  1815. self.relationships_vdb,
  1816. self.text_chunks,
  1817. param,
  1818. global_config,
  1819. hashing_kv=self.llm_response_cache,
  1820. system_prompt=system_prompt,
  1821. chunks_vdb=self.chunks_vdb,
  1822. )
  1823. elif param.mode == "naive":
  1824. query_result = await naive_query(
  1825. query.strip(),
  1826. self.chunks_vdb,
  1827. param,
  1828. global_config,
  1829. hashing_kv=self.llm_response_cache,
  1830. system_prompt=system_prompt,
  1831. )
  1832. elif param.mode == "bypass":
  1833. # Bypass mode: directly use LLM without knowledge retrieval
  1834. # Apply higher priority (8) to entity/relation summary tasks
  1835. use_llm_func = partial(
  1836. global_config["role_llm_funcs"]["query"], _priority=8
  1837. )
  1838. param.stream = True if param.stream is None else param.stream
  1839. response = await use_llm_func(
  1840. query.strip(),
  1841. system_prompt=system_prompt,
  1842. history_messages=param.conversation_history,
  1843. enable_cot=True,
  1844. stream=param.stream,
  1845. )
  1846. if type(response) is str:
  1847. return {
  1848. "status": "success",
  1849. "message": "Bypass mode LLM non streaming response",
  1850. "data": {},
  1851. "metadata": {},
  1852. "llm_response": {
  1853. "content": response,
  1854. "response_iterator": None,
  1855. "is_streaming": False,
  1856. },
  1857. }
  1858. else:
  1859. return {
  1860. "status": "success",
  1861. "message": "Bypass mode LLM streaming response",
  1862. "data": {},
  1863. "metadata": {},
  1864. "llm_response": {
  1865. "content": None,
  1866. "response_iterator": response,
  1867. "is_streaming": True,
  1868. },
  1869. }
  1870. else:
  1871. raise ValueError(f"Unknown mode {param.mode}")
  1872. await self._query_done()
  1873. # Check if query_result is None
  1874. if query_result is None:
  1875. return {
  1876. "status": "failure",
  1877. "message": "Query returned no results",
  1878. "data": {},
  1879. "metadata": {
  1880. "failure_reason": "no_results",
  1881. "mode": param.mode,
  1882. },
  1883. "llm_response": {
  1884. "content": PROMPTS["fail_response"],
  1885. "response_iterator": None,
  1886. "is_streaming": False,
  1887. },
  1888. }
  1889. # Extract structured data from query result
  1890. raw_data = query_result.raw_data or {}
  1891. raw_data["llm_response"] = {
  1892. "content": query_result.content
  1893. if not query_result.is_streaming
  1894. else None,
  1895. "response_iterator": query_result.response_iterator
  1896. if query_result.is_streaming
  1897. else None,
  1898. "is_streaming": query_result.is_streaming,
  1899. }
  1900. return raw_data
  1901. except Exception as e:
  1902. logger.error(f"Query failed: {e}")
  1903. # Return error response
  1904. return {
  1905. "status": "failure",
  1906. "message": f"Query failed: {str(e)}",
  1907. "data": {},
  1908. "metadata": {},
  1909. "llm_response": {
  1910. "content": None,
  1911. "response_iterator": None,
  1912. "is_streaming": False,
  1913. },
  1914. }
  1915. def query_llm(
  1916. self,
  1917. query: str,
  1918. param: QueryParam = QueryParam(),
  1919. system_prompt: str | None = None,
  1920. ) -> dict[str, Any]:
  1921. """
  1922. Synchronous complete query API: returns structured retrieval results with LLM generation.
  1923. This function is the synchronous version of aquery_llm, providing the same functionality
  1924. for users who prefer synchronous interfaces.
  1925. Args:
  1926. query: Query text for retrieval and LLM generation.
  1927. param: Query parameters controlling retrieval and LLM behavior.
  1928. system_prompt: Optional custom system prompt for LLM generation.
  1929. Returns:
  1930. dict[str, Any]: Same complete response format as aquery_llm.
  1931. """
  1932. loop = always_get_an_event_loop()
  1933. return loop.run_until_complete(self.aquery_llm(query, param, system_prompt))
  1934. async def _query_done(self):
  1935. await self.llm_response_cache.index_done_callback()
  1936. async def _update_delete_retry_state(
  1937. self,
  1938. doc_id: str,
  1939. doc_status_data: dict[str, Any],
  1940. *,
  1941. deletion_stage: str,
  1942. doc_llm_cache_ids: list[str],
  1943. error_message: str | None = None,
  1944. failed: bool,
  1945. ) -> dict[str, Any]:
  1946. """Persist deletion retry metadata and return the updated status record."""
  1947. metadata = doc_status_data.get("metadata", {})
  1948. if not isinstance(metadata, dict):
  1949. metadata = {}
  1950. backup_cache_ids = normalize_string_list(
  1951. metadata.get("deletion_llm_cache_ids", []),
  1952. context=f"doc {doc_id} metadata.deletion_llm_cache_ids",
  1953. )
  1954. retry_cache_ids = doc_llm_cache_ids or backup_cache_ids
  1955. updated_metadata = dict(metadata)
  1956. if retry_cache_ids:
  1957. updated_metadata["deletion_llm_cache_ids"] = retry_cache_ids
  1958. updated_metadata["last_deletion_attempt_at"] = datetime.now(
  1959. timezone.utc
  1960. ).isoformat()
  1961. if failed:
  1962. updated_metadata["deletion_failed"] = True
  1963. updated_metadata["deletion_failure_stage"] = deletion_stage
  1964. else:
  1965. updated_metadata.pop("deletion_failed", None)
  1966. updated_metadata.pop("deletion_failure_stage", None)
  1967. updated_status_data = {
  1968. **doc_status_data,
  1969. "updated_at": datetime.now(timezone.utc).isoformat(),
  1970. "metadata": updated_metadata,
  1971. "error_msg": error_message if failed else "",
  1972. }
  1973. await self.doc_status.upsert({doc_id: updated_status_data})
  1974. return updated_status_data
  1975. async def _get_existing_llm_cache_ids(self, cache_ids: list[str]) -> list[str]:
  1976. """Return cache IDs that still exist in cache storage.
  1977. Some KV storage backends only log delete failures and return without
  1978. raising, so callers must verify which records still exist after delete.
  1979. Returns an empty list immediately if cache storage is unavailable.
  1980. Callers must check storage availability independently before treating
  1981. an empty result as a confirmed deletion.
  1982. """
  1983. if not self.llm_response_cache or not cache_ids:
  1984. return []
  1985. try:
  1986. existing_records = await self.llm_response_cache.get_by_ids(cache_ids)
  1987. except Exception as verification_error:
  1988. raise Exception(
  1989. f"Failed to verify LLM cache deletion "
  1990. f"(delete may have succeeded): {verification_error}"
  1991. ) from verification_error
  1992. return [
  1993. cache_id
  1994. for cache_id, record in zip(cache_ids, existing_records)
  1995. if record is not None
  1996. ]
  1997. async def aclear_cache(self) -> None:
  1998. """Clear all cache data from the LLM response cache storage.
  1999. This method clears all cached LLM responses regardless of mode.
  2000. Example:
  2001. # Clear all cache
  2002. await rag.aclear_cache()
  2003. """
  2004. if not self.llm_response_cache:
  2005. logger.warning("No cache storage configured")
  2006. return
  2007. try:
  2008. # Clear all cache using drop method
  2009. success = await self.llm_response_cache.drop()
  2010. if success:
  2011. logger.info("Cleared all cache")
  2012. else:
  2013. logger.warning("Failed to clear all cache")
  2014. await self.llm_response_cache.index_done_callback()
  2015. except Exception as e:
  2016. logger.error(f"Error while clearing cache: {e}")
  2017. def clear_cache(self) -> None:
  2018. """Synchronous version of aclear_cache."""
  2019. return always_get_an_event_loop().run_until_complete(self.aclear_cache())
  2020. async def get_docs_by_status(
  2021. self, status: DocStatus
  2022. ) -> dict[str, DocProcessingStatus]:
  2023. """Get documents by status
  2024. Returns:
  2025. Dict with document id is keys and document status is values
  2026. """
  2027. return await self.doc_status.get_docs_by_status(status)
  2028. async def aget_docs_by_ids(
  2029. self, ids: str | list[str]
  2030. ) -> dict[str, DocProcessingStatus]:
  2031. """Retrieves the processing status for one or more documents by their IDs.
  2032. Args:
  2033. ids: A single document ID (string) or a list of document IDs (list of strings).
  2034. Returns:
  2035. A dictionary where keys are the document IDs for which a status was found,
  2036. and values are the corresponding DocProcessingStatus objects. IDs that
  2037. are not found in the storage will be omitted from the result dictionary.
  2038. """
  2039. if isinstance(ids, str):
  2040. # Ensure input is always a list of IDs for uniform processing
  2041. id_list = [ids]
  2042. elif (
  2043. ids is None
  2044. ): # Handle potential None input gracefully, although type hint suggests str/list
  2045. logger.warning(
  2046. "aget_docs_by_ids called with None input, returning empty dict."
  2047. )
  2048. return {}
  2049. else:
  2050. # Assume input is already a list if not a string
  2051. id_list = ids
  2052. # Return early if the final list of IDs is empty
  2053. if not id_list:
  2054. logger.debug("aget_docs_by_ids called with an empty list of IDs.")
  2055. return {}
  2056. # Create tasks to fetch document statuses concurrently using the doc_status storage
  2057. tasks = [self.doc_status.get_by_id(doc_id) for doc_id in id_list]
  2058. # Execute tasks concurrently and gather the results. Results maintain order.
  2059. # Type hint indicates results can be DocProcessingStatus or None if not found.
  2060. results_list: list[Optional[DocProcessingStatus]] = await asyncio.gather(*tasks)
  2061. # Build the result dictionary, mapping found IDs to their statuses
  2062. found_statuses: dict[str, DocProcessingStatus] = {}
  2063. # Keep track of IDs for which no status was found (for logging purposes)
  2064. not_found_ids: list[str] = []
  2065. # Iterate through the results, correlating them back to the original IDs
  2066. for i, status_obj in enumerate(results_list):
  2067. doc_id = id_list[
  2068. i
  2069. ] # Get the original ID corresponding to this result index
  2070. if status_obj:
  2071. # If a status object was returned (not None), add it to the result dict
  2072. found_statuses[doc_id] = status_obj
  2073. else:
  2074. # If status_obj is None, the document ID was not found in storage
  2075. not_found_ids.append(doc_id)
  2076. # Log a warning if any of the requested document IDs were not found
  2077. if not_found_ids:
  2078. logger.warning(
  2079. f"Document statuses not found for the following IDs: {not_found_ids}"
  2080. )
  2081. # Return the dictionary containing statuses only for the found document IDs
  2082. return found_statuses
  2083. async def _purge_doc_chunks_and_kg(
  2084. self,
  2085. doc_id: str,
  2086. chunk_ids: set[str],
  2087. *,
  2088. pipeline_status: dict,
  2089. pipeline_status_lock: Any,
  2090. ) -> None:
  2091. """Remove a document's chunks and clean up its knowledge-graph contributions.
  2092. Used by:
  2093. - The pipeline resume branch in ``process_document`` when a
  2094. document whose content is already extracted is re-processed
  2095. under different ``process_options``: chunks must be wiped and
  2096. entities/relations rebuilt fresh.
  2097. - Future deletion paths that want a focused "purge KG only"
  2098. operation without the LLM-cache / doc_status / full_docs
  2099. cleanup that ``adelete_by_doc_id`` also performs.
  2100. What this method does:
  2101. 1. Reads ``full_entities`` / ``full_relations`` to identify which
  2102. graph nodes / edges this document contributed to.
  2103. 2. For each affected entity / relation, intersects the doc's
  2104. ``chunk_ids`` with the union of chunk-tracking entries
  2105. (``entity_chunks`` / ``relation_chunks``) and graph
  2106. ``source_id`` lists, then classifies it as either
  2107. *delete-outright* (no remaining sources) or *rebuild*
  2108. (still references chunks from other documents).
  2109. 3. Deletes the chunks themselves from ``chunks_vdb`` and
  2110. ``text_chunks``.
  2111. 4. For *delete-outright* entries: removes the relationship /
  2112. entity from the graph storage, vector storage, and chunk
  2113. tracking.
  2114. 5. Calls :py:meth:`_insert_done` to persist graph changes
  2115. before rebuilding (so the rebuild step sees a consistent
  2116. state).
  2117. 6. Calls :func:`rebuild_knowledge_from_chunks` to rebuild any
  2118. *rebuild* entries from their remaining chunks (so other
  2119. documents that also contributed to the same entity /
  2120. relation keep their data intact).
  2121. 7. Deletes the per-doc ``full_entities`` / ``full_relations``
  2122. index rows so subsequent re-extraction starts fresh.
  2123. Does NOT touch:
  2124. - ``doc_status`` / ``full_docs`` records — caller manages those.
  2125. - ``llm_response_cache`` — orthogonal to KG cleanup.
  2126. - Pipeline busy-flag — assumes the caller already holds the
  2127. pipeline (i.e. this runs inside a pipeline run).
  2128. Idempotent: passing an empty ``chunk_ids`` returns immediately
  2129. without touching storage.
  2130. """
  2131. if not chunk_ids:
  2132. return
  2133. # ---- 1. Analyze affected entities/relations from full_entities/full_relations ----
  2134. entities_to_delete: set[str] = set()
  2135. entities_to_rebuild: dict[str, list[str]] = {}
  2136. relationships_to_delete: set[tuple[str, str]] = set()
  2137. relationships_to_rebuild: dict[tuple[str, str], list[str]] = {}
  2138. entity_chunk_updates: dict[str, list[str]] = {}
  2139. relation_chunk_updates: dict[tuple[str, str], list[str]] = {}
  2140. try:
  2141. doc_entities_data = await self.full_entities.get_by_id(doc_id)
  2142. doc_relations_data = await self.full_relations.get_by_id(doc_id)
  2143. affected_nodes: list[dict[str, Any]] = []
  2144. affected_edges: list[dict[str, Any]] = []
  2145. if doc_entities_data and "entity_names" in doc_entities_data:
  2146. entity_names = doc_entities_data["entity_names"]
  2147. nodes_dict = await self.chunk_entity_relation_graph.get_nodes_batch(
  2148. entity_names
  2149. )
  2150. for entity_name in entity_names:
  2151. node_data = nodes_dict.get(entity_name)
  2152. if node_data:
  2153. if "id" not in node_data:
  2154. node_data["id"] = entity_name
  2155. affected_nodes.append(node_data)
  2156. if doc_relations_data and "relation_pairs" in doc_relations_data:
  2157. relation_pairs = doc_relations_data["relation_pairs"]
  2158. edge_pairs_dicts = [
  2159. {"src": pair[0], "tgt": pair[1]} for pair in relation_pairs
  2160. ]
  2161. edges_dict = await self.chunk_entity_relation_graph.get_edges_batch(
  2162. edge_pairs_dicts
  2163. )
  2164. for pair in relation_pairs:
  2165. src, tgt = pair[0], pair[1]
  2166. edge_data = edges_dict.get((src, tgt))
  2167. if edge_data:
  2168. if "source" not in edge_data:
  2169. edge_data["source"] = src
  2170. if "target" not in edge_data:
  2171. edge_data["target"] = tgt
  2172. affected_edges.append(edge_data)
  2173. except Exception as e:
  2174. logger.error(
  2175. f"[purge] Failed to analyze affected graph elements for {doc_id}: {e}"
  2176. )
  2177. raise Exception(f"Failed to analyze graph dependencies: {e}") from e
  2178. # ---- 2. Classify entities/relations into delete vs rebuild ----
  2179. try:
  2180. for node_data in affected_nodes:
  2181. node_label = node_data.get("entity_id")
  2182. if not node_label:
  2183. continue
  2184. existing_sources: list[str] = []
  2185. graph_sources: list[str] = []
  2186. if self.entity_chunks:
  2187. stored_chunks = await self.entity_chunks.get_by_id(node_label)
  2188. if stored_chunks and isinstance(stored_chunks, dict):
  2189. existing_sources = [
  2190. chunk_id
  2191. for chunk_id in stored_chunks.get("chunk_ids", [])
  2192. if chunk_id
  2193. ]
  2194. if node_data.get("source_id"):
  2195. graph_sources = [
  2196. chunk_id
  2197. for chunk_id in node_data["source_id"].split(GRAPH_FIELD_SEP)
  2198. if chunk_id
  2199. ]
  2200. if not existing_sources:
  2201. existing_sources = graph_sources
  2202. if not existing_sources:
  2203. entities_to_delete.add(node_label)
  2204. entity_chunk_updates[node_label] = []
  2205. continue
  2206. remaining_sources = subtract_source_ids(existing_sources, chunk_ids)
  2207. graph_references_deleted_chunks = bool(
  2208. graph_sources and set(graph_sources) & chunk_ids
  2209. )
  2210. if not remaining_sources:
  2211. entities_to_delete.add(node_label)
  2212. entity_chunk_updates[node_label] = []
  2213. elif (
  2214. remaining_sources != existing_sources
  2215. or graph_references_deleted_chunks
  2216. ):
  2217. entities_to_rebuild[node_label] = remaining_sources
  2218. entity_chunk_updates[node_label] = remaining_sources
  2219. async with pipeline_status_lock:
  2220. log_message = (
  2221. f"[purge] {doc_id}: {len(entities_to_rebuild)} entity(ies) "
  2222. f"to rebuild, {len(entities_to_delete)} to delete"
  2223. )
  2224. logger.info(log_message)
  2225. pipeline_status["latest_message"] = log_message
  2226. pipeline_status["history_messages"].append(log_message)
  2227. for edge_data in affected_edges:
  2228. src = edge_data.get("source")
  2229. tgt = edge_data.get("target")
  2230. if not src or not tgt or "source_id" not in edge_data:
  2231. continue
  2232. edge_tuple = tuple(sorted((src, tgt)))
  2233. if (
  2234. edge_tuple in relationships_to_delete
  2235. or edge_tuple in relationships_to_rebuild
  2236. ):
  2237. continue
  2238. existing_sources = []
  2239. graph_sources = []
  2240. if self.relation_chunks:
  2241. storage_key = make_relation_chunk_key(src, tgt)
  2242. stored_chunks = await self.relation_chunks.get_by_id(storage_key)
  2243. if stored_chunks and isinstance(stored_chunks, dict):
  2244. existing_sources = [
  2245. chunk_id
  2246. for chunk_id in stored_chunks.get("chunk_ids", [])
  2247. if chunk_id
  2248. ]
  2249. if edge_data.get("source_id"):
  2250. graph_sources = [
  2251. chunk_id
  2252. for chunk_id in edge_data["source_id"].split(GRAPH_FIELD_SEP)
  2253. if chunk_id
  2254. ]
  2255. if not existing_sources:
  2256. existing_sources = graph_sources
  2257. if not existing_sources:
  2258. relationships_to_delete.add(edge_tuple)
  2259. relation_chunk_updates[edge_tuple] = []
  2260. continue
  2261. remaining_sources = subtract_source_ids(existing_sources, chunk_ids)
  2262. graph_references_deleted_chunks = bool(
  2263. graph_sources and set(graph_sources) & chunk_ids
  2264. )
  2265. if not remaining_sources:
  2266. relationships_to_delete.add(edge_tuple)
  2267. relation_chunk_updates[edge_tuple] = []
  2268. elif (
  2269. remaining_sources != existing_sources
  2270. or graph_references_deleted_chunks
  2271. ):
  2272. relationships_to_rebuild[edge_tuple] = remaining_sources
  2273. relation_chunk_updates[edge_tuple] = remaining_sources
  2274. async with pipeline_status_lock:
  2275. log_message = (
  2276. f"[purge] {doc_id}: {len(relationships_to_rebuild)} relation(s) "
  2277. f"to rebuild, {len(relationships_to_delete)} to delete"
  2278. )
  2279. logger.info(log_message)
  2280. pipeline_status["latest_message"] = log_message
  2281. pipeline_status["history_messages"].append(log_message)
  2282. # Update entity/relation chunk-tracking with the remaining sources.
  2283. current_time = int(time.time())
  2284. if entity_chunk_updates and self.entity_chunks:
  2285. entity_upsert_payload = {}
  2286. for entity_name, remaining in entity_chunk_updates.items():
  2287. if not remaining:
  2288. continue
  2289. entity_upsert_payload[entity_name] = {
  2290. "chunk_ids": remaining,
  2291. "count": len(remaining),
  2292. "updated_at": current_time,
  2293. }
  2294. if entity_upsert_payload:
  2295. await self.entity_chunks.upsert(entity_upsert_payload)
  2296. if relation_chunk_updates and self.relation_chunks:
  2297. relation_upsert_payload = {}
  2298. for edge_tuple, remaining in relation_chunk_updates.items():
  2299. if not remaining:
  2300. continue
  2301. storage_key = make_relation_chunk_key(*edge_tuple)
  2302. relation_upsert_payload[storage_key] = {
  2303. "chunk_ids": remaining,
  2304. "count": len(remaining),
  2305. "updated_at": current_time,
  2306. }
  2307. if relation_upsert_payload:
  2308. await self.relation_chunks.upsert(relation_upsert_payload)
  2309. except Exception as e:
  2310. logger.error(
  2311. f"[purge] Failed to process graph analysis results for {doc_id}: {e}"
  2312. )
  2313. raise Exception(f"Failed to process graph dependencies: {e}") from e
  2314. # ---- 3. Delete chunks themselves ----
  2315. try:
  2316. await self.chunks_vdb.delete(chunk_ids)
  2317. await self.text_chunks.delete(chunk_ids)
  2318. async with pipeline_status_lock:
  2319. log_message = (
  2320. f"[purge] {doc_id}: deleted {len(chunk_ids)} chunk(s) from storage"
  2321. )
  2322. logger.info(log_message)
  2323. pipeline_status["latest_message"] = log_message
  2324. pipeline_status["history_messages"].append(log_message)
  2325. except Exception as e:
  2326. logger.error(f"[purge] Failed to delete chunks for {doc_id}: {e}")
  2327. raise Exception(f"Failed to delete document chunks: {e}") from e
  2328. # ---- 4. Delete relationships with no remaining sources ----
  2329. if relationships_to_delete:
  2330. try:
  2331. rel_ids_to_delete = []
  2332. for src, tgt in relationships_to_delete:
  2333. rel_ids_to_delete.extend(
  2334. [
  2335. compute_mdhash_id(src + tgt, prefix="rel-"),
  2336. compute_mdhash_id(tgt + src, prefix="rel-"),
  2337. ]
  2338. )
  2339. await self.relationships_vdb.delete(rel_ids_to_delete)
  2340. await self.chunk_entity_relation_graph.remove_edges(
  2341. list(relationships_to_delete)
  2342. )
  2343. if self.relation_chunks:
  2344. relation_storage_keys = [
  2345. make_relation_chunk_key(src, tgt)
  2346. for src, tgt in relationships_to_delete
  2347. ]
  2348. await self.relation_chunks.delete(relation_storage_keys)
  2349. async with pipeline_status_lock:
  2350. log_message = (
  2351. f"[purge] {doc_id}: deleted "
  2352. f"{len(relationships_to_delete)} relation(s)"
  2353. )
  2354. logger.info(log_message)
  2355. pipeline_status["latest_message"] = log_message
  2356. pipeline_status["history_messages"].append(log_message)
  2357. except Exception as e:
  2358. logger.error(
  2359. f"[purge] Failed to delete relationships for {doc_id}: {e}"
  2360. )
  2361. raise Exception(f"Failed to delete relationships: {e}") from e
  2362. # ---- 5. Delete entities with no remaining sources ----
  2363. if entities_to_delete:
  2364. try:
  2365. nodes_edges_dict = (
  2366. await self.chunk_entity_relation_graph.get_nodes_edges_batch(
  2367. list(entities_to_delete)
  2368. )
  2369. )
  2370. edges_to_delete: set[tuple[str, str]] = set()
  2371. for entity, edges in nodes_edges_dict.items():
  2372. if edges:
  2373. for src, tgt in edges:
  2374. edges_to_delete.add(tuple(sorted((src, tgt))))
  2375. if edges_to_delete:
  2376. rel_ids_to_delete = []
  2377. for src, tgt in edges_to_delete:
  2378. rel_ids_to_delete.extend(
  2379. [
  2380. compute_mdhash_id(src + tgt, prefix="rel-"),
  2381. compute_mdhash_id(tgt + src, prefix="rel-"),
  2382. ]
  2383. )
  2384. await self.relationships_vdb.delete(rel_ids_to_delete)
  2385. if self.relation_chunks:
  2386. relation_storage_keys = [
  2387. make_relation_chunk_key(src, tgt)
  2388. for src, tgt in edges_to_delete
  2389. ]
  2390. await self.relation_chunks.delete(relation_storage_keys)
  2391. logger.info(
  2392. f"[purge] {doc_id}: cleaned {len(edges_to_delete)} residual "
  2393. f"edge(s) from VDB and chunk-tracking storage"
  2394. )
  2395. await self.chunk_entity_relation_graph.remove_nodes(
  2396. list(entities_to_delete)
  2397. )
  2398. entity_vdb_ids = [
  2399. compute_mdhash_id(entity, prefix="ent-")
  2400. for entity in entities_to_delete
  2401. ]
  2402. await self.entities_vdb.delete(entity_vdb_ids)
  2403. if self.entity_chunks:
  2404. await self.entity_chunks.delete(list(entities_to_delete))
  2405. async with pipeline_status_lock:
  2406. log_message = (
  2407. f"[purge] {doc_id}: deleted "
  2408. f"{len(entities_to_delete)} entity(ies)"
  2409. )
  2410. logger.info(log_message)
  2411. pipeline_status["latest_message"] = log_message
  2412. pipeline_status["history_messages"].append(log_message)
  2413. except Exception as e:
  2414. logger.error(f"[purge] Failed to delete entities for {doc_id}: {e}")
  2415. raise Exception(f"Failed to delete entities: {e}") from e
  2416. # ---- 6. Persist pre-rebuild changes ----
  2417. try:
  2418. await self._insert_done()
  2419. except Exception as e:
  2420. logger.error(f"[purge] Failed to persist pre-rebuild changes: {e}")
  2421. raise Exception(f"Failed to persist pre-rebuild changes: {e}") from e
  2422. # ---- 7. Rebuild entities/relations that still have remaining sources ----
  2423. if entities_to_rebuild or relationships_to_rebuild:
  2424. try:
  2425. await rebuild_knowledge_from_chunks(
  2426. entities_to_rebuild=entities_to_rebuild,
  2427. relationships_to_rebuild=relationships_to_rebuild,
  2428. knowledge_graph_inst=self.chunk_entity_relation_graph,
  2429. entities_vdb=self.entities_vdb,
  2430. relationships_vdb=self.relationships_vdb,
  2431. text_chunks_storage=self.text_chunks,
  2432. llm_response_cache=self.llm_response_cache,
  2433. global_config=self._build_global_config(),
  2434. pipeline_status=pipeline_status,
  2435. pipeline_status_lock=pipeline_status_lock,
  2436. entity_chunks_storage=self.entity_chunks,
  2437. relation_chunks_storage=self.relation_chunks,
  2438. )
  2439. except Exception as e:
  2440. logger.error(f"[purge] Failed to rebuild knowledge from chunks: {e}")
  2441. raise Exception(f"Failed to rebuild knowledge graph: {e}") from e
  2442. # ---- 8. Delete per-doc full_entities / full_relations index rows ----
  2443. try:
  2444. await self.full_entities.delete([doc_id])
  2445. await self.full_relations.delete([doc_id])
  2446. except Exception as e:
  2447. logger.error(
  2448. f"[purge] Failed to delete full_entities/full_relations rows for {doc_id}: {e}"
  2449. )
  2450. raise Exception(
  2451. f"Failed to delete from full_entities/full_relations: {e}"
  2452. ) from e
  2453. async def adelete_by_doc_id(
  2454. self, doc_id: str, delete_llm_cache: bool = False
  2455. ) -> DeletionResult:
  2456. """Delete a document and all its related data, including chunks, graph elements.
  2457. This method orchestrates a comprehensive deletion process for a given document ID.
  2458. It ensures that not only the document itself but also all its derived and associated
  2459. data across different storage layers are removed or rebuiled. If entities or relationships
  2460. are partially affected, they will be rebuilded using LLM cached from remaining documents.
  2461. **Concurrency Control Design:**
  2462. This function implements a pipeline-based concurrency control to prevent data corruption:
  2463. 1. **Single Document Deletion** (when WE acquire pipeline):
  2464. - Sets job_name to "Single document deletion" (NOT starting with "deleting")
  2465. - Prevents other adelete_by_doc_id calls from running concurrently
  2466. - Ensures exclusive access to graph operations for this deletion
  2467. 2. **Batch Document Deletion** (when background_delete_documents acquires pipeline):
  2468. - Sets job_name to "Deleting {N} Documents" (starts with "deleting")
  2469. - Allows multiple adelete_by_doc_id calls to join the deletion queue
  2470. - Each call validates the job name to ensure it's part of a deletion operation
  2471. The validation logic `if not job_name.startswith("deleting") or "document" not in job_name`
  2472. ensures that:
  2473. - adelete_by_doc_id can only run when pipeline is idle OR during batch deletion
  2474. - Prevents concurrent single deletions that could cause race conditions
  2475. - Rejects operations when pipeline is busy with non-deletion tasks
  2476. Args:
  2477. doc_id (str): The unique identifier of the document to be deleted.
  2478. delete_llm_cache (bool): Whether to delete cached LLM extraction results
  2479. associated with the document. Defaults to False.
  2480. Returns:
  2481. DeletionResult: An object containing the outcome of the deletion process.
  2482. - `status` (str): "success", "not_found", "not_allowed", or "fail".
  2483. - `doc_id` (str): The ID of the document attempted to be deleted.
  2484. - `message` (str): A summary of the operation's result.
  2485. - `status_code` (int): HTTP status code (e.g., 200, 404, 403, 500).
  2486. - `file_path` (str | None): The file path of the deleted document, if available.
  2487. """
  2488. # Get pipeline status shared data and lock for validation
  2489. pipeline_status = await get_namespace_data(
  2490. "pipeline_status", workspace=self.workspace
  2491. )
  2492. pipeline_status_lock = get_namespace_lock(
  2493. "pipeline_status", workspace=self.workspace
  2494. )
  2495. # Track whether WE acquired the pipeline
  2496. we_acquired_pipeline = False
  2497. # Check and acquire pipeline if needed
  2498. async with pipeline_status_lock:
  2499. if not pipeline_status.get("busy", False):
  2500. # Pipeline is idle - WE acquire it for this deletion
  2501. we_acquired_pipeline = True
  2502. pipeline_status.update(
  2503. {
  2504. "busy": True,
  2505. "job_name": "Single document deletion",
  2506. "job_start": datetime.now(timezone.utc).isoformat(),
  2507. "docs": 1,
  2508. "batchs": 1,
  2509. "cur_batch": 0,
  2510. "request_pending": False,
  2511. "cancellation_requested": False,
  2512. "latest_message": f"Starting deletion for document: {doc_id}",
  2513. }
  2514. )
  2515. # Initialize history messages
  2516. pipeline_status["history_messages"][:] = [
  2517. f"Starting deletion for document: {doc_id}"
  2518. ]
  2519. else:
  2520. # Pipeline already busy - verify it's a deletion job
  2521. job_name = pipeline_status.get("job_name", "").lower()
  2522. if not job_name.startswith("deleting") or "document" not in job_name:
  2523. return DeletionResult(
  2524. status="not_allowed",
  2525. doc_id=doc_id,
  2526. message=f"Deletion not allowed: current job '{pipeline_status.get('job_name')}' is not a document deletion job",
  2527. status_code=403,
  2528. file_path=None,
  2529. )
  2530. # Pipeline is busy with deletion - proceed without acquiring
  2531. deletion_operations_started = False
  2532. deletion_fully_completed = False
  2533. in_final_delete_stage = False
  2534. original_exception = None
  2535. doc_llm_cache_ids: list[str] = []
  2536. deletion_stage = "initializing"
  2537. doc_status_data: dict[str, Any] | None = None
  2538. file_path: str | None = None
  2539. async with pipeline_status_lock:
  2540. log_message = f"Starting deletion process for document {doc_id}"
  2541. logger.info(log_message)
  2542. pipeline_status["latest_message"] = log_message
  2543. pipeline_status["history_messages"].append(log_message)
  2544. try:
  2545. # 1. Get the document status and related data
  2546. doc_status_data = await self.doc_status.get_by_id(doc_id)
  2547. if not doc_status_data:
  2548. logger.warning(f"Document {doc_id} not found")
  2549. return DeletionResult(
  2550. status="not_found",
  2551. doc_id=doc_id,
  2552. message=f"Document {doc_id} not found.",
  2553. status_code=404,
  2554. file_path="",
  2555. )
  2556. file_path = doc_status_data.get("file_path")
  2557. # Check document status and log warning for non-completed documents
  2558. raw_status = doc_status_data.get("status")
  2559. try:
  2560. doc_status = DocStatus(raw_status)
  2561. except ValueError:
  2562. doc_status = raw_status
  2563. if doc_status != DocStatus.PROCESSED:
  2564. if doc_status == DocStatus.PENDING:
  2565. warning_msg = (
  2566. f"Deleting {doc_id} {file_path}(previous status: PENDING)"
  2567. )
  2568. elif doc_status == DocStatus.PROCESSING:
  2569. warning_msg = (
  2570. f"Deleting {doc_id} {file_path}(previous status: PROCESSING)"
  2571. )
  2572. elif doc_status == DocStatus.PREPROCESSED:
  2573. warning_msg = (
  2574. f"Deleting {doc_id} {file_path}(previous status: PREPROCESSED)"
  2575. )
  2576. elif doc_status == DocStatus.FAILED:
  2577. warning_msg = (
  2578. f"Deleting {doc_id} {file_path}(previous status: FAILED)"
  2579. )
  2580. else:
  2581. status_text = (
  2582. doc_status.value
  2583. if isinstance(doc_status, DocStatus)
  2584. else str(doc_status)
  2585. )
  2586. warning_msg = (
  2587. f"Deleting {doc_id} {file_path}(previous status: {status_text})"
  2588. )
  2589. logger.info(warning_msg)
  2590. # Update pipeline status for monitoring
  2591. async with pipeline_status_lock:
  2592. pipeline_status["latest_message"] = warning_msg
  2593. pipeline_status["history_messages"].append(warning_msg)
  2594. # 2. Get chunk IDs from document status
  2595. metadata = doc_status_data.get("metadata", {})
  2596. if not isinstance(metadata, dict):
  2597. metadata = {}
  2598. metadata_cache_ids = normalize_string_list(
  2599. metadata.get("deletion_llm_cache_ids", []),
  2600. context=f"doc {doc_id} metadata.deletion_llm_cache_ids",
  2601. )
  2602. chunk_ids = set(
  2603. normalize_string_list(
  2604. doc_status_data.get("chunks_list", []),
  2605. context=f"doc {doc_id} chunks_list",
  2606. )
  2607. )
  2608. if not chunk_ids:
  2609. logger.warning(f"No chunks found for document {doc_id}")
  2610. # Mark that deletion operations have started
  2611. deletion_operations_started = True
  2612. # A prior failed deletion may have collected LLM cache IDs before the
  2613. # chunks were removed. If delete_llm_cache is requested and persisted IDs
  2614. # exist, clean them up now before removing the doc/status entries.
  2615. if delete_llm_cache and metadata_cache_ids:
  2616. if not self.llm_response_cache:
  2617. no_cache_msg = (
  2618. f"Cannot delete LLM cache for document {doc_id}: "
  2619. "cache storage is unavailable"
  2620. )
  2621. logger.error(no_cache_msg)
  2622. async with pipeline_status_lock:
  2623. pipeline_status["latest_message"] = no_cache_msg
  2624. pipeline_status["history_messages"].append(no_cache_msg)
  2625. raise Exception(no_cache_msg)
  2626. try:
  2627. deletion_stage = "delete_llm_cache"
  2628. await self.llm_response_cache.delete(metadata_cache_ids)
  2629. remaining_cache_ids = await self._get_existing_llm_cache_ids(
  2630. metadata_cache_ids
  2631. )
  2632. if remaining_cache_ids:
  2633. raise Exception(
  2634. f"{len(remaining_cache_ids)} LLM cache entries still exist after delete"
  2635. )
  2636. logger.info(
  2637. "Cleaned up %d LLM cache entries from prior attempt for document %s",
  2638. len(metadata_cache_ids),
  2639. doc_id,
  2640. )
  2641. except Exception as cache_err:
  2642. raise Exception(
  2643. f"Failed to delete LLM cache for document {doc_id}: {cache_err}"
  2644. ) from cache_err
  2645. try:
  2646. # Still need to delete the doc status and full doc.
  2647. # Delete doc_status first: if full_docs.delete fails on retry, the
  2648. # doc_status record is already gone so the retry finds no record and
  2649. # treats the document as already deleted rather than creating a zombie.
  2650. deletion_stage = "delete_doc_entries"
  2651. await self.doc_status.delete([doc_id])
  2652. await self.full_docs.delete([doc_id])
  2653. except Exception as e:
  2654. logger.error(
  2655. f"Failed to delete document {doc_id} with no chunks: {e}"
  2656. )
  2657. raise Exception(f"Failed to delete document entry: {e}") from e
  2658. async with pipeline_status_lock:
  2659. log_message = (
  2660. f"Document deleted without associated chunks: {doc_id}"
  2661. )
  2662. logger.info(log_message)
  2663. pipeline_status["latest_message"] = log_message
  2664. pipeline_status["history_messages"].append(log_message)
  2665. deletion_fully_completed = True
  2666. return DeletionResult(
  2667. status="success",
  2668. doc_id=doc_id,
  2669. message=log_message,
  2670. status_code=200,
  2671. file_path=file_path,
  2672. )
  2673. # Mark that deletion operations have started
  2674. deletion_operations_started = True
  2675. if chunk_ids:
  2676. # Always collect/persist cache IDs for chunk-backed documents, even when
  2677. # this call does not request cache deletion. If a delete fails after the
  2678. # chunks/graph have already been removed, a later retry may turn on
  2679. # delete_llm_cache=True, and doc_status metadata is then the only durable
  2680. # place left to recover the cache keys for cleanup.
  2681. deletion_stage = "collect_llm_cache"
  2682. doc_llm_cache_ids = list(metadata_cache_ids)
  2683. if not self.text_chunks:
  2684. logger.info(
  2685. "Skipping LLM cache id collection for document %s because text chunk storage is unavailable",
  2686. doc_id,
  2687. )
  2688. else:
  2689. try:
  2690. chunk_data_list = await self.text_chunks.get_by_ids(
  2691. list(chunk_ids)
  2692. )
  2693. seen_cache_ids: set[str] = set(doc_llm_cache_ids)
  2694. for chunk_data in chunk_data_list:
  2695. if not chunk_data or not isinstance(chunk_data, dict):
  2696. continue
  2697. cache_ids = chunk_data.get("llm_cache_list", [])
  2698. if not isinstance(cache_ids, list):
  2699. continue
  2700. for cache_id in cache_ids:
  2701. if (
  2702. isinstance(cache_id, str)
  2703. and cache_id
  2704. and cache_id not in seen_cache_ids
  2705. ):
  2706. doc_llm_cache_ids.append(cache_id)
  2707. seen_cache_ids.add(cache_id)
  2708. except Exception as cache_collect_error:
  2709. logger.error(
  2710. "Failed to collect LLM cache ids for document %s: %s",
  2711. doc_id,
  2712. cache_collect_error,
  2713. )
  2714. raise Exception(
  2715. f"Failed to collect LLM cache ids for document {doc_id}: {cache_collect_error}"
  2716. ) from cache_collect_error
  2717. if doc_llm_cache_ids:
  2718. try:
  2719. doc_status_data = await self._update_delete_retry_state(
  2720. doc_id,
  2721. doc_status_data,
  2722. deletion_stage=deletion_stage,
  2723. doc_llm_cache_ids=doc_llm_cache_ids,
  2724. failed=False,
  2725. )
  2726. except Exception as status_write_error:
  2727. logger.error(
  2728. "Failed to persist LLM cache IDs for document %s to retry state: %s",
  2729. doc_id,
  2730. status_write_error,
  2731. )
  2732. # Describe whether this is a fresh attempt or a retry so
  2733. # operators can tell whether prior partial deletions exist.
  2734. attempt_context = (
  2735. "retry — prior partial deletions may exist"
  2736. if metadata_cache_ids
  2737. else "deletion not yet started"
  2738. )
  2739. raise Exception(
  2740. f"Failed to persist LLM cache IDs for document {doc_id} "
  2741. f"({attempt_context}): {status_write_error}"
  2742. ) from status_write_error
  2743. logger.info(
  2744. "Collected %d LLM cache entries for document %s",
  2745. len(doc_llm_cache_ids),
  2746. doc_id,
  2747. )
  2748. else:
  2749. logger.info("No LLM cache entries found for document %s", doc_id)
  2750. # 4. Analyze entities and relationships that will be affected
  2751. entities_to_delete = set()
  2752. entities_to_rebuild = {} # entity_name -> remaining chunk id list
  2753. relationships_to_delete = set()
  2754. relationships_to_rebuild = {} # (src, tgt) -> remaining chunk id list
  2755. entity_chunk_updates: dict[str, list[str]] = {}
  2756. relation_chunk_updates: dict[tuple[str, str], list[str]] = {}
  2757. try:
  2758. deletion_stage = "analyze_graph_dependencies"
  2759. # Get affected entities and relations from full_entities and full_relations storage
  2760. doc_entities_data = await self.full_entities.get_by_id(doc_id)
  2761. doc_relations_data = await self.full_relations.get_by_id(doc_id)
  2762. affected_nodes = []
  2763. affected_edges = []
  2764. # Get entity data from graph storage using entity names from full_entities
  2765. if doc_entities_data and "entity_names" in doc_entities_data:
  2766. entity_names = doc_entities_data["entity_names"]
  2767. # get_nodes_batch returns dict[str, dict], need to convert to list[dict]
  2768. nodes_dict = await self.chunk_entity_relation_graph.get_nodes_batch(
  2769. entity_names
  2770. )
  2771. for entity_name in entity_names:
  2772. node_data = nodes_dict.get(entity_name)
  2773. if node_data:
  2774. # Ensure compatibility with existing logic that expects "id" field
  2775. if "id" not in node_data:
  2776. node_data["id"] = entity_name
  2777. affected_nodes.append(node_data)
  2778. # Get relation data from graph storage using relation pairs from full_relations
  2779. if doc_relations_data and "relation_pairs" in doc_relations_data:
  2780. relation_pairs = doc_relations_data["relation_pairs"]
  2781. edge_pairs_dicts = [
  2782. {"src": pair[0], "tgt": pair[1]} for pair in relation_pairs
  2783. ]
  2784. # get_edges_batch returns dict[tuple[str, str], dict], need to convert to list[dict]
  2785. edges_dict = await self.chunk_entity_relation_graph.get_edges_batch(
  2786. edge_pairs_dicts
  2787. )
  2788. for pair in relation_pairs:
  2789. src, tgt = pair[0], pair[1]
  2790. edge_key = (src, tgt)
  2791. edge_data = edges_dict.get(edge_key)
  2792. if edge_data:
  2793. # Ensure compatibility with existing logic that expects "source" and "target" fields
  2794. if "source" not in edge_data:
  2795. edge_data["source"] = src
  2796. if "target" not in edge_data:
  2797. edge_data["target"] = tgt
  2798. affected_edges.append(edge_data)
  2799. except Exception as e:
  2800. logger.error(f"Failed to analyze affected graph elements: {e}")
  2801. raise Exception(f"Failed to analyze graph dependencies: {e}") from e
  2802. try:
  2803. # Process entities
  2804. for node_data in affected_nodes:
  2805. node_label = node_data.get("entity_id")
  2806. if not node_label:
  2807. continue
  2808. existing_sources: list[str] = []
  2809. graph_sources: list[str] = []
  2810. if self.entity_chunks:
  2811. stored_chunks = await self.entity_chunks.get_by_id(node_label)
  2812. if stored_chunks and isinstance(stored_chunks, dict):
  2813. existing_sources = [
  2814. chunk_id
  2815. for chunk_id in stored_chunks.get("chunk_ids", [])
  2816. if chunk_id
  2817. ]
  2818. if node_data.get("source_id"):
  2819. graph_sources = [
  2820. chunk_id
  2821. for chunk_id in node_data["source_id"].split(
  2822. GRAPH_FIELD_SEP
  2823. )
  2824. if chunk_id
  2825. ]
  2826. if not existing_sources:
  2827. existing_sources = graph_sources
  2828. if not existing_sources:
  2829. # No chunk references means this entity should be deleted
  2830. entities_to_delete.add(node_label)
  2831. entity_chunk_updates[node_label] = []
  2832. continue
  2833. remaining_sources = subtract_source_ids(existing_sources, chunk_ids)
  2834. # `existing_sources` comes from chunk-tracking storage when available, but
  2835. # graph `source_id` can still be stale after a failed prior delete. If the
  2836. # graph still references any chunk being deleted in this attempt, force a
  2837. # rebuild/delete so the graph metadata gets synchronized instead of being
  2838. # left untouched with orphaned source references.
  2839. graph_references_deleted_chunks = bool(
  2840. graph_sources and set(graph_sources) & chunk_ids
  2841. )
  2842. if not remaining_sources:
  2843. entities_to_delete.add(node_label)
  2844. entity_chunk_updates[node_label] = []
  2845. elif (
  2846. remaining_sources != existing_sources
  2847. or graph_references_deleted_chunks
  2848. ):
  2849. entities_to_rebuild[node_label] = remaining_sources
  2850. entity_chunk_updates[node_label] = remaining_sources
  2851. else:
  2852. logger.info(f"Untouch entity: {node_label}")
  2853. async with pipeline_status_lock:
  2854. log_message = f"Found {len(entities_to_rebuild)} affected entities"
  2855. logger.info(log_message)
  2856. pipeline_status["latest_message"] = log_message
  2857. pipeline_status["history_messages"].append(log_message)
  2858. # Process relationships
  2859. for edge_data in affected_edges:
  2860. # source target is not in normalize order in graph db property
  2861. src = edge_data.get("source")
  2862. tgt = edge_data.get("target")
  2863. if not src or not tgt or "source_id" not in edge_data:
  2864. continue
  2865. edge_tuple = tuple(sorted((src, tgt)))
  2866. if (
  2867. edge_tuple in relationships_to_delete
  2868. or edge_tuple in relationships_to_rebuild
  2869. ):
  2870. continue
  2871. existing_sources: list[str] = []
  2872. graph_sources: list[str] = []
  2873. if self.relation_chunks:
  2874. storage_key = make_relation_chunk_key(src, tgt)
  2875. stored_chunks = await self.relation_chunks.get_by_id(
  2876. storage_key
  2877. )
  2878. if stored_chunks and isinstance(stored_chunks, dict):
  2879. existing_sources = [
  2880. chunk_id
  2881. for chunk_id in stored_chunks.get("chunk_ids", [])
  2882. if chunk_id
  2883. ]
  2884. if edge_data.get("source_id"):
  2885. graph_sources = [
  2886. chunk_id
  2887. for chunk_id in edge_data["source_id"].split(
  2888. GRAPH_FIELD_SEP
  2889. )
  2890. if chunk_id
  2891. ]
  2892. if not existing_sources:
  2893. existing_sources = graph_sources
  2894. if not existing_sources:
  2895. # No chunk references means this relationship should be deleted
  2896. relationships_to_delete.add(edge_tuple)
  2897. relation_chunk_updates[edge_tuple] = []
  2898. continue
  2899. remaining_sources = subtract_source_ids(existing_sources, chunk_ids)
  2900. # Same as the entity path above: even when relation chunk-tracking is already
  2901. # correct, the graph edge may still carry a stale `source_id` that mentions a
  2902. # chunk deleted in this attempt. Treat that as an affected relation so retry
  2903. # deletion can repair the graph metadata rather than skipping it as "untouched".
  2904. graph_references_deleted_chunks = bool(
  2905. graph_sources and set(graph_sources) & chunk_ids
  2906. )
  2907. if not remaining_sources:
  2908. relationships_to_delete.add(edge_tuple)
  2909. relation_chunk_updates[edge_tuple] = []
  2910. elif (
  2911. remaining_sources != existing_sources
  2912. or graph_references_deleted_chunks
  2913. ):
  2914. relationships_to_rebuild[edge_tuple] = remaining_sources
  2915. relation_chunk_updates[edge_tuple] = remaining_sources
  2916. else:
  2917. logger.info(f"Untouch relation: {edge_tuple}")
  2918. async with pipeline_status_lock:
  2919. log_message = (
  2920. f"Found {len(relationships_to_rebuild)} affected relations"
  2921. )
  2922. logger.info(log_message)
  2923. pipeline_status["latest_message"] = log_message
  2924. pipeline_status["history_messages"].append(log_message)
  2925. current_time = int(time.time())
  2926. deletion_stage = "update_chunk_tracking"
  2927. if entity_chunk_updates and self.entity_chunks:
  2928. entity_upsert_payload = {}
  2929. for entity_name, remaining in entity_chunk_updates.items():
  2930. if not remaining:
  2931. # Empty entities are deleted alongside graph nodes later
  2932. continue
  2933. entity_upsert_payload[entity_name] = {
  2934. "chunk_ids": remaining,
  2935. "count": len(remaining),
  2936. "updated_at": current_time,
  2937. }
  2938. if entity_upsert_payload:
  2939. await self.entity_chunks.upsert(entity_upsert_payload)
  2940. if relation_chunk_updates and self.relation_chunks:
  2941. relation_upsert_payload = {}
  2942. for edge_tuple, remaining in relation_chunk_updates.items():
  2943. if not remaining:
  2944. # Empty relations are deleted alongside graph edges later
  2945. continue
  2946. storage_key = make_relation_chunk_key(*edge_tuple)
  2947. relation_upsert_payload[storage_key] = {
  2948. "chunk_ids": remaining,
  2949. "count": len(remaining),
  2950. "updated_at": current_time,
  2951. }
  2952. if relation_upsert_payload:
  2953. await self.relation_chunks.upsert(relation_upsert_payload)
  2954. except Exception as e:
  2955. logger.error(f"Failed to process graph analysis results: {e}")
  2956. raise Exception(f"Failed to process graph dependencies: {e}") from e
  2957. # Data integrity is ensured by allowing only one process to hold pipeline at a time(no graph db lock is needed anymore)
  2958. # 5. Delete chunks from storage
  2959. if chunk_ids:
  2960. try:
  2961. deletion_stage = "delete_chunks"
  2962. await self.chunks_vdb.delete(chunk_ids)
  2963. await self.text_chunks.delete(chunk_ids)
  2964. async with pipeline_status_lock:
  2965. log_message = (
  2966. f"Successfully deleted {len(chunk_ids)} chunks from storage"
  2967. )
  2968. logger.info(log_message)
  2969. pipeline_status["latest_message"] = log_message
  2970. pipeline_status["history_messages"].append(log_message)
  2971. except Exception as e:
  2972. logger.error(f"Failed to delete chunks: {e}")
  2973. raise Exception(f"Failed to delete document chunks: {e}") from e
  2974. # 6. Delete relationships that have no remaining sources
  2975. if relationships_to_delete:
  2976. try:
  2977. deletion_stage = "delete_relationships"
  2978. # Delete from relation vdb
  2979. rel_ids_to_delete = []
  2980. for src, tgt in relationships_to_delete:
  2981. rel_ids_to_delete.extend(
  2982. [
  2983. compute_mdhash_id(src + tgt, prefix="rel-"),
  2984. compute_mdhash_id(tgt + src, prefix="rel-"),
  2985. ]
  2986. )
  2987. await self.relationships_vdb.delete(rel_ids_to_delete)
  2988. # Delete from graph
  2989. await self.chunk_entity_relation_graph.remove_edges(
  2990. list(relationships_to_delete)
  2991. )
  2992. # Delete from relation_chunks storage
  2993. if self.relation_chunks:
  2994. relation_storage_keys = [
  2995. make_relation_chunk_key(src, tgt)
  2996. for src, tgt in relationships_to_delete
  2997. ]
  2998. await self.relation_chunks.delete(relation_storage_keys)
  2999. async with pipeline_status_lock:
  3000. log_message = f"Successfully deleted {len(relationships_to_delete)} relations"
  3001. logger.info(log_message)
  3002. pipeline_status["latest_message"] = log_message
  3003. pipeline_status["history_messages"].append(log_message)
  3004. except Exception as e:
  3005. logger.error(f"Failed to delete relationships: {e}")
  3006. raise Exception(f"Failed to delete relationships: {e}") from e
  3007. # 7. Delete entities that have no remaining sources
  3008. if entities_to_delete:
  3009. try:
  3010. deletion_stage = "delete_entities"
  3011. # Batch get all edges for entities to avoid N+1 query problem
  3012. nodes_edges_dict = (
  3013. await self.chunk_entity_relation_graph.get_nodes_edges_batch(
  3014. list(entities_to_delete)
  3015. )
  3016. )
  3017. # Debug: Check and log all edges before deleting nodes
  3018. edges_to_delete = set()
  3019. edges_still_exist = 0
  3020. for entity, edges in nodes_edges_dict.items():
  3021. if edges:
  3022. for src, tgt in edges:
  3023. # Normalize edge representation (sorted for consistency)
  3024. edge_tuple = tuple(sorted((src, tgt)))
  3025. edges_to_delete.add(edge_tuple)
  3026. if (
  3027. src in entities_to_delete
  3028. and tgt in entities_to_delete
  3029. ):
  3030. logger.warning(
  3031. f"Edge still exists: {src} <-> {tgt}"
  3032. )
  3033. elif src in entities_to_delete:
  3034. logger.warning(
  3035. f"Edge still exists: {src} --> {tgt}"
  3036. )
  3037. else:
  3038. logger.warning(
  3039. f"Edge still exists: {src} <-- {tgt}"
  3040. )
  3041. edges_still_exist += 1
  3042. if edges_still_exist:
  3043. logger.warning(
  3044. f"⚠️ {edges_still_exist} entities still has edges before deletion"
  3045. )
  3046. # Clean residual edges from VDB and storage before deleting nodes
  3047. if edges_to_delete:
  3048. # Delete from relationships_vdb
  3049. rel_ids_to_delete = []
  3050. for src, tgt in edges_to_delete:
  3051. rel_ids_to_delete.extend(
  3052. [
  3053. compute_mdhash_id(src + tgt, prefix="rel-"),
  3054. compute_mdhash_id(tgt + src, prefix="rel-"),
  3055. ]
  3056. )
  3057. await self.relationships_vdb.delete(rel_ids_to_delete)
  3058. # Delete from relation_chunks storage
  3059. if self.relation_chunks:
  3060. relation_storage_keys = [
  3061. make_relation_chunk_key(src, tgt)
  3062. for src, tgt in edges_to_delete
  3063. ]
  3064. await self.relation_chunks.delete(relation_storage_keys)
  3065. logger.info(
  3066. f"Cleaned {len(edges_to_delete)} residual edges from VDB and chunk-tracking storage"
  3067. )
  3068. # Delete from graph (edges will be auto-deleted with nodes)
  3069. await self.chunk_entity_relation_graph.remove_nodes(
  3070. list(entities_to_delete)
  3071. )
  3072. # Delete from vector vdb
  3073. entity_vdb_ids = [
  3074. compute_mdhash_id(entity, prefix="ent-")
  3075. for entity in entities_to_delete
  3076. ]
  3077. await self.entities_vdb.delete(entity_vdb_ids)
  3078. # Delete from entity_chunks storage
  3079. if self.entity_chunks:
  3080. await self.entity_chunks.delete(list(entities_to_delete))
  3081. async with pipeline_status_lock:
  3082. log_message = (
  3083. f"Successfully deleted {len(entities_to_delete)} entities"
  3084. )
  3085. logger.info(log_message)
  3086. pipeline_status["latest_message"] = log_message
  3087. pipeline_status["history_messages"].append(log_message)
  3088. except Exception as e:
  3089. logger.error(f"Failed to delete entities: {e}")
  3090. raise Exception(f"Failed to delete entities: {e}") from e
  3091. # Persist changes to graph database before entity and relationship rebuild
  3092. try:
  3093. deletion_stage = "persist_pre_rebuild_changes"
  3094. await self._insert_done()
  3095. except Exception as e:
  3096. logger.error(f"Failed to persist pre-rebuild changes: {e}")
  3097. raise Exception(f"Failed to persist pre-rebuild changes: {e}") from e
  3098. # 8. Rebuild entities and relationships from remaining chunks
  3099. if entities_to_rebuild or relationships_to_rebuild:
  3100. try:
  3101. deletion_stage = "rebuild_knowledge_graph"
  3102. await rebuild_knowledge_from_chunks(
  3103. entities_to_rebuild=entities_to_rebuild,
  3104. relationships_to_rebuild=relationships_to_rebuild,
  3105. knowledge_graph_inst=self.chunk_entity_relation_graph,
  3106. entities_vdb=self.entities_vdb,
  3107. relationships_vdb=self.relationships_vdb,
  3108. text_chunks_storage=self.text_chunks,
  3109. llm_response_cache=self.llm_response_cache,
  3110. global_config=self._build_global_config(),
  3111. pipeline_status=pipeline_status,
  3112. pipeline_status_lock=pipeline_status_lock,
  3113. entity_chunks_storage=self.entity_chunks,
  3114. relation_chunks_storage=self.relation_chunks,
  3115. )
  3116. except Exception as e:
  3117. logger.error(f"Failed to rebuild knowledge from chunks: {e}")
  3118. raise Exception(f"Failed to rebuild knowledge graph: {e}") from e
  3119. # 9. Delete LLM cache while the document status still exists so a failure
  3120. # remains retryable via the same doc_id.
  3121. log_message = f"Document {doc_id} successfully deleted"
  3122. if delete_llm_cache and doc_llm_cache_ids:
  3123. if not self.llm_response_cache:
  3124. log_message = (
  3125. f"Cannot delete LLM cache for document {doc_id}: "
  3126. "cache storage is unavailable"
  3127. )
  3128. logger.error(log_message)
  3129. async with pipeline_status_lock:
  3130. pipeline_status["latest_message"] = log_message
  3131. pipeline_status["history_messages"].append(log_message)
  3132. raise Exception(log_message)
  3133. try:
  3134. deletion_stage = "delete_llm_cache"
  3135. await self.llm_response_cache.delete(doc_llm_cache_ids)
  3136. # Some storage implementations do not raise on delete errors and
  3137. # instead only log internally, so confirm the cache entries are
  3138. # actually gone before deleting the document/status records.
  3139. remaining_cache_ids = await self._get_existing_llm_cache_ids(
  3140. doc_llm_cache_ids
  3141. )
  3142. if remaining_cache_ids:
  3143. doc_llm_cache_ids = remaining_cache_ids
  3144. raise Exception(
  3145. f"{len(remaining_cache_ids)} LLM cache entries still exist after delete"
  3146. )
  3147. cache_log_message = f"Successfully deleted {len(doc_llm_cache_ids)} LLM cache entries for document {doc_id}"
  3148. logger.info(cache_log_message)
  3149. async with pipeline_status_lock:
  3150. pipeline_status["latest_message"] = cache_log_message
  3151. pipeline_status["history_messages"].append(cache_log_message)
  3152. log_message = cache_log_message
  3153. except Exception as cache_delete_error:
  3154. log_message = (
  3155. f"Failed to delete LLM cache for document {doc_id}: "
  3156. f"{cache_delete_error}"
  3157. )
  3158. logger.error(log_message)
  3159. logger.error(traceback.format_exc())
  3160. async with pipeline_status_lock:
  3161. pipeline_status["latest_message"] = log_message
  3162. pipeline_status["history_messages"].append(log_message)
  3163. raise Exception(log_message) from cache_delete_error
  3164. # 10. Delete from full_entities and full_relations storage
  3165. try:
  3166. deletion_stage = "delete_doc_graph_metadata"
  3167. await self.full_entities.delete([doc_id])
  3168. await self.full_relations.delete([doc_id])
  3169. except Exception as e:
  3170. logger.error(f"Failed to delete from full_entities/full_relations: {e}")
  3171. raise Exception(
  3172. f"Failed to delete from full_entities/full_relations: {e}"
  3173. ) from e
  3174. # 11. Delete original document and status.
  3175. # doc_status is deleted first so that if full_docs.delete fails, a retry
  3176. # finds no doc_status record and treats the document as already gone,
  3177. # rather than finding a doc_status that points to a missing full_docs entry.
  3178. try:
  3179. deletion_stage = "delete_doc_entries"
  3180. in_final_delete_stage = True
  3181. await self.doc_status.delete([doc_id])
  3182. await self.full_docs.delete([doc_id])
  3183. except Exception as e:
  3184. logger.error(f"Failed to delete document and status: {e}")
  3185. raise Exception(f"Failed to delete document and status: {e}") from e
  3186. deletion_fully_completed = True
  3187. return DeletionResult(
  3188. status="success",
  3189. doc_id=doc_id,
  3190. message=log_message,
  3191. status_code=200,
  3192. file_path=file_path,
  3193. )
  3194. except Exception as e:
  3195. original_exception = e
  3196. error_message = f"Error while deleting document {doc_id}: {e}"
  3197. logger.error(error_message)
  3198. logger.error(traceback.format_exc())
  3199. try:
  3200. # Do not attempt to write retry state if doc_status was already deleted:
  3201. # upsert would re-create the record as a zombie. All earlier stages still
  3202. # have doc_status intact and can safely update it, even if some chunk/graph
  3203. # data has already been removed.
  3204. if doc_status_data is not None and not in_final_delete_stage:
  3205. doc_status_data = await self._update_delete_retry_state(
  3206. doc_id,
  3207. doc_status_data,
  3208. deletion_stage=deletion_stage,
  3209. doc_llm_cache_ids=doc_llm_cache_ids,
  3210. error_message=error_message,
  3211. failed=True,
  3212. )
  3213. except Exception as status_update_error:
  3214. logger.error(
  3215. "Failed to update deletion retry state for document %s: %s",
  3216. doc_id,
  3217. status_update_error,
  3218. )
  3219. logger.error(traceback.format_exc())
  3220. error_message = (
  3221. f"{error_message}. Additionally, failed to persist retry state: "
  3222. f"{status_update_error}. Manual cleanup may be required."
  3223. )
  3224. return DeletionResult(
  3225. status="fail",
  3226. doc_id=doc_id,
  3227. message=error_message,
  3228. status_code=500,
  3229. file_path=file_path,
  3230. )
  3231. finally:
  3232. # ALWAYS ensure persistence if any deletion operations were started
  3233. if deletion_operations_started:
  3234. try:
  3235. await self._insert_done()
  3236. except Exception as persistence_error:
  3237. persistence_error_msg = f"Failed to persist data after deletion attempt for {doc_id}: {persistence_error}"
  3238. logger.error(persistence_error_msg)
  3239. logger.error(traceback.format_exc())
  3240. if deletion_fully_completed:
  3241. # All deletion stages succeeded; the flush error is a post-cleanup
  3242. # concern. Do not override the success result already returned.
  3243. logger.error(
  3244. "Post-deletion persistence flush failed for %s, "
  3245. "but deletion completed successfully: %s",
  3246. doc_id,
  3247. persistence_error,
  3248. )
  3249. elif original_exception is None:
  3250. # Deletion stages were in-flight but the try-block return was never
  3251. # reached; treat the persistence failure as the primary error.
  3252. return DeletionResult(
  3253. status="fail",
  3254. doc_id=doc_id,
  3255. message=f"Deletion completed but failed to persist changes: {persistence_error}",
  3256. status_code=500,
  3257. file_path=file_path,
  3258. )
  3259. # If there was an original exception, log the persistence error but
  3260. # don't override it — the original error result was already returned.
  3261. else:
  3262. logger.debug(
  3263. f"No deletion operations were started for document {doc_id}, skipping persistence"
  3264. )
  3265. # Release pipeline only if WE acquired it
  3266. if we_acquired_pipeline:
  3267. async with pipeline_status_lock:
  3268. pipeline_status["busy"] = False
  3269. pipeline_status["cancellation_requested"] = False
  3270. completion_msg = (
  3271. f"Deletion process completed for document: {doc_id}"
  3272. )
  3273. pipeline_status["latest_message"] = completion_msg
  3274. pipeline_status["history_messages"].append(completion_msg)
  3275. logger.info(completion_msg)
  3276. async def adelete_by_entity(self, entity_name: str) -> DeletionResult:
  3277. """Asynchronously delete an entity and all its relationships.
  3278. Args:
  3279. entity_name: Name of the entity to delete.
  3280. Returns:
  3281. DeletionResult: An object containing the outcome of the deletion process.
  3282. """
  3283. from lightrag.utils_graph import adelete_by_entity
  3284. return await adelete_by_entity(
  3285. self.chunk_entity_relation_graph,
  3286. self.entities_vdb,
  3287. self.relationships_vdb,
  3288. entity_name,
  3289. )
  3290. def delete_by_entity(self, entity_name: str) -> DeletionResult:
  3291. """Synchronously delete an entity and all its relationships.
  3292. Args:
  3293. entity_name: Name of the entity to delete.
  3294. Returns:
  3295. DeletionResult: An object containing the outcome of the deletion process.
  3296. """
  3297. loop = always_get_an_event_loop()
  3298. return loop.run_until_complete(self.adelete_by_entity(entity_name))
  3299. async def adelete_by_relation(
  3300. self, source_entity: str, target_entity: str
  3301. ) -> DeletionResult:
  3302. """Asynchronously delete a relation between two entities.
  3303. Args:
  3304. source_entity: Name of the source entity.
  3305. target_entity: Name of the target entity.
  3306. Returns:
  3307. DeletionResult: An object containing the outcome of the deletion process.
  3308. """
  3309. from lightrag.utils_graph import adelete_by_relation
  3310. return await adelete_by_relation(
  3311. self.chunk_entity_relation_graph,
  3312. self.relationships_vdb,
  3313. source_entity,
  3314. target_entity,
  3315. )
  3316. def delete_by_relation(
  3317. self, source_entity: str, target_entity: str
  3318. ) -> DeletionResult:
  3319. """Synchronously delete a relation between two entities.
  3320. Args:
  3321. source_entity: Name of the source entity.
  3322. target_entity: Name of the target entity.
  3323. Returns:
  3324. DeletionResult: An object containing the outcome of the deletion process.
  3325. """
  3326. loop = always_get_an_event_loop()
  3327. return loop.run_until_complete(
  3328. self.adelete_by_relation(source_entity, target_entity)
  3329. )
  3330. async def get_processing_status(self) -> dict[str, int]:
  3331. """Get current document processing status counts
  3332. Returns:
  3333. Dict with counts for each status
  3334. """
  3335. return await self.doc_status.get_status_counts()
  3336. async def aget_docs_by_track_id(
  3337. self, track_id: str
  3338. ) -> dict[str, DocProcessingStatus]:
  3339. """Get documents by track_id
  3340. Args:
  3341. track_id: The tracking ID to search for
  3342. Returns:
  3343. Dict with document id as keys and document status as values
  3344. """
  3345. return await self.doc_status.get_docs_by_track_id(track_id)
  3346. async def get_entity_info(
  3347. self, entity_name: str, include_vector_data: bool = False
  3348. ) -> dict[str, str | None | dict[str, str]]:
  3349. """Get detailed information of an entity"""
  3350. from lightrag.utils_graph import get_entity_info
  3351. return await get_entity_info(
  3352. self.chunk_entity_relation_graph,
  3353. self.entities_vdb,
  3354. entity_name,
  3355. include_vector_data,
  3356. )
  3357. async def get_relation_info(
  3358. self, src_entity: str, tgt_entity: str, include_vector_data: bool = False
  3359. ) -> dict[str, str | None | dict[str, str]]:
  3360. """Get detailed information of a relationship"""
  3361. from lightrag.utils_graph import get_relation_info
  3362. return await get_relation_info(
  3363. self.chunk_entity_relation_graph,
  3364. self.relationships_vdb,
  3365. src_entity,
  3366. tgt_entity,
  3367. include_vector_data,
  3368. )
  3369. async def aedit_entity(
  3370. self,
  3371. entity_name: str,
  3372. updated_data: dict[str, str],
  3373. allow_rename: bool = True,
  3374. allow_merge: bool = False,
  3375. ) -> dict[str, Any]:
  3376. """Asynchronously edit entity information.
  3377. Updates entity information in the knowledge graph and re-embeds the entity in the vector database.
  3378. Also synchronizes entity_chunks_storage and relation_chunks_storage to track chunk references.
  3379. Args:
  3380. entity_name: Name of the entity to edit
  3381. updated_data: Dictionary containing updated attributes, e.g. {"description": "new description", "entity_type": "new type"}
  3382. allow_rename: Whether to allow entity renaming, defaults to True
  3383. allow_merge: Whether to merge into an existing entity when renaming to an existing name
  3384. Returns:
  3385. Dictionary containing updated entity information
  3386. """
  3387. from lightrag.utils_graph import aedit_entity
  3388. return await aedit_entity(
  3389. self.chunk_entity_relation_graph,
  3390. self.entities_vdb,
  3391. self.relationships_vdb,
  3392. entity_name,
  3393. updated_data,
  3394. allow_rename,
  3395. allow_merge,
  3396. self.entity_chunks,
  3397. self.relation_chunks,
  3398. )
  3399. def edit_entity(
  3400. self,
  3401. entity_name: str,
  3402. updated_data: dict[str, str],
  3403. allow_rename: bool = True,
  3404. allow_merge: bool = False,
  3405. ) -> dict[str, Any]:
  3406. loop = always_get_an_event_loop()
  3407. return loop.run_until_complete(
  3408. self.aedit_entity(entity_name, updated_data, allow_rename, allow_merge)
  3409. )
  3410. async def aedit_relation(
  3411. self, source_entity: str, target_entity: str, updated_data: dict[str, Any]
  3412. ) -> dict[str, Any]:
  3413. """Asynchronously edit relation information.
  3414. Updates relation (edge) information in the knowledge graph and re-embeds the relation in the vector database.
  3415. Also synchronizes the relation_chunks_storage to track which chunks reference this relation.
  3416. Args:
  3417. source_entity: Name of the source entity
  3418. target_entity: Name of the target entity
  3419. updated_data: Dictionary containing updated attributes, e.g. {"description": "new description", "keywords": "new keywords"}
  3420. Returns:
  3421. Dictionary containing updated relation information
  3422. """
  3423. from lightrag.utils_graph import aedit_relation
  3424. return await aedit_relation(
  3425. self.chunk_entity_relation_graph,
  3426. self.entities_vdb,
  3427. self.relationships_vdb,
  3428. source_entity,
  3429. target_entity,
  3430. updated_data,
  3431. self.relation_chunks,
  3432. )
  3433. def edit_relation(
  3434. self, source_entity: str, target_entity: str, updated_data: dict[str, Any]
  3435. ) -> dict[str, Any]:
  3436. loop = always_get_an_event_loop()
  3437. return loop.run_until_complete(
  3438. self.aedit_relation(source_entity, target_entity, updated_data)
  3439. )
  3440. async def acreate_entity(
  3441. self, entity_name: str, entity_data: dict[str, Any]
  3442. ) -> dict[str, Any]:
  3443. """Asynchronously create a new entity.
  3444. Creates a new entity in the knowledge graph and adds it to the vector database.
  3445. Args:
  3446. entity_name: Name of the new entity
  3447. entity_data: Dictionary containing entity attributes, e.g. {"description": "description", "entity_type": "type"}
  3448. Returns:
  3449. Dictionary containing created entity information
  3450. """
  3451. from lightrag.utils_graph import acreate_entity
  3452. return await acreate_entity(
  3453. self.chunk_entity_relation_graph,
  3454. self.entities_vdb,
  3455. self.relationships_vdb,
  3456. entity_name,
  3457. entity_data,
  3458. )
  3459. def create_entity(
  3460. self, entity_name: str, entity_data: dict[str, Any]
  3461. ) -> dict[str, Any]:
  3462. loop = always_get_an_event_loop()
  3463. return loop.run_until_complete(self.acreate_entity(entity_name, entity_data))
  3464. async def acreate_relation(
  3465. self, source_entity: str, target_entity: str, relation_data: dict[str, Any]
  3466. ) -> dict[str, Any]:
  3467. """Asynchronously create a new relation between entities.
  3468. Creates a new relation (edge) in the knowledge graph and adds it to the vector database.
  3469. Args:
  3470. source_entity: Name of the source entity
  3471. target_entity: Name of the target entity
  3472. relation_data: Dictionary containing relation attributes, e.g. {"description": "description", "keywords": "keywords"}
  3473. Returns:
  3474. Dictionary containing created relation information
  3475. """
  3476. from lightrag.utils_graph import acreate_relation
  3477. return await acreate_relation(
  3478. self.chunk_entity_relation_graph,
  3479. self.entities_vdb,
  3480. self.relationships_vdb,
  3481. source_entity,
  3482. target_entity,
  3483. relation_data,
  3484. )
  3485. def create_relation(
  3486. self, source_entity: str, target_entity: str, relation_data: dict[str, Any]
  3487. ) -> dict[str, Any]:
  3488. loop = always_get_an_event_loop()
  3489. return loop.run_until_complete(
  3490. self.acreate_relation(source_entity, target_entity, relation_data)
  3491. )
  3492. async def amerge_entities(
  3493. self,
  3494. source_entities: list[str],
  3495. target_entity: str,
  3496. merge_strategy: dict[str, str] = None,
  3497. target_entity_data: dict[str, Any] = None,
  3498. ) -> dict[str, Any]:
  3499. """Asynchronously merge multiple entities into one entity.
  3500. Merges multiple source entities into a target entity, handling all relationships,
  3501. and updating both the knowledge graph and vector database.
  3502. Args:
  3503. source_entities: List of source entity names to merge
  3504. target_entity: Name of the target entity after merging
  3505. merge_strategy: Merge strategy configuration, e.g. {"description": "concatenate", "entity_type": "keep_first"}
  3506. Supported strategies:
  3507. - "concatenate": Concatenate all values (for text fields)
  3508. - "keep_first": Keep the first non-empty value
  3509. - "keep_last": Keep the last non-empty value
  3510. - "join_unique": Join all unique values (for fields separated by delimiter)
  3511. target_entity_data: Dictionary of specific values to set for the target entity,
  3512. overriding any merged values, e.g. {"description": "custom description", "entity_type": "PERSON"}
  3513. Returns:
  3514. Dictionary containing the merged entity information
  3515. """
  3516. from lightrag.utils_graph import amerge_entities
  3517. return await amerge_entities(
  3518. self.chunk_entity_relation_graph,
  3519. self.entities_vdb,
  3520. self.relationships_vdb,
  3521. source_entities,
  3522. target_entity,
  3523. merge_strategy,
  3524. target_entity_data,
  3525. self.entity_chunks,
  3526. self.relation_chunks,
  3527. )
  3528. def merge_entities(
  3529. self,
  3530. source_entities: list[str],
  3531. target_entity: str,
  3532. merge_strategy: dict[str, str] = None,
  3533. target_entity_data: dict[str, Any] = None,
  3534. ) -> dict[str, Any]:
  3535. loop = always_get_an_event_loop()
  3536. return loop.run_until_complete(
  3537. self.amerge_entities(
  3538. source_entities, target_entity, merge_strategy, target_entity_data
  3539. )
  3540. )
  3541. async def aexport_data(
  3542. self,
  3543. output_path: str,
  3544. file_format: Literal["csv", "excel", "md", "txt"] = "csv",
  3545. include_vector_data: bool = False,
  3546. ) -> None:
  3547. """
  3548. Asynchronously exports all entities, relations, and relationships to various formats.
  3549. Args:
  3550. output_path: The path to the output file (including extension).
  3551. file_format: Output format - "csv", "excel", "md", "txt".
  3552. - csv: Comma-separated values file
  3553. - excel: Microsoft Excel file with multiple sheets
  3554. - md: Markdown tables
  3555. - txt: Plain text formatted output
  3556. - table: Print formatted tables to console
  3557. include_vector_data: Whether to include data from the vector database.
  3558. """
  3559. from lightrag.utils import aexport_data as utils_aexport_data
  3560. await utils_aexport_data(
  3561. self.chunk_entity_relation_graph,
  3562. self.entities_vdb,
  3563. self.relationships_vdb,
  3564. output_path,
  3565. file_format,
  3566. include_vector_data,
  3567. )
  3568. def export_data(
  3569. self,
  3570. output_path: str,
  3571. file_format: Literal["csv", "excel", "md", "txt"] = "csv",
  3572. include_vector_data: bool = False,
  3573. ) -> None:
  3574. """
  3575. Synchronously exports all entities, relations, and relationships to various formats.
  3576. Args:
  3577. output_path: The path to the output file (including extension).
  3578. file_format: Output format - "csv", "excel", "md", "txt".
  3579. - csv: Comma-separated values file
  3580. - excel: Microsoft Excel file with multiple sheets
  3581. - md: Markdown tables
  3582. - txt: Plain text formatted output
  3583. - table: Print formatted tables to console
  3584. include_vector_data: Whether to include data from the vector database.
  3585. """
  3586. try:
  3587. loop = asyncio.get_event_loop()
  3588. except RuntimeError:
  3589. loop = asyncio.new_event_loop()
  3590. asyncio.set_event_loop(loop)
  3591. loop.run_until_complete(
  3592. self.aexport_data(output_path, file_format, include_vector_data)
  3593. )
  3594. # `addon_params` is declared as an InitVar on the dataclass so it can still be
  3595. # passed through LightRAG(addon_params=...). InitVars are not stored as
  3596. # instance attributes, which frees the name to be installed here as a property
  3597. # that routes reads/writes through the observable `_addon_params` store.
  3598. # Declaring it as both a dataclass field and a property is not supported by
  3599. # @dataclass, so the property is attached after class creation.
  3600. LightRAG.addon_params = property( # type: ignore[attr-defined]
  3601. LightRAG._get_addon_params,
  3602. LightRAG._set_runtime_addon_params,
  3603. )