o
    Zh0z                     @   s  d Z ddlZddlZddlZddlmZ ddlZddlZddlZddl	Z	ddl
Z
ddlZddlmZmZmZmZmZ ddlmZmZ ddlZddlZddlmZ ddlmZ ddlZddlmZmZmZ d	d
l m!Z! d	dl m"Z" d	dl m#Z# d	dl$m%Z%m&Z&m'Z' e	(dZ)dZ*de+e,e,f ddfddZ-de&de+e,ef de&fddZ.de,de,de,fddZ/dee,df deee,f fddZ0dedefd d!Z1eG d"d# d#Z2G d$d% d%e!jZ3G d&d' d'Z4G d(d) d)ej5Z6G d*d+ d+ej7Z8G d,d- d-Z9dS ).zBase client for calling HTTP APIs sending and receiving JSON.

The BaseApiClient is intended to be a private module and is subject to change.
    N)	dataclass)AnyAsyncIteratorOptionalTupleUnion)urlparse
urlunparse)Credentials)Request)	BaseModelFieldValidationError   )_common)errors)version)HttpOptionsHttpOptionsDictHttpOptionsOrDictzgoogle_genai._api_clienti   headersreturnc                 C   s   dt j }dtj  d  }| d| }d| v r,|| d vr,| d  d| 7  < nd| vr4|| d< d| v rK|| d vrK| d  d| 7  < dS d| vrU|| d< dS dS )z1Appends the telemetry header to the headers dict.zgoogle-genai-sdk/z
gl-python/r    z
user-agentzx-goog-api-clientN)r   __version__syssplit)r   Zlibrary_labelZlanguage_labelZversion_header_value r   O/var/www/html/lang_env/lib/python3.10/site-packages/google/genai/_api_client.py_append_library_version_headers1   s   r   optionspatch_optionsc                 C   s   t  }||  | D ]/\}}t|tr3t||d tr3i ||< || | |  || | q|d ur;|||< q|d rFt|d  |S )Nr   )r   updateitems
isinstancedictgetr   )r   r    Zcopy_optionZ	patch_keyZpatch_valuer   r   r   _patch_http_optionsF   s"   
r&   base_urlpathc                 C   sZ   t | }|jdr|jd d n|j}|dr|dd  n|}t|j|d | dS )N/r   )r(   )r   r(   endswith
startswithr	   _replace)r'   r(   Zparsed_base	base_pathr   r   r   _join_url_path^   s   
r/   projectc                 C   s0   t jjdgd\}}| s|} | std|| fS )z-Loads google auth credentials and project id.z.https://www.googleapis.com/auth/cloud-platform)Zscopesz@Could not resolve project using application default credentials.)googleauthdefault
ValueError)r0   credentialsZloaded_project_idr   r   r   
_load_authi   s   
r6   r5   c                 C   s   |  t  | S N)refreshr   )r5   r   r   r   _refresh_authz   s   r9   c                   @   sV   e Zd ZU eeef ed< eed< eed< eeeef ef ed< dZ	e
e ed< dS )HttpRequestr   urlmethoddataNtimeout)__name__
__module____qualname__r$   str__annotations__r   objectbytesr>   r   floatr   r   r   r   r:      s   
 r:   c                   @   sF   e Zd ZU edddZeeeef  ed< edddZ	ee
 ed< dS )BaseResponseNz!The http headers of the response.)r3   descriptionhttp_headersz!The json payload of the response.json_payload)r?   r@   rA   r   rI   r   r$   rB   rC   rJ   r   r   r   r   r   rG      s   
 rG   c                	   @   s   e Zd Z		ddeeeef ejf deeef deee	f fddZ
dd Zd	d
 ZedefddZdd Zdee fddZdd Zdeeef fddZdS )HttpResponseNr   response_streambyte_streamc                 C   s   d| _ || _|| _|| _d S )N   )status_coder   rL   rM   )selfr   rL   rM   r   r   r   __init__   s   
