a
    bg                     @  s  U d dl mZ d dlZd dlZd dlZd dlZd dlmZ d dlm	Z	m
Z
mZmZmZmZmZmZmZmZmZmZmZ d dlZd dlmZ d dlmZ d dlmZ d dlmZ e	rd dl Z ed	 Z!e"ee!Z#d
e$d< ed Z%e"ee%Z&de$d< dZ'dZ(dZ)e)d Z*g dZ+g dZ,ddi gZ-e.e/Z0ddddddddZ1dddddZ2dd dd!d"Z3dhdd%d&d'd(d)Z4G d*d+ d+eZ5diddd,d-d.d.d-d.d.d/d0
d1d2Z6djddd-d,d,d-d4d4d5d4d4d6d7d8d9Z7dd:d;d<d=d>Z8d?g dfddd:d-d@dAdBdCZ9dDddEdFdGZ:dDddHdIdJZ;dKdLdMdNdOZ<dPdd5dQdRdSZ=d&dddTdUdVdWZ>dXdYdZd[d\Z?dkdd]d]d/d^d_d`Z@ddDdadbdcZAd6dddedfdgZBdS )l    )annotationsN)deepcopy)TYPE_CHECKINGAnyCallableDictIterableListLiteralOptionalSizedTupleTypeUnionget_args)Document)
Embeddings)VectorStore)maximal_marginal_relevance)L2ZIPzList[DISTANCE_METRICS]AVAILABLE_DISTANCE_METRICS)ZTileDBDenseZTileDBSparse	FaissFlatZFaissIVFFlatZFlinngzList[ENGINES]AVAILABLE_ENGINESZ	langchain          )	_distanceidcontent)r   r   blobzMissing propertyr   strNone)xyx_namey_namereturnc                 C  sZ   t | trVt |trVt| t|krVt| d| d| dt|  d| dt| dS )z
    Check that sizes of two variables are the same

    Args:
        x: Variable to compare
        y: Variable to compare
        x_name: Name for variable x
        y_name: Name for variable y
    z and z% expected to be equal length but len(z)=z	 and len(N)
isinstancer   len
ValueError)r"   r#   r$   r%    r*   s/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/vectorstores/vdms.py_len_check_if_sized>   s    
$r,   List[Document])resultsr&   c                 C  s   dd t | D S )Nc                 S  s   g | ]\}}|qS r*   r*   .0doc_r*   r*   r+   
<listcomp>Q       z$_results_to_docs.<locals>.<listcomp>)_results_to_docs_and_scores)r.   r*   r*   r+   _results_to_docsP   s    r6   List[Tuple[Document, float]]c              
   C  s   g }z| d \}}t |dkrd|d v rd|d d v r|d d d }|D ]Z}t|d d}|d }tD ]}||v rl||= qldd | D }	|t||	d	|f qNW n4 ty }
 ztd
|
  W Y d }
~
n
d }
~
0 0 |S )Nr   FindDescriptorentitiesr   
   r   c                 S  s   i | ]\}}|t vr||qS r*   )INVALID_METADATA_VALUE)r0   ZmkeyZmvalr*   r*   r+   
<dictcomp>e   s   z/_results_to_docs_and_scores.<locals>.<dictcomp>)page_contentmetadataz2No results returned. Error while parsing results: )	r(   roundINVALID_DOC_METADATA_KEYSitemsappendr   	Exceptionloggerwarning)r.   Z	final_res	responsesZblobsZresult_entitiesentZdistanceZtxt_contentsppropser*   r*   r+   r5   T   s2    


&r5   	localhost  int	vdms.vdms)hostportr&   c                 C  s@   zddl }W n ty&   tdY n0 |  }|| | |S )zVDMS client for the VDMS server.

    Args:
        host: IP or hostname of VDMS server
        port: Port to connect to VDMS server
    r   NzOCould not import vdms python package. Please install it with `pip install vdms.)vdmsImportErrorconnect)rO   rP   rQ   clientr*   r*   r+   VDMS_Clients   s    
