NgtdZddlmZddlZddlZddlZddlmZmZm Z m Z m Z m Z m Z ddlZddlmZddlmZddlmZddlmZdd lmZed gZd Zd Zd Zejej d j!Z"ej#ej dj!Z$ej%j!Z&ddZ'ddZ(ddZ)ddZ*ddZ+GddeZ,dS)z&Wrapper around TileDB vector database.) annotationsN)AnyDictIterableListMappingOptionalTuple)Document) Embeddings guard_import) VectorStore)maximal_marginal_relevance euclidean documentsvectorsuint64float32returnrc>tdtdfS)z@Import tiledb-vector-search if available, otherwise raise error.tiledb.vector_searchtiledbr c/var/www/html/ai-engine/env/lib/python3.11/site-packages/langchain_community/vectorstores/tiledb.pydependable_tiledb_importrs% +,,X rgroupstrc&|tjS) Get the URI of the vector index.)VECTOR_INDEX_NAMEurirs rget_vector_index_uri_from_groupr%#s " # ''rc&|tjS)zGet the URI of the documents array from group. Args: group: TileDB group object. Returns: URI of the documents array. )DOCUMENTS_ARRAY_NAMEr#r$s r"get_documents_array_uri_from_groupr((s % & **rr#c|dtS)r!/)r"r#s rget_vector_index_urir,4s ' '% ' ''rc|dtS)z#Get the URI of the documents array.r*)r'r+s rget_documents_array_urir.9s * *( * **rceZdZdZdddddddUdZedVdZddeddWd%Zddd&d'dXd)Z ddd&d'dYd+Z dZd[d-Z dZd\d.Z dd&d/dd0d]d2Z d^d_d3Z d^d`d4Zed5dd6dad=Zedded>dd?d@dbdHZ dcdddJZ dedfdLZeddedMd>dd?fdgdNZedded>dd?d@dhdQZeedddRdidSZdjdTZdS)kTileDBa2TileDB vector store. To use, you should have the ``tiledb-vector-search`` python package installed. Example: .. code-block:: python from langchain_community import TileDB embeddings = OpenAIEmbeddings() db = TileDB(embeddings, index_uri, metric) NF)vector_index_uridocs_array_uriconfig timestampallow_dangerous_deserialization embeddingr index_urirmetricr2r3r4Optional[Mapping[str, Any]]r5rr6boolkwargsc `|std||_|j|_||_||_||_tdtd} } | |5| |jd} |dkr|nt| |_ |dkr|nt| |_ | | |j d} | jd|_| ||_|jdkr+| jjd |j |j|jd | |_n5|jd kr*| jjd |j |j|jd | |_d d d d S#1swxYwYd S) aInitialize with necessary components. Args: allow_dangerous_deserialization: whether to allow deserialization of the data which involves loading data using pickle. data can be modified by malicious actors to deliver a malicious payload that results in execution of arbitrary code on your machine. aTileDB relies on pickle for serialization and deserialization. This can be dangerous if the data is intercepted and/or modified by malicious actors prior to being de-serialized. If you are sure that the data is safe from modification, you can set allow_dangerous_deserialization=True to proceed. Loading of compromised data using pickle can result in execution of arbitrary code on your machine.rr ctx_or_configrr1 index_typeFLAT)r#r4r5IVF_FLATNr) ValueErrorr7 embed_queryembedding_functionr8r9r4r scope_ctxGroupr%r2r(r3closemetagetrAr5 flat_index FlatIndex vector_indexivf_flat_index IVFFlatIndex)selfr7r8r9r2r3r4r5r6r< tiledb_vsr index_grouprs r__init__zTileDB.__init__Ls;,/ 2 #"+"7"   / 0 0  " "   F  3 3   ,,t~s;;K$r))! 4[AA  !"R''7 DD        LL!6<z0TileDB.process_index_results..sC"U 5$(?(?J%Urc3VK|]#\}}j||vV$dSrW)rhrK)rmrnro result_docs r z/TileDB.process_index_results..sQ"U'++C00E9rc*g|]\}}|k ||fSrr)rmdocscorer\s r z0TileDB.process_index_results..s+PPPeu7O7Oe 7O7O7Or)ropenr3r5r4zip MAX_UINT64 MAX_FLOAT_32lenrDrKr rpickleloadsnparraytolistastypeuint8tobytesrhitemsallappendrI)rQr]r_rZr[r\rdocs docs_arrayidxrvrupickled_metadatarhrrs ` @rprocess_index_resultszTileDB.process_index_resultss1*h''[[  t{!  c6** 1 1JCaxxEQJJj  Ul%:%:S/C{c#f+..!33 !S3!S!Sc!S!STTT"wwz22 !s3v;q>/B/BCCCJ+!<H-446677>>rxHHPPRR'/ #!&,llnn&,llnn5KKU 3444 Z/0000PPPPtPPPBQBxrrZr[fetch_krc d|vr|d}nt}|jjt jt j|tjgtjfd||n|i|\}}||d|d|||S)a[Return docs most similar to query. Args: embedding: Embedding vector to look up documents similar to. k: Number of Documents to return. Defaults to 4. filter (Optional[Dict[str, Any]]): Filter by metadata. Defaults to None. fetch_k: (Optional[int]) Number of Documents to fetch before filtering. Defaults to 20. **kwargs: kwargs to be passed to similarity search. Can include: nprobe: Optional, number of partitions to check if using IVF_FLAT index score_threshold: Optional, a floating point value to filter the resulting set of retrieved docs Returns: List of documents most similar to the query text and distance in float for each. Lower score represents more similarity. r\rZNrr]r_r[rZr\) pop MAX_FLOATrNqueryrrrrr) rQr7rZr[rr<r\dis r&similarity_search_with_score_by_vectorz-TileDB.similarity_search_with_score_by_vectors4  & &$jj):;;OO'O&t & Hbhy))00<<= > > E Ebj Q Q  >aaw   1 ))!QqT&A*   rrc T||}|j|f|||d|}|S)aReturn docs most similar to query. Args: query: Text to look up documents similar to. k: Number of Documents to return. Defaults to 4. filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None. fetch_k: (Optional[int]) Number of Documents to fetch before filtering. Defaults to 20. Returns: List of documents most similar to the query text with Distance as float. Lower score represents more similarity. r)rFr)rQrrZr[rr<r7rs rsimilarity_search_with_scorez#TileDB.similarity_search_with_scoresR,++E22 :t:         rList[Document]c >|j|f|||d|}d|DS)aReturn docs most similar to embedding vector. Args: embedding: Embedding to look up documents similar to. k: Number of Documents to return. Defaults to 4. filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None. fetch_k: (Optional[int]) Number of Documents to fetch before filtering. Defaults to 20. Returns: List of Documents most similar to the embedding. rcg|]\}}|Srrrmru_s rrwz6TileDB.similarity_search_by_vector..5222Q222r)r)rQr7rZr[rr<docs_and_scoress rsimilarity_search_by_vectorz"TileDB.similarity_search_by_vectorsO(F$E        32/2222rc >|j|f|||d|}d|DS)aReturn docs most similar to query. Args: query: Text to look up documents similar to. k: Number of Documents to return. Defaults to 4. filter: (Optional[Dict[str, str]]): Filter by metadata. Defaults to None. fetch_k: (Optional[int]) Number of Documents to fetch before filtering. Defaults to 20. Returns: List of Documents most similar to the query. rcg|]\}}|Srrrs rrwz,TileDB.similarity_search..Nrr)r)rQrrZr[rr<rs rsimilarity_searchzTileDB.similarity_search7sJ(<$;  vw  :@  32/2222r?rZr lambda_multr[rc dd|vr|d}nt}jjt jt j|tjgtjfd||n|dzi|\}} | d|d|||n|dz|} fd| D} tt j|gtj| || } g} | D]}| | || S) azReturn docs and their similarity scores selected using the maximal marginal relevance. Maximal marginal relevance optimizes for similarity to query AND diversity among selected documents. Args: embedding: Embedding to look up documents similar to. k: Number of Documents to return. Defaults to 4. fetch_k: Number of Documents to fetch before filtering to pass to MMR algorithm. lambda_mult: Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5. Returns: List of Documents and similarity scores selected by maximal marginal relevance and score for each. r\rZNrrcbg|]+\}}j|jgd,Sr)r7embed_documentsri)rmrurrQs rrwzMTileDB.max_marginal_relevance_search_with_score_by_vector..}sE   FLc1DN * *C,<+= > >q A   rdtype)rZr) rrrNrrrrrrrr)rQr7rZrrr[r<r\r_indicesresultsrX mmr_selectedrrs` r2max_marginal_relevance_search_with_score_by_vectorz9TileDB.max_marginal_relevance_search_with_score_by_vectorPst:  & &$jj):;;OO'O1$+1 Hbhy))00<<= > > E Ebj Q Q  ggWq[    ,, !9ggWq[+ -      PW   2 Hi[ 3 3 3 #     / /A  " "71: . . . .rc @|j|f||||d|}d|DS)aReturn docs selected using the maximal marginal relevance. Maximal marginal relevance optimizes for similarity to query AND diversity among selected documents. Args: embedding: Embedding to look up documents similar to. k: Number of Documents to return. Defaults to 4. fetch_k: Number of Documents to fetch before filtering to pass to MMR algorithm. lambda_mult: Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5. Returns: List of Documents selected by maximal marginal relevance. rcg|]\}}|Srrrs rrwzBTileDB.max_marginal_relevance_search_by_vector..rr)r)rQr7rZrrr[r<rs r'max_marginal_relevance_search_by_vectorz.TileDB.max_marginal_relevance_search_by_vectorsR4R$Q  #      32/2222rc V||}|j|f||||d|}|S)aReturn docs selected using the maximal marginal relevance. Maximal marginal relevance optimizes for similarity to query AND diversity among selected documents. Args: query: Text to look up documents similar to. k: Number of Documents to return. Defaults to 4. fetch_k: Number of Documents to fetch before filtering (if needed) to pass to MMR algorithm. lambda_mult: Number between 0 and 1 that determines the degree of diversity among the results with 0 corresponding to maximum diversity and 1 to minimum diversity. Defaults to 0.5. Returns: List of Documents selected by maximal marginal relevance. r)rFr) rQrrZrrr[r<r7rs rmax_marginal_relevance_searchz$TileDB.max_marginal_relevance_searchsU4++E22 ;t;  #       rT) metadatasr4rA dimensions vector_typenp.dtyperNonecvtdtd}}||5 ||n#|j$r} | d} ~ wwxYw||d} t | j} t| j} |dkr|j | |||n$|dkr|j | |||| | t| d d td z ftjtj } || }|d tjdd}|g}|r7|dtjd}||||dd|}|j | || | t0| ddddS#1swxYwYdS)Nrrr>wrB)r#rrr4rC)nameidr)rdomainrrgU1T)rrvarrhF)rsparseallows_duplicatesattrs)rrG group_create TileDBErrorrHr,r#r.rLcreaterOaddr"DimrzrrrDomainAttrrr ArraySchemaArrayr'rI)clsr8rArrrr4rRrerrrr2docs_uridimdom text_attrr metadata_attrschemas rrz TileDB.creates / 0 0  " "   F  3 3/ /  ##I....%     LLC00E3EI>> .uy99HV##$++() +! , z))(//() +! 0 II&->I ? ? ? **:>*hry))C --$$C rx~~4 PPIKE , & 28QU V V  ]+++''"' (F L  & 1 1 1 IIh%9I : : : KKMMM_/ / / / / / / / / / / / / / / / / / s5H.A  H. AAAGH..H25H2rBr)rr]r9rAr4index_timestamptexts List[str]rXList[List[float]]Optional[List[dict]]Optional[List[str]]rc |tvr'td|dtttdtd} } t j|tj}||||j d|j |du| | | 5|stdt|}t|}| d |D}t j|tj}| jjd||||| d kr| nd| d | | |d 5}|Qt jt'|tj }t)t'|D]}|||<i}t j||d<|ot jt'|gt, }d }|D]<}t jt1j|tj ||<|dz }=||d<|||<dddn #1swxYwYdddn #1swxYwY|d|||| d| S)NzUnsupported distance metric: z. Expected one of rrr)r8rArrrr4r>z3embeddings must be provided to build a TileDB indexc bg|],}ttjdtdz -Srrrrandomrandintrzrmrs rrwz!TileDB.__from..?s0MMM!s6>!Z!^<<==MMMrr)rAr8 input_vectors external_idsrr4rrrgrh)r7r8r9r4r) INDEX_METRICSrDrlrrrrrrshaperrGr,r.r ingestioningestrxzerosr|rangeemptyobject frombufferr}dumpsr)rrrXr7r8rr]r9rAr4rr<rRrrr2rrArdatarrhs r__fromz TileDB.__fromsk  & &=F=='+M':':==  / 0 0  " " ,,33BJ??  !$*1-%+t+      F  3 3$ '$ ' X !VWWW3I>> .y99H{MMuMMM8C==// ::L &I  & %*+)3Ba3G3GT      Xs++ 'q'#%8CJJbi#H#H#HL"3u::..,,*+ Q!xV ($&Hc)nn-=V$L$L$LMA$-+-="L22"(,,, a(Q'4D$"&,# ' ' ' ' ' ' ' ' ' ' ' ' ' ' ''$ '$ '$ '$ '$ '$ '$ '$ '$ '$ '$ '$ '$ '$ '$ 'Js        s8B&I/)C#I I/I I/I I//I36I3Optional[bool]c tj|tj}|j||dkr|nddS)amDelete by vector ID or other criteria. Args: ids: List of ids to delete. timestamp: Optional timestamp to delete with. **kwargs: Other keyword arguments that subclasses might use. Returns: Optional[bool]: True if deletion is successful, False otherwise, None if not implemented. rN)rr5T)rrrrrN delete_batch)rQr]r5r<rs rdeletez TileDB.deleteesZx}}++BI66  &&%i1nnRV '   tr Iterable[str]c td}|jt|}| d|D}t j|tj}t jt|d} tt|D]+} t j|| tj | | <,|j | ||dkr|ndi} t j|| d<|ot jt|gt} d} |D]<} t jt!j| tj| | <| d z } =| | d <||jd |dkr|nd|j }| ||<||S) aRun more texts through the embeddings and add to the vectorstore. Args: texts: Iterable of strings to add to the vectorstore. metadatas: Optional list of metadatas associated with the texts. ids: Optional ids of each text object. timestamp: Optional timestamp to write new texts with. kwargs: vectorstore specific parameters Returns: List of ids from adding the texts into the vectorstore. rNc bg|],}ttjdtdz -Srrrs rrwz$TileDB.add_texts..s0IIIa3v~aa8899IIIrOrr)rrr5rgrrhrrf)rr7rrlrrrrrr|rrrN update_batchrrr}rrrxr3r4rI)rQrrr]r5r<rrXrrrrrrhrs r add_textszTileDB.add_textszs(h''^33DKK@@ ;II5IIICx}}++BI66 (C OOC888s:'' C CA*Q-rzBBBGAJJ &&%#,>>iit '   xV  Hc)nn%5VDDDMA%  #%=h1G1Grx#X#X#X a Q,D [[   #,>>iit; !  $( <  rz/tmp/tiledb_arrayc  `g} ||} |jd|| |||||||| d | S)aConstruct a TileDB index from raw documents. Args: texts: List of documents to index. embedding: Embedding function to use. metadatas: List of metadata dictionaries to associate with documents. ids: Optional ids of each text object. metric: Metric to use for indexing. Defaults to "euclidean". index_uri: The URI to write the TileDB arrays index_type: Optional, Vector index type ("FLAT", IVF_FLAT") config: Optional, TileDB config index_timestamp: Optional, timestamp to write new texts with. Example: .. code-block:: python from langchain_community import TileDB from langchain_community.embeddings import OpenAIEmbeddings embeddings = OpenAIEmbeddings() index = TileDB.from_texts(texts, embeddings) rrXr7rr]r9r8rAr4rr)r _TileDB__from) rrr7rr]r9r8rAr4rr<rXs r from_textszTileDB.from_textsseF ..u55 sz  !!+      rtext_embeddingsList[Tuple[str, List[float]]]c bd|D} d|D} |jd| | |||||||| d | S)aConstruct TileDB index from embeddings. Args: text_embeddings: List of tuples of (text, embedding) embedding: Embedding function to use. index_uri: The URI to write the TileDB arrays metadatas: List of metadata dictionaries to associate with documents. metric: Optional, Metric to use for indexing. Defaults to "euclidean". index_type: Optional, Vector index type ("FLAT", IVF_FLAT") config: Optional, TileDB config index_timestamp: Optional, timestamp to write new texts with. Example: .. code-block:: python from langchain_community import TileDB from langchain_community.embeddings import OpenAIEmbeddings embeddings = OpenAIEmbeddings() text_embeddings = embeddings.embed_documents(texts) text_embedding_pairs = list(zip(texts, text_embeddings)) db = TileDB.from_embeddings(text_embedding_pairs, embeddings) cg|] }|d Srrrmts rrwz*TileDB.from_embeddings.. s///!1///rcg|] }|d S)rrrs rrwz*TileDB.from_embeddings.. s444qad444rrr)r) rrr7r8rr]r9rAr4rr<rrXs rfrom_embeddingszTileDB.from_embeddingsssJ0////44O444 sz  !!+      r)r9r4r5c |d|||||d|S)a}Load a TileDB index from a URI. Args: index_uri: The URI of the TileDB vector index. embedding: Embeddings to use when generating queries. metric: Optional, Metric to use for indexing. Defaults to "euclidean". config: Optional, TileDB config timestamp: Optional, timestamp to use for opening the arrays. )r7r8r9r4r5rr)rr8r7r9r4r5r<s rloadz TileDB.loads;(s        rc 4|jjdi||_dS)Nr)rNconsolidate_updates)rQr<s rrzTileDB.consolidate_updates7s&AD-AKKFKKr)r7r r8rr9rr2rr3rr4r:r5rr6r;r<r)rrU) r]r^r_r`rZrar[rbr\rcrrd) r7r`rZrar[rbrrar<rrrd) rrrZrar[rbrrar<rrrd)rYNr) r7r`rZrar[rbrrar<rrr) rrrZrar[rbrrar<rrr)r7r`rZrarrarrcr[rbr<rrrd)rYrrN)r7r`rZrarrarrcr[rbr<rrr)rrrZrarrarrcr[rbr<rrr)r8rrArrrarrrr;r4r:rr)rrrXrr7r r8rrrr]rr9rrArr4r:rrar<rrr0)Nr)r]rr5rar<rrr)NNr) rrrrr]rr5rar<rrr)rrr7r rrr]rr9rr8rrArr4r:rrar<rrr0)rrr7r r8rrrr]rr9rrArr4r:rrar<rrr0)r8rr7r r9rr4r:r5rr<rrr0)r<rrr)__name__ __module__ __qualname____doc__rTpropertyrXrrrrrrrrr classmethodrDEFAULT_METRICrrrrr r rrrrr0r0>s@  &!# .205HHHHHHTX+/!*777777z+/ % % % % % % V+/ F+/ 33333@+/ 33333: +/999999| +/ "3"3"3"3"3N +/ #####J.2=====[=~+/#'$ .2 N N N N N [N bAB0+/#' 55555n +/#'$, .2 0 0 0 0 [0 d+/#'$ .2 3 3 3 3 3 [3 j %.2     [ 8LLLLLLrr0)rr)rrrr)r#rrr)-r __future__rr}rsystypingrrrrrr r numpyrlangchain_core.documentsr langchain_core.embeddingsr langchain_core.utilsrlangchain_core.vectorstoresr&langchain_community.vectorstores.utilsr frozensetrrr'r"iinformaxrzfinfor{ float_inforrr%r(r,r.r0rrrr$s,,"""""" FFFFFFFFFFFFFFFFFF------000000------333333MMMMMM ;-(( " RXhbhx(( ) ) - rx++,,0 N  (((( + + + +(((( ++++ z Lz Lz Lz Lz L[z Lz Lz Lz Lz Lr