zHttpResponse.__init__c                 C   s   |   | _| S r7   )async_segmentssegment_iteratorrP   r   r   r   	__aiter__   s   
zHttpResponse.__aiter__c                    s(   z	| j  I d H W S  ty   tw r7   )rS   	__anext__StopIterationStopAsyncIterationrT   r   r   r   rV      s   zHttpResponse.__anext__r   c                 C   s   | j d sdS t| j d S )Nr    )rL   jsonloadsrT   r   r   r   rZ      s   
zHttpResponse.jsonc                 c   s    t | jtr| jD ]}|rt|ni V  q
d S | jd u r%g E d H  d S | j D ]!}|rKt |ts8|d}|drE|t	dd  }t|V  q*d S )Nutf-8data: )
r#   rL   listrZ   r[   
iter_linesrB   decoder,   len)rP   chunkr   r   r   segments   s    




zHttpResponse.segmentsc                 C  s   t | jtr| jD ]}|rt|ni V  q
d S | jd u r-g 2 z	3 d H W }|V  q 6 d S t| jdra| j 2 z%3 d H W }|r]t |tsJ|d}|	drW|t
dd  }t|V  q86 d S td)Naiter_linesr\   r]   z!Error parsing streaming response.)r#   rL   r^   rZ   r[   hasattrrd   rB   r`   r,   ra   r4   )rP   rb   cr   r   r   rR      s,   




zHttpResponse.async_segmentsc                 c   s>    t | jtr| jE d H  d S | jd u rg E d H  d S td)Nz8Byte segments are not supported for streaming responses.)r#   rM   r^   r4   rT   r   r   r   byte_segments   s   
zHttpResponse.byte_segmentsresponse_payloadc                 C   s0   t | d t| D ]}tt| |||< q	d S )NrS   )delattrdircopydeepcopygetattr)rP   rh   	attributer   r   r   _copy_to_dict   s   
zHttpResponse._copy_to_dict)NN)r?   r@   rA   r   r$   rB   httpxZHeadersr   rE   rQ   rU   rV   propertyrZ   rc   r   rR   rg   rD   ro   r   r   r   r   rK      s$    


rK   c                       4   e Zd ZdZdeddf fddZd	ddZ  ZS )
SyncHttpxClientzSync httpx client.kwargsr   Nc                    "   | dd t jdi | dS zInitializes the httpx client.Zfollow_redirectsTNr   
setdefaultsuperrQ   rP   rt   	__class__r   r   rQ         zSyncHttpxClient.__init__c                 C   s.   | j rdS z|   W dS  ty   Y dS w )zCloses the httpx client.N)	is_closedclose	ExceptionrT   r   r   r   __del__   s   zSyncHttpxClient.__del__r   Nr?   r@   rA   __doc__r   rQ   r   __classcell__r   r   r{   r   rs          rs   c                       rr   )
AsyncHttpxClientzAsync httpx client.rt   r   Nc                    ru   rv   rw   rz   r{   r   r   rQ     r}   zAsyncHttpxClient.__init__c                 C   s8   | j rd S zt |   W d S  ty   Y d S w r7   )r~   asyncioget_running_loopcreate_taskacloser   rT   r   r   r   r     s   zAsyncHttpxClient.__del__r   r   r   r   r{   r   r     r   r   c                   @   sx  e Zd ZdZ						d;dee dee deejj	j
 dee dee dee fd	d
Zdd ZdefddZdefddZ	d<dededeeef dee def
ddZ	d=dededefddZ	d=dedefddZdefddZ	d<dededeeef dee fd d!Z	d<dededeeef dee fd"d#Z	d<dededeeef dee deeef f
d$d%Z	d<dededeeef dee fd&d'Zd(eeejf d)ed*e deeef fd+d,Z!d-ejd)ed*e deeef fd.d/Z"defd0d1Z#d(eeejf d)ed*e deeef fd2d3Z$d-eeje%j&f d)ed*e deeef fd4d5Z'defd6d7Z(d8e)j*fd9d:Z+dS )>BaseApiClientz8Client for calling HTTP APIs sending and receiving JSON.Nvertexaiapi_keyr5   r0   locationhttp_optionsc              
   C   s  || _ | j d u rtjdd dv rd| _ |s|r |r td|r(|r(tdt|trJz	t	|
 }W n tyI } ztd| d }~ww t|trS|
 }tjdd }	tjd	d }