rU   c                   @  s|  e Zd ZdZdeddddddddd	d
ddddddZeddddZdddddZddddddZ	dddddZ
dd d!d"d#Zd$dd%d&Zeedfdd'd'd(dd)d*d+d,Zddddd-d.d/d0d1d2Zddd
d	dd3d4d5Zddd6d7d8d9d:d;Zddd<d=d<d>d?d@dAdBZdddDdDddEdFdGZg dCfdHdIdDdJdKdLdMZdddddd-ddNdOdPZdd/dIddQdRdSZddedTfddUd.d'dDdddVdWdXZddd.dUd.d'dDdddZd[d\Zddefd]dUd.d'ddd^d_d`ZdefddddUd'dddadbdcZdddddedfdgZdd'dhdidjZddkdldmdnZdd.d>dodd8dpdqdrZ dddd(dIdDdsdtdudvZ!eeddddCfddd'd'dwd(dxd8dydz	d{d|Z"ddd}d~dZ#e$ddeefdddd.d'ddd dddZ%e$dddeefddddUd.d'ddd d	ddZ&dddgfddodddJdddZ'eeddfdd'd'dddddddZ(eeddfd d'd'dddddddZ)eeddfdd'd'dddd)dddZ*eeddfd d'd'dddd)dddZ+ddeedddCfdd>d'd'd7d7d8ddd	ddZ,eedfdd'd'dddd*ddZ-eedfd d'd'ddddddZ.eedfdd'd'ddd)d*ddZ/eedfd d'd'ddd)dddZ0dddddddZ1dddddddZ2dS )VDMSa  Intel Lab's VDMS for vector-store workloads.

    To use, you should have both:
    - the ``vdms`` python package installed
    - a host (str) and port (int) associated with a deployed VDMS Server

    Visit https://github.com/IntelLabs/vdms/wiki more information.

    IT IS HIGHLY SUGGESTED TO NORMALIZE YOUR DATA.

    Args:
        client: VDMS Client used to connect to VDMS server
        collection_name: Name of data collection [Default: langchain]
        distance_strategy: Method used to calculate distances. VDMS supports
            "L2" (euclidean distance) or "IP" (inner product) [Default: L2]
        engine: Underlying implementation for indexing and computing distances.
            VDMS supports TileDBDense, TileDBSparse, FaissFlat, FaissIVFFlat,
            and Flinng [Default: FaissFlat]
        embedding: Any embedding function implementing
            `langchain_core.embeddings.Embeddings` interface.
        relevance_score_fn: Function for obtaining relevance score

    Example:
        .. code-block:: python

            from langchain_huggingface import HuggingFaceEmbeddings
            from langchain_community.vectorstores.vdms import VDMS, VDMS_Client

            model_name = "sentence-transformers/all-mpnet-base-v2"
            vectorstore = VDMS(
                client=VDMS_Client("localhost", 55555),
                embedding=HuggingFaceEmbeddings(model_name=model_name),
                collection_name="langchain-demo",
                distance_strategy="L2",
                engine="FaissFlat",
            )
    Nr   r   )	embeddingcollection_namedistance_strategyenginerelevance_score_fnembedding_dimensionsrN   zOptional[Embeddings]r    DISTANCE_METRICSENGINESz"Optional[Callable[[float], float]]Optional[int]r!   )rT   rW   rX   rY   rZ   r[   r\   r&   c                C  sD   || _ || _|| _|| _| || || _| j|| j| jd| _d S )NrZ   metric)_clientsimilarity_search_enginerY   rW   _check_required_inputsoverride_relevance_score_fnadd_set_collection_name)selfrT   rW   rX   rY   rZ   r[   r\   r*   r*   r+   __init__   s    zVDMS.__init__)r&   c                 C  s   | j S N)rW   rh   r*   r*   r+   
embeddings   s    zVDMS.embeddings	List[str]zList[List[float]])textsr&   c                 C  s0   t | jtr| j|S d}|d7 }t|d S )Nz*Must provide `embedding` which is expectedz to be an Embeddings object)r'   rW   r   Zembed_documentsr)   )rh   rn   p_strr*   r*   r+   _embed_documents   s
    zVDMS._embed_documentsr   )pathskwargsr&   c                 K  s8   | j d ur,t| j dr,| j jf d|i|S tdd S )Nembed_videorq   z:Must provide `embedding` which has attribute `embed_video`)rW   hasattrrs   r)   )rh   rq   rr   r*   r*   r+   _embed_video   s
    zVDMS._embed_video)urisr&   c                 C  s0   | j d ur$t| j dr$| j j|dS tdd S )Nembed_imagerv   z:Must provide `embedding` which has attribute `embed_image`)rW   rt   rw   r)   )rh   rv   r*   r*   r+   _embed_image   s
    zVDMS._embed_imagezList[float])textr&   c                 C  s$   t | jtr| j|S tdd S )NzEMust provide `embedding` which is expected to be an Embeddings object)r'   rW   r   embed_queryr)   )rh   rz   r*   r*   r+   _embed_query   s
    zVDMS._embed_queryzCallable[[float], float]c                 C  s<   | j dur| j S | j dv r&dd S td| j ddS )a8  
        The 'correct' relevance function
        may differ depending on a few things, including:
        - the distance / similarity metric used by the VectorStore
        - the scale of your embeddings (OpenAI's are unit normed. Many others are not!)
        - embedding dimensionality
        - etc.
        N)ipl2c                 S  s   | S rj   r*   )r"   r*   r*   r+   <lambda>   r4   z1VDMS._select_relevance_score_fn.<locals>.<lambda>z=No supported normalization function for distance_strategy of z;.Consider providing relevance_score_fn to VDMS constructor.)re   rY   lowerr)   rk   r*   r*   r+   _select_relevance_score_fn   s    	
zVDMS._select_relevance_score_fnrM   zOptional[Dict[str, Any]]r7   )querykfetch_kfilterrr   r&   c           
      K  sp   | j du rd|d< | jf ||||d|}g }|D ]6\}}	| j du rV|||	f q4|||  |	f q4|S )z?Return docs and their similarity scores on a scale from 0 to 1.NTnormalize_distance)r   r   r   r   )re   similarity_search_with_scorerB   )
