
    Ng*
                         d dl mZ d dlmZmZ d dlZd dlmZ d dl	m
Z
 d dlmZmZ d dlmZ erd dlmZ  G d	 d
e          Ze G d de                      ZdS )    )	dataclass)TYPE_CHECKINGListN)	SecretStr)Element)BaseEmbeddingEncoderEmbeddingConfig)requires_dependenciesBedrockEmbeddingsc                   d    e Zd ZU eed<   eed<   dZeed<    eg dd          dd
            ZdS )BedrockEmbeddingConfigaws_access_key_idaws_secret_access_keyz	us-west-2region_name)boto3numpylangchain_communitybedrock)extrasreturnr   c                     dd l }ddlm} |                    d| j                                        | j                                        | j                  } ||          }|S )Nr   r   zbedrock-runtime)service_namer   r   r   )client)r   langchain_community.embeddingsr   r   r   get_secret_valuer   r   )selfr   r   bedrock_runtimebedrock_clients        V/var/www/html/ai-engine/env/lib/python3.11/site-packages/unstructured/embed/bedrock.py
get_clientz!BedrockEmbeddingConfig.get_client   s     	DDDDDD,,*"4EEGG"&"<"M"M"O"O(	 ' 
 
 +*/BBB    N)r   r   )	__name__
__module____qualname__r   __annotations__r   strr
   r!    r"   r    r   r      s}             $$$$"K"""111    	   r"   r   c                       e Zd ZU eed<   dee         fdZd Zd Z	d Z
d Zdee         dee         fd	Zdee         fd
ZdS )BedrockEmbeddingEncoderconfigr   c                 .    |                      d          S )NQ)query)embed_queryr   s    r    get_exemplary_embeddingz/BedrockEmbeddingEncoder.get_exemplary_embedding.   s    c***r"   c                 .    |                                   d S N)
initializer0   s    r    __post_init__z%BedrockEmbeddingEncoder.__post_init__1   s    r"   c                 R    |                                  }t          j        |          S r3   )r1   npshaper   exemplary_embeddings     r    num_of_dimensionsz)BedrockEmbeddingEncoder.num_of_dimensions4   s%    "::<<x+,,,r"   c                     |                                  }t          j        t          j                            |          d          S )Ng      ?)r1   r7   iscloselinalgnormr9   s     r    is_unit_vectorz&BedrockEmbeddingEncoder.is_unit_vector8   s5    "::<<z")..)<==sCCCr"   c                     | j                                         }t          j        |                    |                    S r3   )r+   r!   r7   arrayr/   )r   r.   r   s      r    r/   z#BedrockEmbeddingEncoder.embed_query<   s3    //11x22599:::r"   elementsc                     | j                                         }|                    d |D                       }|                     ||          }|S )Nc                 ,    g | ]}t          |          S r(   )r'   ).0es     r    
<listcomp>z;BedrockEmbeddingEncoder.embed_documents.<locals>.<listcomp>B   s    4N4N4NSVV4N4N4Nr"   )r+   r!   embed_documents_add_embeddings_to_elements)r   rC   r   
embeddingselements_with_embeddingss        r    rI   z'BedrockEmbeddingEncoder.embed_documents@   sS    //11#334N4NX4N4N4NOO
#'#C#CHj#Y#Y ''r"   c                     t          |          t          |          k    sJ g }t          |          D ]'\  }}||         |_        |                    |           (|S r3   )len	enumeraterK   append)r   rC   rK   elements_w_embeddingielements         r    rJ   z3BedrockEmbeddingEncoder._add_embeddings_to_elementsF   sh    8}}J////!#H-- 	1 	1JAw!+AG ''0000r"   N)r#   r$   r%   r   r&   r   floatr1   r5   r;   r@   r/   r   rI   rJ   r(   r"   r    r*   r*   *   s         """"+e + + + +  - - -D D D; ; ;(W ($w- ( ( ( (4=      r"   r*   )dataclassesr   typingr   r   r   r7   pydanticr   unstructured.documents.elementsr   unstructured.embed.interfacesr   r	   unstructured.utilsr
   r   r   r   r*   r(   r"   r    <module>r[      s   ! ! ! ! ! ! & & & & & & & &                P O O O O O O O 4 4 4 4 4 4 A@@@@@@    _   2 ! ! ! ! !2 ! ! ! ! !r"   