tjd
d }|pk|	| _|pp|
| _|pu|| _|| _t | _t | _| j r|rtd d | _n/|
s|	r|rtd d | _d | _n|s|r|rtd d | _n|
s|	r|rtd d | _| js| jstd d\}| _| js|| _| jr| js| jstd| js| jdkrd| jd< n
d| j d| jd< d| jd< n| jstdd| jd< d| jd< ddi| jd< | jr$| j| jd d< |r/t| j|| _nt| jd  t | _t | _d S )NZGOOGLE_GENAI_USE_VERTEXAI0)true1TzNProject/location and API key are mutually exclusive in the client initializer.zICredentials and API key are mutually exclusive in the client initializer.zInvalid http_options: ZGOOGLE_CLOUD_PROJECTZGOOGLE_CLOUD_LOCATIONZGOOGLE_API_KEYzoThe user provided Google Cloud credentials will take precedence over the API key from the environment variable.zrThe user provided Vertex AI API key will take precedence over the project/location from the environment variables.zqThe user provided project/location will take precedence over the Vertex AI API key from the environment variable.zyThe project/location from the environment variables will take precedence over the API key from the environment variables.r0   zIProject and location or API key must be set when using the Vertex AI API.globalz"https://aiplatform.googleapis.com/r'   zhttps://z-aiplatform.googleapis.com/Zv1beta1api_versionzMissing key inputs argument! To use the Google AI API,provide (`api_key`) arguments. To use the Google Cloud API, provide (`vertexai`, `project` & `location`) arguments.z*https://generativelanguage.googleapis.com/Zv1betazContent-Typezapplication/jsonr   zx-goog-api-key)r   osenvironr%   lowerr4   r#   r$   r   Zmodel_validate
model_dumpr   r0   r   r   _credentialsr   _http_optionsr   Lock
_auth_lockloggerinfor6   r&   r   rs   _httpx_clientr   _async_httpx_client)rP   r   r   r5   r0   r   r   Zvalidated_http_optionseZenv_projectZenv_locationZenv_api_keyr   r   r   rQ     s   	







zBaseApiClient.__init__c                 C   s   t | jd }|jdd S )Nr'   wss)scheme)r   r   r-   geturl)rP   Z	url_partsr   r   r   _websocket_base_url  s   z!BaseApiClient._websocket_base_urlr   c                 C   sd   | j st| jd\| _ }| js|| _| j r.| j js| j js"t| j  | j js*td| j jS td)z/Retrieves the access token for the credentials.r   0Could not resolve API token from the environment)r   r6   r0   expiredtokenr9   RuntimeErrorrP   r0   r   r   r   _access_token  s   
zBaseApiClient._access_tokenc              	      s  | j s8| j4 I dH " | j s#tjt| jdI dH \| _ }| js#|| _W d  I dH  n1 I dH s3w   Y  | j r~| j jsC| j jsr| j4 I dH  | j jsS| j js]tt| j I dH  W d  I dH  n1 I dH smw   Y  | j jszt	d| j jS t	d)z>Retrieves the access token for the credentials asynchronously.Nr   r   )
r   r   r   	to_threadr6   r0   r   r   r9   r   r   r   r   r   _async_access_token  s0   ((z!BaseApiClient._async_access_tokenhttp_methodr(   request_dictc                 C   s   dd |  D }|D ]}||= q|r(t|tr!t| j| }n
t| j|}n| j}d}| jr;|dkr;|dr;d}| jrU|dsU|sU| jsUd| j	 d| j
 d	| }t|d
d|ddd	 | }	|dd }
|
rs|
d }
nd }
t||	|d ||
dS )Nc                 S   s   g | ]	}| d r|qS )_)r,   ).0keyr   r   r   
<listcomp>  s    z0BaseApiClient._build_request.<locals>.<listcomp>Fr%   zpublishers/google/modelsTz	projects/z/locations/r)   r'   rY   r   r>   g     @@r   )r<   r;   r   r=   r>   )keysr#   r   r&   r   r   r   r,   r   r0   r   r/   r%   r:   )rP   r   r(   r   r   Zkeys_to_deleter   Zpatched_http_optionsZquery_vertex_base_modelsr;   Ztimeout_in_secondsr   r   r   _build_request  sZ   