rh   r   r   r   r   rr   docs_and_scoresZdocs_and_rel_scoresr1   Zscorer*   r*   r+   (_similarity_search_with_relevance_scores  s$    	

z-VDMS._similarity_search_with_relevance_scoresz1Optional[Union[List[None], List[Dict[str, Any]]]]zOptional[List[str]]r	   )rX   rn   rl   	metadatasidsr&   c                 C  s   t ||dd |d ur|ndd |D }t ||dd |d urB|ndd |D }t ||dd g }g }g }t||||D ]J\}	}
}}| j||	|
||d\}}|d urx|| || || qx| ||\}}|S )	Nrn   rl   c                 S  s   g | ]}d qS rj   r*   r0   r2   r*   r*   r+   r3   -  r4   zVDMS.add.<locals>.<listcomp>r   c                 S  s   g | ]}t t qS r*   r    uuiduuid4r   r*   r*   r+   r3   0  r4   r   r>   rW   documentr   )r,   zip_VDMS__get_add_queryrB   _VDMS__run_vdms_query)rh   rX   rn   rl   r   r   all_queries	all_blobsinserted_idsmetaembr1   r   r   r   responseresponse_arrayr*   r*   r+   add#  s$    



zVDMS.add)rX   rZ   ra   r&   c              	   C  sR   t d|| jt|d|t|d|d}| |g\}}d|d v rNtd| |S )NAddDescriptorSetvaluer`   ZFailedCommandr   zFailed to add collection )_add_descriptorsetembedding_dimensiongetattrr   r)   )rh   rX   rZ   ra   r   r   r2   r*   r*   r+   rf   D  s    

zVDMS.add_setzUnion[None, List[str]]zUnion[None, Dict[str, Any]]bool)rX   r   constraintsr&   c                 C  s   g }g }|  |}d|i}|du r0dddgi}nddg|d< |durTd|d g|d< td|ddddd||d		}|| | ||\}	}
td
|dd}| |g|\}}d|	d v S )zA
        Deletes entire collection if id is not provided
        listN	_deletion==   r   r   r8   labelrefrI   linkk_neighborsr   r.   FindDescriptorSetT)
storeIndex)_VDMS__get_properties_add_descriptorrB   r   r   )rh   rX   r   r   r   r   collection_propertiesr.   r   r   r   ZresponseSetr2   r*   r*   r+   Z__deleteY  s6    	

zVDMS.__deletezOptional[Any]Union[List[float], None]Optional[str]z4Tuple[Dict[str, Dict[str, Any]], Union[bytes, None]])rX   r>   rW   r   r   r&   c                 C  s   |d u ri }njd|i}t | j||\}}|rxdd |d d  D }	d| d}
|
d7 }
t|
 td	|	  |d fS |r|| |d
vr||d< | D ]}|| jvr| j| qtd|d d |d d d d d	}t	|}||fS )Nr   c                 S  s   i | ]\}}||d  qS )r*   )r0   Zprop_keyZprop_valr*   r*   r+   r<     s   z(VDMS.__get_add_query.<locals>.<dictcomp>r8   r   z[!] Embedding with id (z) exists in DB;z#Therefore, skipped and not insertedz	Skipped values are: )N r   ZAddDescriptorr   )
_check_descriptor_exists_by_idrb   rA   printupdatekeysr   rB   r   embedding2bytes)rh   rX   r>   rW   r   r   rI   Z	id_existsr   Zskipped_valuepstrr   r   r*   r*   r+   Z__get_add_query  sN    

zVDMS.__get_add_queryFOptional[bool]rX   unique_entitydeletionr&   c                 C  sJ   t |||d}| |g\}}t|dkr>t|d d}ntt}|S )N)r   r   r   ,)_find_property_entityr   r(   
_bytes2strsplitr   DEFAULT_PROPERTIES)rh   rX   r   r   Z
find_queryr   Zresponse_blobr   r*   r*   r+   Z__get_properties  s    zVDMS.__get_propertiesz
List[Dict]zOptional[List]Tuple[Any, Any])r   r   print_last_responser&   c                 C  s2   | j ||\}}t||}|r*| j   ||fS rj   )rb   r   _check_valid_responser   )rh   r   r   r   r   r   r2   r*   r*   r+   Z__run_vdms_query  s
    

zVDMS.__run_vdms_query)rX   r   	documentsrl   r   r&   c                 C  s  t ||dd t ||dd |dur(|ndd |D }t ||dd | |}g }t||||D ]\}}	}
}d| ji}d	d
dgi}|durd
|g|d< td|ddddd||d	}| |g\}}| j|||	|
|d\}}|dur`| |g|g\}}|| q`| ||| j dS )z
        Updates (find, delete, add) a collection based on id.
        If more than one collection returned with id, error occuers
        r   r   rl   Nc                 S  s   g | ]}d qS rj   r*   r   r*   r*   r+   r3     r4   z!VDMS.__update.<locals>.<listcomp>r   r   r   r   r   r   r8   r   r   )	r,   r   r   r   r   r   r   rB   _VDMS__update_properties)rh   rX   r   r   rl   r   
orig_propsZupdated_idsr   r   r1   r   r.   r   r   r   r   r   r*   r*   r+   Z__update  sH    


zVDMS.__update)rX   current_collection_propertiesnew_collection_propertiesr&   c           
      C  s\   |d urXt |}|D ]}||vr|| q||krXt|d|d\}}| ||g\}}	d S )Nr   )command_typeall_properties)r   rB   _build_property_queryr   )
rh   rX   r   r   Zold_collection_propertiespropr   blob_arrr   r2   r*   r*   r+   Z__update_properties  s    
zVDMS.__update_propertiesTzOptional[List[dict]])rv   r   r   
batch_sizeadd_pathrr   r&   c                   s    fdd|D }|r:|r:t |D ]\}}	|	|| d< q"n |rZg }|D ]}	|d|	i qF|durf|ndd |D } j|d}
|du rdd |D }ndd |D } jf ||
|||d	| |S )
a  Run more images through the embeddings and add to the vectorstore.

        Images are added as embeddings (AddDescriptor) instead of separate
        entity (AddImage) within VDMS to leverage similarity search capability

        Args:
            uris: List of paths to the images to add to the vectorstore.
            metadatas: Optional list of metadatas associated with the images.
            ids: Optional list of unique IDs.
            batch_size (int): Number of concurrent requests to send to the server.
            add_path: Bool to add image path as metadata

        Returns:
            List of ids from adding images into the vectorstore.
        c                   s   g | ]} j |d qS ))
image_path)encode_image)r0   urirk   r*   r+   r3   C  r4   z#VDMS.add_images.<locals>.<listcomp>r   Nc                 S  s   g | ]}t t qS r*   r   r   r*   r*   r+   r3   N  r4   rx   c                 S  s   g | ]}i qS r*   r*   r   r*   r*   r+   r3   T  r4   c                 S  s   g | ]}t |qS r*   _validate_vdms_propertiesr0   mr*   r*   r+   r3   V  r4   rn   rl   r   r   r   )	enumeraterB   ry   add_from)rh   rv   r   r   r   r   rr   Z	b64_textsmidxr   rl   r*   rk   r+   
add_images*  s.    zVDMS.add_imagesr   )rq   rn   r   r   r   r   rr   r&   c                 K  s   |du rdd |D }|r>|r>t |D ]\}}	|	|| d< q&n |r^g }|D ]}	|d|	i qJ|durj|ndd |D }| jf d|i|}
|du rdd |D }| jf ||
|||d| |S )	a  Run videos through the embeddings and add to the vectorstore.

        Videos are added as embeddings (AddDescriptor) instead of separate
        entity (AddVideo) within VDMS to leverage similarity search capability

        Args:
            paths: List of paths to the videos to add to the vectorstore.
            metadatas: Optional list of text associated with the videos.
            metadatas: Optional list of metadatas associated with the videos.
            ids: Optional list of unique IDs.
            batch_size (int): Number of concurrent requests to send to the server.
            add_path: Bool to add video path as metadata

        Returns:
            List of ids from adding videos into the vectorstore.
        Nc                 S  s   g | ]}d qS )r   r*   r   r*   r*   r+   r3   }  r4   z#VDMS.add_videos.<locals>.<listcomp>Z
video_pathc                 S  s   g | ]}t t qS r*   r   r   r*   r*   r+   r3     r4   rq   c                 S  s   g | ]}i qS r*   r*   r   r*   r*   r+   r3     r4   r   )r   rB   ru   r   )rh   rq   rn   r   r   r   r   rr   r   pathrl   r*   r*   r+   
add_videosb  s.    zVDMS.add_videoszIterable[str])rn   r   r   r   rr   r&   c                 K  sn   t |}|du rdd |D }| |}|du r@dd |D }ndd |D }| jf |||||d|}|S )a  Run more texts through the embeddings and add to the vectorstore.

        Args:
            texts: List of strings to add to the vectorstore.
            metadatas: Optional list of metadatas associated with the texts.
            ids: Optional list of unique IDs.
            batch_size (int): Number of concurrent requests to send to the server.

        Returns:
            List of ids from adding the texts into the vectorstore.
        Nc                 S  s   g | ]}t t qS r*   r   r   r*   r*   r+   r3     r4   z"VDMS.add_texts.<locals>.<listcomp>c                 S  s   g | ]}i qS r*   r*   r   r*   r*   r+   r3     r4   c                 S  s   g | ]}t |qS r*   r   r   r*   r*   r+   r3     r4   r   )r   rp   r   )rh   rn   r   r   r   rr   rl   r   r*   r*   r+   	add_texts  s"    
zVDMS.add_texts)rn   rl   r   r   r   rr   r&   c                 K  s   |  | j}g }tdt||D ]j}	t|	| t|}
||	|
 }||	|
 }||	|
 }|rj||	|
 }| j| j||||d}|| q | | j|| j |S )Nr   )rl   rn   r   r   )	r   rg   ranger(   minr   extendr   r   )rh   rn   rl   r   r   r   rr   r   r   Z	start_idxZend_idxZbatch_textsZbatch_embedding_vectorsZ	batch_idsZbatch_metadatasresultr*   r*   r+   r     s*    

zVDMS.add_fromzUnion[int, None])rX   r\   r&   c                 C  s  | j  std| jtvr$td| jtvr6td| jd u rHtd|d urX|| _n| jd urt	| jdrt
| d| _nf| jd urt	| jdst	| jdrt	| jd	rz| jjjj| _W q ty   td
Y q0 ntd
| |}t	| dr
| j| n|| _d S )Nz_VDMS client must be connected to a VDMS server.Please use VDMS_Client to establish a connectionz-distance_strategy must be either 'L2' or 'IP'z]engine must be either 'TileDBDense', 'TileDBSparse', 'FaissFlat', 'FaissIVFFlat', or 'Flinng'Must provide embedding functionr{   zThis is a sample sentence.rw   rs   modelz>Embedding dimension needed. Please define embedding_dimensionsr   )rb   Zis_connectedr)   rY   r   rc   r   rW   r   rt   r(   r|   r   Ztoken_embeddingZembedding_dimr   r   r   )rh   rX   r\   Zcurrent_propsr*   r*   r+   rd     sL    








zVDMS._check_required_inputs)rX   r&   c                 C  sX   g }g }ddgd}t d|d d d d d d |d	}|| | ||\}}|d d d S )Nr   r   )countr   r8   r   r   returned)r   rB   r   )rh   rX   r   r   r.   r   r   r   r*   r*   r+   r      s"    
z
VDMS.countbytes)base64_imager&   c                 C  s
   t |S rj   )base64	b64decode)rh   r   r*   r*   r+   decode_image6  s    zVDMS.decode_imageOptional[Dict])r   rX   r   rr   r&   c                 K  s"   |dur|n| j }| j|||dS )zDelete by ID. These are the IDs in the vectorstore.

        Args:
            ids: List of ids to delete.

        Returns:
            Optional[bool]: True if deletion is successful,
            False otherwise, None if not implemented.
        N)r   r   )rg   _VDMS__delete)rh   r   rX   r   rr   namer*   r*   r+   delete9  s    zVDMS.deletez(Tuple[List[Dict[str, Any]], List, float])setnamer   r.   r   	normalizer&   c                 C  s\   d}d}t ||||d}| |g|\}	}
|rR||	d v rR|	d | d d d }|	|
|fS )Nr   r8   )r   r.   r   r9   r   r   )r   r   )rh   r   r   r.   r   r   max_distcommand_strr   r   r   r*   r*   r+   get_k_candidatesL  s    zVDMS.get_k_candidatesOptional[dict]zOptional[List[float]]z!Tuple[List[Dict[str, Any]], List])	r   r   r   r   r   r.   query_embeddingr   r&   c	                 C  s6  g }	t |}
|
d ur|	|
 |d u rD| j||||	|d\}}}n|d u rXddgi}n.d|vrldg|d< nd|d vr|d d t||||d}| |g\}}||d v r|d | d dkrdd |d | d	 D }ng g fS | j||||	|d\}}}||d vs6||d v r>|d | d dkr>g g fS g }|d | d	 D ]2}|d |v rn|| t||krR qqR||d | d	< t||d | d< t||k rd
}t| |r.|dtjfv rdn|}t	|d | d	 D ]6\}}|d | |d< |d |d | d	 | d< q||fS )N)r   r   r   r   r.   r   r   c                 S  s   g | ]}|d  qS )r   r*   )r0   rG   r*   r*   r+   r3     s   z0VDMS.get_descriptor_response.<locals>.<listcomp>r9   z4Returned items < k_neighbors; Try increasing fetch_kg      ?r   )
r   rB   r   r   r   r(   r   npinfr   )rh   r   r   r   r   r   r.   r   r   r   r   r   r   r   r   Zids_of_interestZnew_entitiesrG   ro   Zent_idxr*   r*   r+   get_descriptor_responsec  sl    

 



zVDMS.get_descriptor_response)r   r&   c                 C  sF   t |d(}| }t|dW  d    S 1 s80    Y  d S )Nrbzutf-8)openreadr   	b64encodedecode)rh   r   fr   r*   r*   r+   r     s    zVDMS.encode_imagez
Type[VDMS]r-   )clsr   rW   r   r   rX   rr   r&   c              	   K  s4   |d }| j |dd |D dd |D ||||dS )a  Create a VDMS vectorstore from a list of documents.

        Args:
            collection_name (str): Name of the collection to create.
            documents (List[Document]): List of documents to add to vectorstore.
            embedding (Embeddings): Embedding function. Defaults to None.
            ids (Optional[List[str]]): List of document IDs. Defaults to None.
            batch_size (int): Number of concurrent requests to send to the server.

        Returns:
            VDMS: VDMS vectorstore.
        rT   c                 S  s   g | ]
}|j qS r*   r=   r0   r1   r*   r*   r+   r3     r4   z'VDMS.from_documents.<locals>.<listcomp>c                 S  s   g | ]
}|j qS r*   )r>   r  r*   r*   r+   r3     r4   )rT   rn   r   rW   r   r   rX   )
from_texts)r  r   rW   r   r   rX   rr   rT   r*   r*   r+   from_documents  s    zVDMS.from_documents)	r  rn   rW   r   r   r   rX   rr   r&   c           
      K  sB   |d }| |||d}	|du r,dd |D }|	j ||||d |	S )aH  Create a VDMS vectorstore from a raw documents.

        Args:
            texts (List[str]): List of texts to add to the collection.
            embedding (Embeddings): Embedding function. Defaults to None.
            metadatas (Optional[List[dict]]): List of metadatas. Defaults to None.
            ids (Optional[List[str]]): List of document IDs. Defaults to None.
            batch_size (int): Number of concurrent requests to send to the server.
            collection_name (str): Name of the collection to create.

        Returns:
            VDMS: VDMS vectorstore.
        rT   )rX   rW   rT   Nc                 S  s   g | ]}t t qS r*   r   r   r*   r*   r+   r3     r4   z#VDMS.from_texts.<locals>.<listcomp>)rn   r   r   r   )r   )
r  rn   rW   r   r   r   rX   rr   rT   Zvdms_collectionr*   r*   r+   r    s    zVDMS.from_textsr>   )rX   r   limitincluder&   c                 C  s~   g }g }ddi}|dur ||d< d|v r:|  |}||d< d|v rJd|d	< td
|d||d}	||	 | ||\}
}|
|fS )a  Gets the collection.
        Get embeddings and their associated data from the data store.
        If no constraints provided returns all embeddings up to limit.

        Args:
            constraints: A dict used to filter results by.
                   E.g. `{"color" : ["==", "red"], "price": [">", 4.00]}`. Optional.
            limit: The number of documents to return. Optional.
            include: A list of what to include in the results.
                     Can contain `"embeddings"`, `"metadatas"`, `"documents"`.
                     Ids are always included.
                     Defaults to `["metadatas", "documents"]`. Optional.
        r   r   Nr
  r>   r   rl   Tr   r8   )r   r   r.   )r   r   rB   r   )rh   rX   r   r
  r  r   r   r.   r   r   r   r   r*   r*   r+   get   s(    

zVDMS.getg      ?floatzOptional[Dict[str, List]])r   r   r   lambda_multr   rr   r&   c           
      K  s   | j du rtdtj|s6t| j dr6| |}n|tj|rbt| j drb| j|gdd }nPtj|rt| j dr| j|gdd }n$d	| d
}|d7 }|d7 }t|| j	|||||d}	|	S aR  Return docs selected using the maximal marginal relevance.
        Maximal marginal relevance optimizes for similarity to query AND diversity
        among selected documents.

        Args:
            query (str): Query to look up. Text or path for image or video.
            k: Number of Documents to return. Defaults to 4.
            fetch_k: Number of Documents to fetch 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.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents selected by maximal marginal relevance.
        NzBFor MMR search, you must specify an embedding function oncreation.r{   rw   rx   r   rs   rq   (Could not generate embedding for query ''.9If using path for image or video, verify embedding model 6has callable functions 'embed_image' or 'embed_video'.)r  r   )
rW   r)   osr   isfilert   r|   ry   ru   'max_marginal_relevance_search_by_vector)
rh   r   r   r   r  r   rr   Zembedding_vector	error_msgdocsr*   r*   r+   max_marginal_relevance_search2  s,    
z"VDMS.max_marginal_relevance_search)rW   r   r   r  r   rr   r&   c                   s   | j |g||g dd}t|d d dkr0g S dd |d d D }ttj|tjd|||d t|}	 fd	dt|	D }
|
S d
S )H  Return 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 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.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents selected by maximal marginal relevance.
        r   r   Z	distancesrl   query_embeddings	n_resultsr   r  r   r   c                 S  s   g | ]}t t|qS r*   r   _bytes2embeddingr0   r   r*   r*   r+   r3     s   z@VDMS.max_marginal_relevance_search_by_vector.<locals>.<listcomp>Zdtyper   r  c                   s   g | ]\}}| v r|qS r*   r*   )r0   irZmmr_selectedr*   r+   r3     s   N)query_collection_embeddingsr(   r   r   arrayfloat32r6   r   rh   rW   r   r   r  r   rr   r.   Zembedding_list
candidatesZselected_resultsr*   r'  r+   r  h  s,    

z,VDMS.max_marginal_relevance_search_by_vectorc           
      K  s   | j du rtdtj|s6t| j dr6| |}n|tj|rbt| j drb| j|gdd }nPtj|rt| j dr| j|gdd }n$d	| d
}|d7 }|d7 }t|| j	|||||d}	|	S r  )
rW   r)   r  r   r  rt   r|   ry   ru   2max_marginal_relevance_search_with_score_by_vector)
rh   r   r   r   r  r   rr   rW   r  r  r*   r*   r+   (max_marginal_relevance_search_with_score  s,    
z-VDMS.max_marginal_relevance_search_with_scorec                   s   | j |g||g dd}t|d d dkr0g S dd |d d D }ttj|tjd|||d t|}	 fd	dt|	D }
|
S d
S )r  r  r  r   r   c                 S  s   g | ]}t t|qS r*   r   r"  r*   r*   r+   r3     s   zKVDMS.max_marginal_relevance_search_with_score_by_vector.<locals>.<listcomp>r#  r$  c                   s$   g | ]\}\}}| v r||fqS r*   r*   )r0   r%  r&  sr'  r*   r+   r3     s   N)r(  r(   r   r   r)  r*  r5   r   r+  r*   r'  r+   r-    s,    

z7VDMS.max_marginal_relevance_search_with_score_by_vectorzOptional[List[List[float]]]z!List[Tuple[Dict[str, Any], List]])	r  rX   r  r   r   r.   r   rr   r&   c                 K  s   g }	|d u r| j }|d u r|	S |ddg}
|d u rLd|
v rL| jd|
v d}|D ]0}| jd|||||||d\}}|	||g qP|	S )Nr  r   rl   )r   r   r8   )r   r   r   r.   r   r   )rg   r  r   r   rB   )rh   r  rX   r  r   r   r.   r   rr   Zall_responsesr  Zqembr   r   r*   r*   r+   r(  
  s.    

z VDMS.query_collection_embeddingsc                 K  s(   | j |f|||d|}dd |D S )a  Run similarity search with VDMS.

        Args:
            query (str): Query to look up. Text or path for image or video.
            k (int): Number of results to return. Defaults to 3.
            fetch_k (int): Number of candidates to fetch for knn (>= k).
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List[Document]: List of documents most similar to the query text.
        )r   r   r   c                 S  s   g | ]\}}|qS r*   r*   r/   r*   r*   r+   r3   I  r4   z*VDMS.similarity_search.<locals>.<listcomp>)r   )rh   r   r   r   r   rr   r   r*   r*   r+   similarity_search3  s    zVDMS.similarity_search)rW   r   r   r   rr   r&   c                 K  s$   | j f |g|||d|}t|S )a  Return docs most similar to embedding vector.
        Args:
            embedding (List[float]): Embedding to look up documents similar to.
            k (int): Number of Documents to return. Defaults to 3.
            fetch_k (int): Number of candidates to fetch for knn (>= k).
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.
        Returns:
            List of Documents most similar to the query vector.
        r  r  r   r   )r(  r6   rh   rW   r   r   r   rr   r.   r*   r*   r+   similarity_search_by_vectorK  s    z VDMS.similarity_search_by_vectorc           	      K  s   | j du rtdntj|s8t| j dr8| |}n|tj|rdt| j drd| j|gdd }nPtj|rt| j dr| j|gdd }n$d	| d
}|d7 }|d7 }t|| j	f |g|||d|}t
|S )aE  Run similarity search with VDMS with distance.

        Args:
            query (str): Query to look up. Text or path for image or video.
            k (int): Number of results to return. Defaults to 3.
            fetch_k (int): Number of candidates to fetch for knn (>= k).
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List[Tuple[Document, float]]: List of documents most similar to
            the query text and cosine distance in float for each.
            Lower score represents more similarity.
        Nr   r{   rw   rx   r   rs   r  r  r  r  r  r1  )rW   r)   r  r   r  rt   r|   ry   ru   r(  r5   )	rh   r   r   r   r   rr   r   r  r.   r*   r*   r+   r   f  s*    

z!VDMS.similarity_search_with_scorec                 K  s$   | j f |g|||d|}t|S )a6  
        Return docs most similar to embedding vector and similarity score.

        Args:
            embedding (List[float]): Embedding to look up documents similar to.
            k (int): Number of Documents to return. Defaults to 3.
            fetch_k (int): Number of candidates to fetch for knn (>= k).
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List[Tuple[Document, float]]: List of documents most similar to
            the query text. Lower score represents more similarity.
        r1  )r(  r5   r2  r*   r*   r+   &similarity_search_with_score_by_vector  s    z+VDMS.similarity_search_with_score_by_vectorr   )rX   document_idr   r&   c                 C  s   |  ||g|gS )zUpdate a document in the collection.

        Args:
            document_id (str): ID of the document to update.
            document (Document): Document to update.
        )update_documents)rh   rX   r5  r   r*   r*   r+   update_document  s    	zVDMS.update_document)rX   r   r   r&   c                 C  s>   dd |D }dd |D }|  |}| j|||||d dS )zUpdate a document in the collection.

        Args:
            ids (List[str]): List of ids of the document to update.
            documents (List[Document]): List of documents to update.
        c                 S  s   g | ]
}|j qS r*   r  r0   r   r*   r*   r+   r3     r4   z)VDMS.update_documents.<locals>.<listcomp>c                 S  s   g | ]}t |jqS r*   )r   r>   r8  r*   r*   r+   r3     s   )r   rl   r   N)rp   _VDMS__update)rh   rX   r   r   rz   r>   rl   r*   r*   r+   r6    s    	
zVDMS.update_documents)NN)r   r   )NN)NNNN)FF)N)NNNr   T)NNN)NNF)3__name__
__module____qualname____doc__DEFAULT_COLLECTION_NAMEri   propertyrl   rp   ru   ry   r|   r   	DEFAULT_KDEFAULT_FETCH_Kr   r   rf   r   r   r   r   r9  r   DEFAULT_INSERT_BATCH_SIZEr   r   r   r   rd   r   r   r   r   r   r   classmethodr	  r  r  r  r  r.  r-  r(  r0  r3  r   r4  r7  r6  r*   r*   r*   r+   rV      s  * 	"  $    0    ;   ;;     ;.%8       L ""+59:79 ,1!rV   r   r_   r   zDict[str, Dict[str, Any]])
r   r   r   r   rI   r   r   r   r.   r&   c	                 C  s   d|i}	d| v r|r||	d< |d ur,||	d< |t vr<||	d< d| v rT|d urT||	d< d| v rp|d urpt||	d< d| v r|t vr||	d	< d| v r|t vr||	d
< | |	i}
|
S )NsetAddr   _ref
propertiesr   ZFindr   r   r.   )r;   rM   )r   r   r   r   rI   r   r   r   r.   entityr   r*   r*   r+   r     s"    r   Fr   r   zDict[str, Any])r   r   num_dimsrZ   ra   r   rI   r   r   r   r.   r&   c                 C  s   | dkr~t dd ||fD r~||d}|d ur8||d< |d urH||d< |d urX||d< |d i fvrl||d< |d ur||d	< nP| d
krd|i}|r||d< |	d i fvr|	|d< |
d ur|
|d< ntd|  | |i}|S )Nr   c                 s  s   | ]}|d uV  qd S rj   r*   )r0   varr*   r*   r+   	<genexpr>  s   z%_add_descriptorset.<locals>.<genexpr>)r   
dimensionsrZ   ra   rF  rG  r   r   rD  r   r   r.   zUnknown command: )allr)   )r   r   rI  rZ   ra   r   rI   r   r   r   r.   rH  r   r*   r*   r+   r     s6    

r   r	   zTuple[Dict[str, Any], bytes])rX   r   r&   c                 C  sn   t |dkrd|nd}d}i }d|d< d|d< d	| i}d
|d< ||d< ||d< t|}i }|||< ||fS )Nr   r   r   Z	AddEntityrG  classTr   r   zqueryable propertiestyper   )r(   join
_str2bytes)rX   r   Zall_properties_str	querytyperH  rI   	byte_datar   r*   r*   r+   _add_entity_with_blob8  s    rT  findr   )rX   r   r   r   r&   c           	      C  s   g }g }g d}|  |vr0tdd||  dkrTt| dd}|| nt|  dkrt| |\}}|| || nD|  dkrt| dd	}|| t| |\}}|| || ||fS )
N)rU  r   r   z"[!] Invalid type. Choices are : {}r   rU  T)r   r   r   )r   )r   r)   formatrP  r   rB   rT  )	rX   r   r   r   r   r   choicesr   rS  r*   r*   r+   r   N  s&    



r   r   )r   r&   c                 C  s   t j| dd}|S )Nr*  r#  )r   Z
frombuffer)r   r   r*   r*   r+   r!  q  s    r!  )in_bytesr&   c                 C  s   |   S rj   )r  )rX  r*   r*   r+   r   v  s    r   r   rm   )r   r&   c                 C  s   t tdd | D S )Nc                 S  s   g | ]}|  D ]}|qqS r*   )r   )r0   qr   r*   r*   r+   r3   {  r4   z(_get_cmds_from_query.<locals>.<listcomp>)r   rD  )r   r*   r*   r+   _get_cmds_from_queryz  s    rZ  z
List[dict])r   r   r&   c                   s,   t | }t to&t fdd|D }|S )Nc                 3  s>   | ]6}| d  v o4d d  | v o4 d  | d d kV  qdS )r   r   Nr*   )r0   cmdr   r*   r+   rK    s
   z(_check_valid_response.<locals>.<genexpr>)rZ  r'   r   any)r   r   Zcmd_list	valid_resr*   r\  r+   r   ~  s
    r   zTuple[bool, Any])rT   r   r   r&   c           	      C  sJ   dd|gi}t d||dgddd}|g}| |\}}t||}||fS )Nr   r   r8   r   )r   r   r   )r   r   r   )	rT   r   r   r   ZfindDescriptorr   resr2   r^  r*   r*   r+   r     s    

r   r   zUnion[bytes, None])rW   r&   c                 C  s&   d}| dur"t j| dd}| }|S )zConvert embedding to bytes.Nr*  r#  )r   r)  tobytes)rW   r   r   r*   r*   r+   r     s
    r   r   r   c                 C  sz   d}i }d|d< |r||d< i }d|d< d|d< d	g|d
< ||d< i }|rVddg|d< d| g|d< ||d< i }|||< |S )NZ
FindEntityrG  rN  uniqueTr   r   r   r   r   r.   r   r   r   r   r   r*   )rX   r   r   rR  rH  r.   r   r   r*   r*   r+   r     s$    
r   )in_strr&   c                 C  s
   t | S rj   )r    encode)rb  r*   r*   r+   rQ    s    rQ  r   )r>   r&   c                 C  s0   i }|   D ]\}}t|ts||t|< q|S rj   )rA   r'   r   r    )r>   Znew_metadatakeyr   r*   r*   r+   r     s
    
r   )rK   rL   )NNNNNNN)	NNNNNNFNN)FF)C
__future__r   r   loggingr  r   copyr   typingr   r   r   r   r   r	   r
   r   r   r   r   r   r   numpyr   Zlangchain_core.documentsr   Zlangchain_core.embeddingsr   Zlangchain_core.vectorstoresr   Z&langchain_community.vectorstores.utilsr   rQ   r]   r   r   __annotations__r^   r   r>  rB  r@  rA  r   r@   r;   	getLoggerr:  rD   r,   r6   r5   rU   rV   r   r   rT  r   r!  r   rZ  r   r   r   r   rQ  r   r*   r*   r*   r+   <module>   s   <

          a       ")         &7#  