JivxddlZddlZddlZddlZddlmZmZddlmZddl m Z m Z m Z m Z mZmZmZddlZddlZddlmZddlmZmZmZmZddlmZmZmZmZm Z m!Z!m"Z"m#Z#m$Z$m%Z%ddl&m'Z'dd l(m)Z)dd l*m+Z+m,Z,dd l-m.Z.m/Z/m0Z0dd l1m2Z2dd l3m4Z4ddl5m6Z6ddl7m8Z8m9Z9ddl:m;Z;ddle e+e efZ?ee$je"jZBdeCdeDfdZEde'de?de+fdZF dFde'deGdeHdeIdeIdeIdefdZJ dGde'deGd e eGd!e eGdeHdeIdeIdeIdefd"ZK dHde'deGd#eHd$eHd%eDd e eGd!e eGd&eHd'eId(eId)eHd*eId+eIde efd,ZLd-eeee>fd.e=d/eDd0eId1ede>f d2ZMd3eddfd4ZNGd5d6ZOGd7d8e4ZPeQd9k(rdd:lRmSZSeSeeTjjd;z dZYeYjeVdDz dEz Z]eWjeYje`e]yy#eZ$reWjd=d?d@dAid@dAidBdCZYY_wxYw)IN)BufferedIOBase TextIOWrapper)Path)ListOptionalTypeUnion CoroutineAnyTypeVar) BaseModel)retry_if_exceptionstop_after_attemptwait_exponential_jitter AsyncRetrying) ExtractAgent ExtractConfig ExtractJob ExtractRunFileFileData ExtractMode StatusEnum ExtractTargetPaginatedExtractRunsResponse)AsyncLlamaCloud)ApiError)JSONObjectTypeExperimentalWarning)augment_async_errors SourceText FileInput) FileClient) BaseComponent)run_jobs)Field PrivateAttr)DEFAULT_BASE_URL)ThreadPoolExecutorT)extraction_targetextraction_mode exceptionreturnct|tr|jdvSt|tjtj tj fryy)z#Check if an exception is retryable.)iiiiiTF) isinstancer status_codehttpxHTTPStatusError RequestErrorTimeoutException)r-s v/var/www/html/BankruptcyAI-uat/bankruptcy-ai/venv/lib/python3.12/site-packages/llama_cloud_services/extract/extract.py_is_retryable_errorr76sL)X&$$(AAA E))5+=+=u?U?UV  client data_schemac Kt|tr|}n : Bcheck_interval max_timeoutverbosejob_retry_attempts job_max_wait job_jitterrun_retry_attempts run_max_wait run_jitterc Ktj} d} tj|d{|dz }t ||||| d{}|j t jk(rt||||| | | d{S|j t jk(rCtj}|| z |kDrtd||r|dzdk(rtdd d tjd |d |j d |jt||||| | | d{S7777 w)5Wait for and return the results of an extraction job.rTN)rHrJrKz#Timeout while extracting the file: .)endflushzFailure in job: z , status: z , error: )time perf_counterasynciosleeprXstatusrSUCCESSr]PENDING Exceptionprintwarningswarnerror)r9rGr^r_r`rYrZrarbrcrdrerfstart poll_countjobrms r6_wait_for_job_resultr~sZ    EJ mmN+++a '  +!    ::++ +,/%! ZZ:-- -##%CU{[("EfX NOO:?a/cr.  MM"6(*SZZL #))U -/%! ? , (sF0ED7ED:6ED<B+E2D>3E:E<E>Ecoro thread_poolverify httpx_timeoutclient_wrappercdtffd dtffd }|j|jS)z8Run coroutine in a thread with proper client management.r.c2Ktjtjdd}j} |_d{|_|j d{S7#7#|_|j d{7wxYww)Nd)max_keepalive_connectionsmax_connections)rtimeoutlimits)r2 AsyncClientLimits httpx_clientaclose)r9original_clientrrrrs r6 wrapped_coroz#run_in_thread..wrapped_coros""!<<#sS  )55 "*0N ':*9N '--/ ! ! "+:N '--/ ! !sL:B A2 A. A2B(A0)B.A20B2B BBBc tjS#tj$r}t dt ||d}~wtj $r}tdt ||d}~wwxYw)NzRequest timed out: zNetwork error: )rqrunr2r5 TimeoutErrorstr NetworkErrorConnectionError)ers r6run_corozrun_in_thread..run_coross E;;|~. .%% F!4SVH=>A E!! E!OCF8"<=1 D Es A;AA;A66A;)r*submitresult)rrrrrrrs` ``` @r6 run_in_threadrs> " " "EaE   h ' . . 00r8configcX|js |jrtjdt|j r(|j tjk(r td|jr8|j tjtjfvr tdyy)Nz`cite_sources`/`confidence_scores` could greatly increase the size of the response, and slow down the extraction. Results will be available in the `extraction_metadata` field for the extraction run.zU`reasoning` is only supported with BALANCED, MULTIMODAL, or PREMIUM extraction modes.zM`cite_sources` is only supported with MULTIMODAL or PREMIUM extraction modes.) cite_sourcesconfidence_scoresrxryr use_reasoningr,rFASTrABALANCED)rs r6_extraction_config_warningrs f66  S    ! ![%5%5 5g   ! !k&6&6 8L8L%M M_  Nr8cjeZdZdZ d-dededeedeededed ed e d e d ee d ee fdZ e defdZ e defdZe defdZej"deddfdZe defdZej"deddfdZdeeeefdefdZdedefdZdedefdZdedeefdZd.dZ d e!ee"efde!e#e"e#ffd!Z$d e!ee"efde!ee"effd"Z%d e!ee"efde!ee"effd#Z&dede#fd$Z'dedefd%Z(d&eddfd'Z) d/d(ed)ede*fd*Z+defd+Z,d.d,Z-y)0ExtractionAgentzTClass representing a single extraction agent with methods for extraction operations.Nr9agentrYrZr^r_ num_workers show_progressr`rrc F||_||_||_||_||_||_||_||_| |_| |_ | |_ d|_ d|_ ttdtj xsddz|_t%||||_y)Nrjri max_workers)_client_agent _project_id_organization_idr^r_rrrr_verbose _data_schema_configr)minos cpu_count _thread_poolr# _file_client) selfr9rrYrZr^r_rrr`rrs r6__init__zExtractionAgent.__init__s  % /,&&* * 9=37 .B!41 9: 'vz?Kr8r.c.|jjSN)ridrs r6rzExtractionAgent.ids{{~~r8c.|jjSr)rnamers r6rzExtractionAgent.names{{r8c^|js|jjS|jSr)rrr:rs r6r:zExtractionAgent.data_schema s%.2.?.?t{{&&VTEVEVVr8r:cX|jt|j||_yr)_run_in_threadrFrr)rr:s r6r:zExtractionAgent.data_schema$s%!// T\\; 7 r8c^|js|jjS|jSr)rrrrs r6rzExtractionAgent.config+s!)-t{{!!G4<<Gr8rc(t|||_yr)rr)rrs r6rzExtractionAgent.config/s"6* r8rct||j|j|j|jj Sz=Run coroutine in a separate thread to avoid event loop issues)rrrrr_client_wrapperrrs r6rzExtractionAgent._run_in_thread4s8     KK    LL ( (   r8 file_inputcTK|jj|d{S7w)a"Upload a file for extraction. Args: file_input: The file to upload (path, bytes, or file-like object) Raises: ValueError: If filename is not provided for bytes input or for file-like objects without a name attribute. Nrupload_contentrrs r6 upload_filezExtractionAgent.upload_file>s%&&55jAAAA (&(cTK|jj|d{S7w)z8Upload a file from various input types using FileClient.Nrrs r6 _upload_filezExtractionAgent._upload_fileJs#&&55jAAAArrGcKt|j||j|j|j|j |j dddddd d{S7w)rh< r9rGr^r_r`rYrZrarbrcrdrerfN)r~rr^r_rrrrrGs r6r~z$ExtractionAgent._wait_for_job_resultNsd)<<..((MM'' 11      AA!AA!c|j|jjj|j|j |j |_y)zPersist the extraction agent's schema and config to the database. Returns: ExtractionAgent: The updated extraction agent )extraction_agent_idr:rN)rrrBupdate_extraction_agentrr:rrrs r6savezExtractionAgent.save`sJ )) LL & & > >$(GG ,,{{ ?   r8filesc JK t|ts|g}d}nd}|Dcgc]}|j|}}t5t ||j d|j d{}dddDcgc]S}|jjj|j|j|j|jU}}t5t ||j d|j d{}ddd|jrRt|D]C\}}t|tt fr t|nd} t#d | d |jE|rd SScc}w7#1swYxYwcc}w7#1swYxYww) z Queue multiple files for extraction. Args: files (Union[FileInput, List[FileInput]]): The files to extract Returns: Union[ExtractJob, List[ExtractJob]]: The queued extraction jobs TFzUploading filesworkersdescrN)rfile_iddata_schema_overrideconfig_overridezCreating extraction jobszz Queued file extraction for file z under job_id r)r0listrr r%rrrrBrun_jobrr:rrziprrrw) rr single_filefile upload_tasksuploaded_files job_tasks extract_jobsr} file_reprs r6queue_extractionz ExtractionAgent.queue_extractionns C%&GEKKCIDT6ykPSPVPVxX  #.|A?<?GC*     sqF#E; F#'F+F,F0 F# 99d9o56 9 z4 ++ , 9 V V V    C D +.  $' % B#Br8rceZdZUdZedZeed<edZeed<edd Z e ed <ed d Z e ed <eddddZ e ed<edd Z eed<edd Zeed<edd Zeeed<edd Zeeed<eZeed<eZeed <eZeeed!<eZeeed"< d?deedeed e d e de ded$eed%eedeedeedeffd& Zd'eeeefd(efd)Z d@d*ed+e d,ee!d(e"fd-Z# dAd*eed.eed(e"fd/Z$d(e%e"fd0Z&d1ed(d#fd2Z'd3ed(ee(fd4Z) dAd5eed6ee*ee+fd(efd7Z,d8e-d(e*e.ee/ffd9Z0d+e d,e!d:e*e-e%e-fd(e*e1e%e1ffd;Z2d+e d,e!d:e*e-e%e-fd(e*e(e%e(ffd<Z3d+e d,e!d:e*e-e%e-fd(e*e(e%e(ffd=Z4dBd>Z5xZ6S)C LlamaExtractz:Factory class for creating and managing extraction agents.z%The API key for the LlamaExtract API.) descriptionapi_keyz%The base URL of the LlamaExtract API.base_urlriz;The interval in seconds to check if the extraction is done.)defaultr"r^rzDThe maximum timeout in seconds to wait for the extraction to finish.r_rrrjzAThe number of workers to use sending API requests for extraction.)r%gtltr"rTz-Show progress when extracting multiple files.rFz*Show verbose output when extracting files.r`zSimple SSL verification option.rrzTimeout for the httpx client.r _async_clientrrrNrYrZc .|s#tjdd}| td|stjddxst}t |||||||||| | |  t j| | |_| |_ | |_ t|j|j|j|_tt!dtj"xsdd z |_|stjd d}||_||_y) NLLAMA_CLOUD_API_KEYzThe API key is required.LLAMA_CLOUD_BASE_URL) r#r$r^r_rrrYrZrrr`)rr)tokenr$rrjrirrLLAMA_CLOUD_PROJECT_ID)rgetenvrAr(superrr2r _httpx_clientrrrr#r$r(r)rrrrr) rr#r$r^r_rrrYrZrrr` __class__s r6rzLlamaExtract.__init__9sii 5tRBRH )##'!+'  #..fmT *,,,]]++  /B!41 9: #;TBJ% /r8rr.ct||j|j|j|jj Sr)rrrrr(rrs r6rzLlamaExtract._run_in_threadms:     KK       . .   r8rr:rc 2| t|nt}t|trn,t |t r|j }n td|j|jjj|j|j|||}t|j||j|j|j|j |j"|j$|j&|j(|j* S)a]Create a new extraction agent. Args: name (str): The name of the extraction agent data_schema (SchemaInput): The data schema for the extraction agent config (Optional[ExtractConfig]): The extraction config for the agent Returns: ExtractionAgent: The created extraction agent r<)rYrZrr:r r9rrYrZr^r_rrr`rr)rDEFAULT_EXTRACT_CONFIGr0r=r?r r@rArr(rBcreate_extraction_agentrrrr^r_rrr`rr)rrr:rrs r6 create_agentzLlamaExtract.create_agentws   &v .+F k4 (   Y /%779KM ##    , , D D++ $ 5 5' E  %%'' 11..((((,,LL;;,,  r8rc 4||tjd|r6|j|jjj |}nN|rA|j|jjj ||j}n tdt|j||j|j|j|j|j|j|j|j |j" S)aGet extraction agents by name or extraction agent ID. Args: name (Optional[str]): Filter by name extraction_agent_id (Optional[str]): Filter by extraction agent ID Returns: ExtractionAgent: The extraction agent zJBoth name and extraction_agent_id are provided. Using extraction_agent_id.r)rrYz4Either name or extraction_agent_id must be provided.r4)rxryrr(rBget_extraction_agentget_extraction_agent_by_namerrArrr^r_rrr`rr)rrrrs r6 get_agentzLlamaExtract.get_agents  >d. MM\  ''""00EE(*FE ''""00MM#//NEST T%%'' 11..((((,,LL;;,,  r8c|j|jjj|j}|Dcgc]|}t |j||j|j |j|j|j|j|j|j|j ~c}Scc}w)z%List all available extraction agents.)rYr4)rr(rBlist_extraction_agentsrrrr^r_rrr`rr)ragentsrs r6 list_agentszLlamaExtract.list_agentss$$    , , C C++ D  (   ))++ $ 5 5#22 ,, ,,"00 {{"00    sBCagent_idcn|j|jjj|y)zzDelete an extraction agent by ID. Args: agent_id (str): ID of the extraction agent to delete r9N)rr(rBdelete_extraction_agent)rrAs r6 delete_agentzLlamaExtract.delete_agents5     , , D D$, E  r8rGcKt|j||j|j|j|j |j dddddd d{S7w)rhrrrrN)r~r(r^r_r`rrrs r6r~z!LlamaExtract._wait_for_job_resultsf)%%..((LL'' 11      rfilename file_pathc dddddddddd d d }d }|r$t|jj}n%|r#t|jj}|r ||vr||S|jDcgc]}|d d  }}d j t |}|r|d d }t d|d|t d|cc}w)z9Determine MIME type for a file based on filename or path.z text/plainztext/csvzapplication/jsonz text/htmlz text/markdownzapplication/pdfzGapplication/vnd.openxmlformats-officedocument.wordprocessingml.documentz image/pngz image/jpeg) z.txtz.csvz.jsonz.htmlz.htmz.mdz.pdfz.docxz.pngz.jpgz.jpegNriz, zUnsupported file type: 'z'. Supported formats are: zaCould not determine file type. Please provide a filename with one of these supported extensions: )rsuffixlowerkeysjoinsortedrA) rrFrG MIME_TYPE_MAP extensionextsupported_extensionssupported_listext_without_dots r6_get_mime_typezLlamaExtract._get_mime_types!' "%^ ! $ X--335I Y..446I m3 + +4A3E3E3GHCABHH6*>#?@ 'mO*?*;<**8)9;  suCtDE  Is9 Crct|tr|St|tr|j |jS|jt|jt r|j}|j }n:t|jttfrSt|jd5}|j}ddd|j xst|j}nt|jttfrt|jdr?|jj}t|tr|jd}n|}n td|j xst!|jdd}n!tdt#|jt%j&j)d}|j+|}t-|| Std t|ttfrht|d5}|j}dddt%j&j)d}|j+| }t-|| St|t r td t|ttfrt|dr|j}t|tr|jd}n|}t%j&|j)d}t!|dd}|j+|}t-|| Stdtd t#|#1swYxYw#1swY=xYw)zFConvert FileInput to FileData or text string for stateless extraction.Nrbreadzutf-8z#File object must have a read methodrzUnsupported file type: )rF)data mime_typez0SourceText must have either text_content or file)rGziCannot determine file type from raw bytes. Please use SourceText with a filename, or provide a file path.zUnsupported file input type: )r0rr! text_contentrbytesrFrropenrWrrr encoderAgetattrr>base64 b64encodedecoderTr)rrrXrFfcontent encoded_datarYs r6_convert_file_to_file_dataz'LlamaExtract._convert_file_to_file_dataHs j$ '  j* -&&2!...,joou5%??D)22H #t=joot4( vvx()22Jc*//6JH .-1PQz7",//"6"6"8%gs3#*>>'#:D#*D()NOO)22 g"7H%'>tJOO?T>U%VWW &//5<>'2D"D%//5<B)DD-DDcK|j|||d{}t|trYg}|D]P}|j|jd{}|t d|j|j |R|S|j|jd{}|t d|j|S77t7$w)a@Run stateless extraction and wait for results. Args: data_schema: The schema defining what data to extract config: The extraction configuration files: File(s) to extract from Returns: ExtractRun or list of ExtractRuns with the extraction results Nz(Failed to get extraction result for job )rr0rr~r RuntimeErrorri)rr:rrrrunsr}rs r6rzLlamaExtract.aextracts **;FF dD !D ! 55cff==;&B366(K C  !K11$''::C{"%MdggY#WXXJ!G >;s4C C9C CAC %C &!C C  C cF|j|j|||S)aVRun stateless extraction and wait for results (synchronous version). Args: data_schema: The schema defining what data to extract config: The extraction configuration files: File(s) to extract from Returns: ExtractRun or list of ExtractRuns with the extraction results r)rr:rrs r6rzLlamaExtract.extracts" ""4==fe#LMMr8cv t|dr|jjdyy#t$rYywxYwr r rs r6rzLlamaExtract.__del__rr) NNrirrTNNTrFr)NNr)7rrrrr&r#r__annotations__r$r^rr_rrrr`rrrrr'r(rrr)rrrr r r*rrrrr7r<rr@rDrr~r rrTr"rrrerrrrr __classcell__)r1s@r6r!r!sD%LMGSM&MNHcNQNCZK W K  "QM4#OGT#"CFHTN&+ ?&M8E?&1]M?2'2}L$4!,K#.&1mhsm3"&"&"$()-!%)+20#203-20 20  20  2020SM20"#2020 2020h 9S#q[#9 a +/ 4 4 !4 ' 4  4 p# 0 sm0  SM0   0 d T/2 2  S  T    *9M (#'0403-0E#t),-0 0dJQ#JQ xd" #JQX-3 -3-3YY/0 -3 z4 ++ , -3^    YY/0  z4 ++ , DN NNYY/0 N z4 ++ , N$r8r!__main__) load_dotenvtestsrXz test-agent)robjectr>string)titlesummary)r> propertiesslidezconocophilips.pdf)rrirr)NNrrirr) rirFNNrrrrrr)arqr_rroiorrpathlibrtypingrrrr r r r rxr2pydanticr tenacityrrrr llama_cloudrrrrrrrrrrrllama_cloud.clientrllama_cloud.core.api_errorr"llama_cloud_services.extract.utilsrrllama_cloud_services.utilsr r!r"!llama_cloud_services.files.clientr#llama_index.core.schemar$llama_index.core.async_utilsr% llama_index.core.bridge.pydanticr&r'llama_index.core.constantsr(concurrent.futuresr)r*rPER_DOC MULTIMODALr5 BaseExceptionrr7rFrrrrXr]r~rrrr!rdotenvru__file__parentdata_dir extractorr<rrvr7rrrDrrwrr8r6rs ,GGG    //SR811?71 CLNDO34 &#++** =T( (*5((. E E EE E  E  EE,!%%)   c]      8 $%); ; ;; ;  ;  ;c];;;;;;;j;|1 CaK 1#1 1 1  1  1D}(]]@ h=hV z"M H~$$++g5>HI  ###6mmHw.1DDEG 588$ 'NI.   &&  $h/ &1     sH!H98H9