zBaseApiClient._build_requestFhttp_requeststreamc                 C   s$  d }| j r,| js,d|   |jd< | jr | jjr | jj|jd< |jr)t|jnd }n|jrDt	|jt
sA|jr>t|jnd }n|j}|rp| jj|j|j||j|jd}| jj||d}tj| t|j|rk|S |jgS | jj|j|j|j||jd}tj| t|j|r|S |jgS NzBearer Authorizationzx-goog-user-project)r<   r;   contentr   r>   r   r<   r;   r   r   r>   )r   r   r   r   r   quota_project_idr=   rZ   dumpsr#   rE   r   build_requestr<   r;   r>   sendr   APIErrorraise_for_responserK   textrequestrP   r   r   r=   Zhttpx_requestresponser   r   r   _request  sP   

zBaseApiClient._requestc                    s8  d }| j r0| js0d|  I d H  |jd< | jr$| jjr$| jj|jd< |jr-t|jnd }n|jrHt	|jt
sE|jrBt|jnd }n|j}|rw| jj|j|j||j|jd}| jj||dI d H }tj| t|j|rr|S |jgS | jj|j|j|j||jdI d H }tj| t|j|r|S |jgS r   )r   r   r   r   r   r   r=   rZ   r   r#   rE   r   r   r<   r;   r>   r   r   r   r   rK   r   r   r   r   r   r   _async_request>  sX   

zBaseApiClient._async_requestc                 C   s.   t  }t| jtr| j | _|| j |S r7   )r   r#   r   r   r   r!   )rP   Zcopiedr   r   r   get_read_only_http_optionso  s
   z(BaseApiClient.get_read_only_http_optionsc                 C   s@   |  ||||}| j|dd}|j}|st|jdjddS |S )NFr   rI   TZby_alias)r   r   rZ   rG   r   r   )rP   r   r(   r   r   r   r   json_responser   r   r   r   v  s   zBaseApiClient.requestc                 c   s8    |  ||||}| j|dd}| D ]}|V  qd S )NTr   )r   r   rc   )rP   r   r(   r   r   r   Zsession_responserb   r   r   r   request_streamed  s   zBaseApiClient.request_streamedc                    sH   |  ||||}| j|ddI d H }|j}|s"t|jdjddS |S )NFr   r   r   Tr   )r   r   rZ   rG   r   r   )rP   r   r(   r   r   r   resultr   r   r   r   async_request  s   zBaseApiClient.async_requestc                    s8   |  ||||}| j|ddI d H   fdd}| S )NTr   c                    s     2 z	3 d H W } | V  q6 d S r7   r   )rb   r   r   r   async_generator  s   z=BaseApiClient.async_request_streamed.<locals>.async_generator)r   r   )rP   r   r(   r   r   r   r   r   r   r   async_request_streamed  s   z$BaseApiClient.async_request_streamed	file_path
upload_urlupload_sizec                 C   sV   t |tjr| |||S t|d}| |||W  d   S 1 s$w   Y  dS )a  Transfers a file to the given URL.

    Args:
      file_path: The full path to the file or a file like object inherited from
        io.BytesIO. If the local file path is not found, an error will be
        raised.
      upload_url: The URL to upload the file to.
      upload_size: The size of file content to be uploaded, this will have to
        match the size requested in the resumable upload request.

    returns:
          The response json object from the finalize request.
    rbN)r#   ioIOBase
_upload_fdopen)rP   r   r   r   filer   r   r   upload_file  s
   $zBaseApiClient.upload_filer   c           	      C   s   d}	 | t}d}|rt|}d}|| |kr|d7 }| jjd||t|t|d|d}||7 }|jd d	kr:n	||krBtd
q|jd dkrNtd| S )ay  Transfers a file to the given URL.

    Args:
      file: A file like object inherited from io.BytesIO.
      upload_url: The URL to upload the file to.
      upload_size: The size of file content to be uploaded, this will have to
        match the size requested in the resumable upload request.

    returns:
          The response json object from the finalize request.
    r   Tupload
