a bgQ@sddlmZddlZddlZddlZddlmZddlmZm Z m Z m Z m Z m Z mZddlZddlmZddlmZddlmZddlmZdd lmZmZdd lmZdd lmZd d d ddZ ddddZ!GdddeZ"dS)) annotationsN)Path)AnyCallableDictIterableListOptionalTuple)Document) Embeddings guard_import) VectorStore) AddableMixinDocstore)InMemoryDocstore)DistanceStrategyz np.ndarray)xreturncCs$|ttjj|ddddd}|S)z!Normalize vectors to unit length.T)ZaxisZkeepdimsg-q=N)npZclipZlinalgZnorm)rrt/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/vectorstores/scann.py normalizes rrrcCstdS)z= Import `scann` if available, otherwise raise error. scannr rrrrdependable_scann_importsrc @seZdZdZddejdfdddddd d d d d dZdNdddddddddZdOddddddddZdPddddddddZ dQddddd d!Z dRd$d%d&d%dd'd(d)d*Z dSd+d%d&d%dd'd,d-d.Z dTd$d%d&d%dd/d(d0d1Z dUd+d%d&d%dd/d,d2d3ZedVdd4dddd ddd5d6d7ZedWddddddd8d9d:ZedXd;dddddd<d=d>ZdYd+d+d@dAdBdCZedZddDd+dd+d dddEdFdGZdHdIdJdKZd[d+d%d&d%dd'd,dLdMZdS)\ScaNNa`ScaNN` vector store. To use, you should have the ``scann`` python package installed. Example: .. code-block:: python from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_community.vectorstores import ScaNN model_name = "sentence-transformers/all-mpnet-base-v2" db = ScaNN.from_texts( ['foo', 'bar', 'barz', 'qux'], HuggingFaceEmbeddings(model_name=model_name)) db.similarity_search('foo?', k=1) NFr rrzDict[int, str]z"Optional[Callable[[float], float]]boolrz Optional[str]) embeddingindexdocstoreindex_to_docstore_idrelevance_score_fn normalize_L2distance_strategy scann_configc Cs4||_||_||_||_||_||_||_||_dS)z%Initialize with necessary components.N)r r!r"r#r&override_relevance_score_fn _normalize_L2Z _scann_config) selfr r!r"r#r$r%r&r'rrr__init__3s zScaNN.__init__z Iterable[str]zIterable[List[float]]zOptional[List[dict]]zOptional[List[str]]z List[str])texts embeddings metadatasidskwargsrcKs*t|jtstd|jdtddS)NSIf trying to add texts, the underlying docstore should support adding items, which does notz(Updates are not available in ScaNN, yet.) isinstancer"r ValueErrorNotImplementedError)r*r,r-r.r/r0rrrZ__addHs z ScaNN.__add)r,r.r/r0rcKs*|jt|}|j||f||d|S)alRun 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 list of unique IDs. Returns: List of ids from adding the texts into the vectorstore. r.r/)r embed_documentslist _ScaNN__add)r*r,r.r/r0r-rrr add_textsWszScaNN.add_textsz!Iterable[Tuple[str, List[float]]])text_embeddingsr.r/r0rcKsDt|jtstd|jdt|\}}|j||f||d|S)aRun more texts through the embeddings and add to the vectorstore. Args: text_embeddings: Iterable pairs of string and embedding to add to the vectorstore. metadatas: Optional list of metadatas associated with the texts. ids: Optional list of unique IDs. Returns: List of ids from adding the texts into the vectorstore. r1r2r6)r3r"rr4zipr9)r*r;r.r/r0r,r-rrradd_embeddingsls  zScaNN.add_embeddingszOptional[bool])r/r0rcKs tddS)a3Delete by vector ID or other criteria. Args: ids: List of ids to delete. **kwargs: Other keyword arguments that subclasses might use. Returns: Optional[bool]: True if deletion is successful, False otherwise, None if not implemented. z*Deletions are not available in ScaNN, yet.N)r5)r*r/r0rrrdeletes z ScaNN.deletez List[float]intzOptional[Dict[str, Any]]zList[Tuple[Document, float]])r kfilterfetch_kr0rc  sLtj|gtjd}|jr t|}|j||dur4|n|\}}g} t|dD]\} } | dkr`qN|j| } |j | t t st d| d|durdd|D}tfd d |Dr| |d| fqN| |d| fqN|d dur@|jtjtjfvr&tjntjfd d | D} | d|S)aReturn 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: score_threshold: Optional, a floating point value between 0 to 1 to filter the resulting set of retrieved docs Returns: List of documents most similar to the query text and L2 distance in float for each. Lower score represents more similarity. ZdtypeNrrzCould not find document for id z, got cSs&i|]\}}|t|ts|gn|qSr)r3r8.0keyvaluerrr sz@ScaNN.similarity_search_with_score_by_vector..c3s"|]\}}j||vVqdS)N)metadatagetrF)docrr z?ScaNN.similarity_search_with_score_by_vector..score_thresholdcs"g|]\}}|r||fqSrrrGrMZ similarity)cmprPrr s z@ScaNN.similarity_search_with_score_by_vector..)rarrayfloat32r)rr!Zsearch_batched enumerater#r"searchr3r r4itemsallappendrLr&rMAX_INNER_PRODUCTZJACCARDoperatorgele) r*r rBrCrDr0vectorindicesZscoresdocsjiZ_idr)rRrMrPr&similarity_search_with_score_by_vectorsB        z,ScaNN.similarity_search_with_score_by_vectorstr)queryrBrCrDr0rcKs*|j|}|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 L2 distance in float. Lower score represents more similarity. rCrD)r Z embed_queryrd)r*rfrBrCrDr0r rarrrsimilarity_search_with_scores z"ScaNN.similarity_search_with_scorezList[Document]cKs(|j||f||d|}dd|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. rgcSsg|] \}}|qSrrrGrM_rrrrSrOz5ScaNN.similarity_search_by_vector..)rd)r*r rBrCrDr0docs_and_scoresrrrsimilarity_search_by_vectorsz!ScaNN.similarity_search_by_vectorcKs(|j||f||d|}dd|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. rgcSsg|] \}}|qSrrrirrrrS)rOz+ScaNN.similarity_search..)rh)r*rfrBrCrDr0rkrrrsimilarity_searchszScaNN.similarity_searchzList[List[float]])r,r-r r.r/r%r0rcKsJtd}|dtj} |dd} tj|tjd} |r>t| } | durV|j | | } n<| tj krz|j | dd } n|j | dd } g} |durdd |D}t|D]*\}}|r||ni}| t||d qtt|}t|t| krtt|d t| d ttt|| }||| ||fd |i|S)Nrr&r'rEZ dot_productZ squared_l2cSsg|]}ttqSr)reuuiduuid4)rGrjrrrrSQrOz ScaNN.__from..)Z page_contentrKz ids provided for z, documents. Each document should have an id.r%)rrLrEUCLIDEAN_DISTANCErrTrUrscann_ops_pybindZcreate_searcherr[ZbuilderZscore_brute_forcebuildrVrZr dictlen Exceptionrr<values)clsr,r-r r.r/r%r0rr&r'r_r!Z documentsrctextrKZ index_to_idr"rrrZ__from+sJ    z ScaNN.__from)r,r r.r/r0rcKs&||}|j|||f||d|S)aNConstruct ScaNN wrapper from raw documents. This is a user friendly interface that: 1. Embeds documents. 2. Creates an in memory docstore 3. Initializes the ScaNN database This is intended to be a quick way to get started. Example: .. code-block:: python from langchain_community.vectorstores import ScaNN from langchain_community.embeddings import OpenAIEmbeddings embeddings = OpenAIEmbeddings() scann = ScaNN.from_texts(texts, embeddings) r6)r7 _ScaNN__from)rxr,r r.r/r0r-rrr from_textsgs zScaNN.from_textszList[Tuple[str, List[float]]])r;r r.r/r0rcKs8dd|D}dd|D}|j|||f||d|S)aConstruct ScaNN wrapper from raw documents. This is a user friendly interface that: 1. Embeds documents. 2. Creates an in memory docstore 3. Initializes the ScaNN database This is intended to be a quick way to get started. Example: .. code-block:: python from langchain_community.vectorstores import ScaNN from langchain_community.embeddings import OpenAIEmbeddings embeddings = OpenAIEmbeddings() text_embeddings = embeddings.embed_documents(texts) text_embedding_pairs = list(zip(texts, text_embeddings)) scann = ScaNN.from_embeddings(text_embedding_pairs, embeddings) cSsg|] }|dqS)rrrGtrrrrSrOz)ScaNN.from_embeddings..cSsg|] }|dqS)rnrr|rrrrSrOr6)rz)rxr;r r.r/r0r,r-rrrfrom_embeddingsszScaNN.from_embeddingsr!None) folder_path index_namercCst|}|dj|d}|jddd|jt|t|dj|dd$}t|j |j f|Wdn1sv0YdS)zSave ScaNN index, docstore, and index_to_docstore_id to disk. Args: folder_path: folder path to save index, docstore, and index_to_docstore_id to. {index_name}.scannrTexist_okparents{index_name}.pklwbN) rformatmkdirr! serializereopenpickledumpr"r#)r*rrpath scann_pathfrrr save_locals zScaNN.save_local)allow_dangerous_deserialization)rr rrr0rc Ks|s tdt|}|dj|d}|jdddtd}|jt|} t|dj|dd} t | \} } Wd n1s0Y||| | | fi|S) aLoad ScaNN index, docstore, and index_to_docstore_id from disk. Args: folder_path: folder path to load index, docstore, and index_to_docstore_id from. embedding: Embeddings to use when generating queries index_name: for saving with a specific index file name allow_dangerous_deserialization: whether to allow deserialization of the data which involves loading a pickle file. Pickle files can be modified by malicious actors to deliver a malicious payload that results in execution of arbitrary code on your machine. aBThe de-serialization relies loading a pickle file. Pickle files can be modified to deliver a malicious payload that results in execution of arbitrary code on your machine.You will need to set `allow_dangerous_deserialization` to `True` to enable deserialization. If you do this, make sure that you trust the source of the data. For example, if you are loading a file that you created, and know that no one else has modified the file, then this is safe to do. Do not set this to `True` if you are loading a file from an untrusted source (e.g., some random site on the internet.).rrTrrrrbN) r4rrrrrrZ load_searcherrerrload) rxrr rrr0rrrr!rr"r#rrr load_locals$ zScaNN.load_localzCallable[[float], float]rcCs@|jdur|jS|jtjkr"|jS|jtjkr4|jStddS)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. NzJUnknown distance strategy, must be cosine, max_inner_product, or euclidean)r(r&rr[Z%_max_inner_product_relevance_score_fnrqZ_euclidean_relevance_score_fnr4)r*rrr_select_relevance_score_fns   z ScaNN._select_relevance_score_fnc sn|dd|dur$td|j|f|||d|}fdd|D}durjfdd|D}|S)z?Return docs and their similarity scores on a scale from 0 to 1.rPNzLnormalize_score_fn must be provided to ScaNN constructor to normalize scores)rBrCrDcsg|]\}}||fqSrr)rGrMZscore)r$rrrS,szBScaNN._similarity_search_with_relevance_scores..cs g|]\}}|kr||fqSrrrQ)rPrrrS0s)poprr4rh)r*rfrBrCrDr0rkZdocs_and_rel_scoresr)r$rPr(_similarity_search_with_relevance_scoress.   z.ScaNN._similarity_search_with_relevance_scores)NN)NN)NN)N)r?Nr@)r?Nr@)r?Nr@)r?Nr@)NNF)NN)NN)r!)r!)r?Nr@)__name__ __module__ __qualname____doc__rrqr+r9r:r=r>rdrhrlrm classmethodrzr{r~rrrrrrrrr!spC"  ;#&5r)# __future__rr\rropathlibrtypingrrrrrr r numpyrZlangchain_core.documentsr Zlangchain_core.embeddingsr Zlangchain_core.utilsrZlangchain_core.vectorstoresrZ!langchain_community.docstore.baserrZ&langchain_community.docstore.in_memoryrZ&langchain_community.vectorstores.utilsrrrrrrrrs  $