, finalizePOSTzX-Goog-Upload-CommandzX-Goog-Upload-OffsetzContent-Length)r<   r;   r   r   x-goog-upload-statusactiveFAll content has been uploaded, but the upload status is not finalized.final6Failed to upload file: Upload status is not finalized.)	read
CHUNK_SIZEra   r   r   rB   r   r4   rZ   	rP   r   r   r   offsetZ
file_chunk
chunk_sizeZupload_commandr   r   r   r   r     s>   

zBaseApiClient._upload_fdc                 C   s   | j d|i |d}d}|jrt|jtst|j}n|j}| jj|j|j	|j
||jd}tj| t|j
| gdjd S zDownloads the file data.

    Args:
      path: The request path with query params.
      http_options: The http options to use for the request.

    returns:
          The file bytes
    r%   )r(   r   r   Nr   )rM   r   )r   r=   r#   rE   rZ   r   r   r   r<   r;   r   r>   r   r   r   rK   r   rM   rP   r(   r   r   r=   r   r   r   r   download_file  s*   
zBaseApiClient.download_filec              	      s   t |tjr| |||I dH S t|}|dI dH }|4 I dH  | |||I dH W  d  I dH  S 1 I dH s?w   Y  dS )a  Transfers a file asynchronously to the given URL.

    Args:
      file_path: The full path to the file. If the local file path is not found,
        an error will be raised.
      upload_url: The URL to upload the file to.
      upload_size: The size of file content to be uploaded, this will have to
        match the size requested in the resumable upload request.

    returns:
          The response json object from the finalize request.
    Nr   )r#   r   r   _async_upload_fdanyioPathr   )rP   r   r   r   r   fdr   r   r   async_upload_file%  s   
0zBaseApiClient.async_upload_filec           	         s   d}	 t |tjr|t}n|tI dH }d}|r t|}d}|| |kr,|d7 }| jjd|||t|t|ddI dH }||7 }|j	
d	d
krNn	||krVtdq|j	
d	dkrctd| S )a  Transfers a file asynchronously to the given URL.

    Args:
      file: A file like object inherited from io.BytesIO.
      upload_url: The URL to upload the file to.
      upload_size: The size of file content to be uploaded, this will have to
        match the size requested in the resumable upload request.

    returns:
          The response json object from the finalize request.
    r   TNr   r   r   r   )r<   r;   r   r   r   r   r   r   r   )r#   r   r   r   r   ra   r   r   rB   r   r%   r4   rZ   r   r   r   r   r   ?  sD   
zBaseApiClient._async_upload_fdc                    s   | j d|i |d}d}|jrt|jtst|j}n|j}| jj|j|j	|j
||jdI dH }tj| t|j
| gdjd S r   )r   r=   r#   rE   rZ   r   r   r   r<   r;   r   r>   r   r   r   rK   r   rM   r   r   r   r   async_download_filew  s,   
z!BaseApiClient.async_download_fileresponse_modelc                 C   s   d S r7   r   )rP   r   r   r   r   _verify_response  s   zBaseApiClient._verify_response)NNNNNNr7   )F),r?   r@   rA   r   r   boolrB   r1   r2   r5   r
   r   rQ   r   r   r   r$   rD   r:   r   rK   r   r   r   r   r   r   r   r   r   r   r   intr   r   r   r   r   Z	AsyncFiler   r   r   r   r   r   r   r   r   r     s
   
 $

@
1
1












2"



8%r   ):r   r   r   rk   dataclassesr   datetimer   rZ   loggingr   r   typingr   r   r   r   r   urllib.parser   r	   Zgoogle.authr1   Zgoogle.auth.credentialsr
   Zgoogle.auth.transport.requestsr   rp   Zpydanticr   r   r   rY   r   r   r   typesr   r   r   	getLoggerr   r   r$   rB   r   r&   r/   r6   r9   r:   rG   rK   ZClientrs   ZAsyncClientr   r   r   r   r   r   <module>   sV   


"

]