U
    hg                  	   @  s  d dl mZ d dlZd dlZd dlmZmZmZmZm	Z	m
Z
mZmZ d dlmZm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mZmZ dd	lmZ dd
lm Z  ddl!m"Z"m#Z# ddl$m%Z% ddl&m'Z'm(Z( e)ddddddddgZ*e	g e+f Z,e	g df Z-edeej.ej/f dZ0edee"e e#e f dZ1d2dddgZ3G dd de%Z4G d d! d!e(e0e1f Z5G d"d# d#e5ej.e"e f eZ6G d$d% d%e5ej/e#e f eZ7dS )&    )annotationsN)AnyUnionMappingTypeVarCallable	Awaitablecastoverload)Selfoverride   )	NOT_GIVENOmitQueryTimeoutNotGiven)is_given
is_mapping)OpenAIAsyncOpenAI)
model_copy)FinalRequestOptions)StreamAsyncStream)OpenAIError)DEFAULT_MAX_RETRIES
BaseClientz/completionsz/chat/completionsz/embeddingsz/audio/transcriptionsz/audio/translationsz/audio/speechz/images/generationsz/images/editszstr | Awaitable[str]_HttpxClientT)bound_DefaultStreamT <zmissing API key>c                      s"   e Zd Zdd fddZ  ZS )MutuallyExclusiveAuthErrorNonereturnc                   s   t  d d S )NzThe `api_key`, `azure_ad_token` and `azure_ad_token_provider` arguments are mutually exclusive; Only one can be passed at a time)super__init__)self	__class__ 4/tmp/pip-unpacked-wheel-dms8d3uv/openai/lib/azure.pyr)   .   s    z#MutuallyExclusiveAuthError.__init__)__name__
__module____qualname__r)   __classcell__r-   r-   r+   r.   r$   -   s   r$   c                      sZ   e Zd ZU ded< ded< eddddd	d
 fddZeddd fddZ  ZS )BaseAzureClientzhttpx.URL | None_azure_endpoint
str | None_azure_deploymentr   retries_takenr   intzhttpx.Request)optionsr8   r'   c                  sZ   |j tkrJt|jrJ|jd}|d k	rJdt| jjkrJd| |j  |_ t j	||dS )Nmodelz/deploymentsz/deployments/r7   )
url_deployments_endpointsr   Z	json_datagetstrbase_urlpathr(   _build_request)r*   r:   r8   r;   r+   r-   r.   rB   8   s
    zBaseAzureClient._build_requestr?   z	httpx.URL)r<   r'   c                   s`   | j rT| jrT|tkrTt|}|jrP| jjdd |jd }| jj	|dS |S t
 |S )zAdjust the URL if the client was configured with an Azure endpoint + deployment
        and the API feature being called is **not** a deployments-based endpoint
        (i.e. requires /deployments/deployment-name in the URL path).
           /s   /openai/raw_path)r6   r4   r=   httpxURLZis_relative_urlrE   rstriplstrip	copy_withr(   _prepare_url)r*   r<   Z	merge_urlmerge_raw_pathr+   r-   r.   rK   F   s    
zBaseAzureClient._prepare_url)r/   r0   r1   __annotations__r   rB   rK   r2   r-   r-   r+   r.   r3   4   s   
r3   c                      s  e Zd Zeddddddddeeddddddddddddddd	d
ddddddddZeddddddddeedddddddddddddd	d
ddddddddZedddddddeedddddddddddddd	d
ddddddddZdddddddddddeeddddddddddddddddd	d
dddddd fddZedddddddddededdddi ddddddddddd	ddddddddd fd d!ZeZ	dd"d#d$Z
ed%d%d&d'd(Zdd)d*d+d,d-Z  ZS ).AzureOpenAINF)azure_deploymentapi_versionapi_keyazure_ad_tokenazure_ad_token_providerorganizationwebhook_secretwebsocket_base_urltimeoutmax_retriesdefault_headersdefault_queryhttp_client_strict_response_validationr?   r5   zstr | Callable[[], str] | NonezAzureADTokenProvider | Nonestr | httpx.URL | None!float | Timeout | None | NotGivenr9   Mapping[str, str] | NoneMapping[str, object] | Nonezhttpx.Client | Noneboolr%   )azure_endpointrO   rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r'   c                C  s   d S Nr-   )r*   rb   rO   rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r-   r-   r.   r)   Z   s    zAzureOpenAI.__init__)rO   rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r'   c                C  s   d S rc   r-   )r*   rO   rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r-   r-   r.   r)   o   s    )rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   )r@   rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r'   c                C  s   d S rc   r-   )r*   r@   rP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r-   r-   r.   r)      s    )rP   rb   rO   rQ   rR   rS   rT   projectrU   rV   r@   rW   rX   rY   rZ   r[   r\   )rP   rb   rO   rQ   rR   rS   rT   rd   rU   rV   r@   rW   rX   rY   rZ   r[   r\   r'   c                  sd  |dkrt jd}|dkr(t jd}|dkrH|dkrH|dkrHtd|dkr\t jd}|dkrltd|dkr~d|i}n|d|i}|dkr|dkrt jd}|dkrtd	|dk	r|d
 d| }q|d
 d}n|dk	rtd|dkrt}t j||||	|||||||
|d || _	|| _
|| _|rD|nd| _|rZt|nd| _dS )aa  Construct a new synchronous azure openai client instance.

        This automatically infers the following arguments from their corresponding environment variables if they are not provided:
        - `api_key` from `AZURE_OPENAI_API_KEY`
        - `organization` from `OPENAI_ORG_ID`
        - `project` from `OPENAI_PROJECT_ID`
        - `azure_ad_token` from `AZURE_OPENAI_AD_TOKEN`
        - `api_version` from `OPENAI_API_VERSION`
        - `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`

        Args:
            azure_endpoint: Your Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`

            azure_ad_token: Your Azure Active Directory token, https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id

            azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on every request.

            azure_deployment: A model deployment, if given with `azure_endpoint`, sets the base client URL to include `/deployments/{azure_deployment}`.
                Not supported with Assistants APIs.
        NAZURE_OPENAI_API_KEYAZURE_OPENAI_AD_TOKENMissing credentials. Please pass one of `api_key`, `azure_ad_token`, `azure_ad_token_provider`, or the `AZURE_OPENAI_API_KEY` or `AZURE_OPENAI_AD_TOKEN` environment variables.OPENAI_API_VERSION_Must provide either the `api_version` argument or the `OPENAI_API_VERSION` environment variableapi-versionAZURE_OPENAI_ENDPOINTuMust provide one of the `base_url` or `azure_endpoint` arguments, or the `AZURE_OPENAI_ENDPOINT` environment variable//openai/deployments//openai2base_url and azure_endpoint are mutually exclusiverQ   rT   rd   rU   r@   rW   rX   rY   rZ   r[   rV   r\   osenvironr>   r   
ValueErrorrH   API_KEY_SENTINELr(   r)   _api_version_azure_ad_token_azure_ad_token_providerr6   rF   rG   r4   )r*   rP   rb   rO   rQ   rR   rS   rT   rd   rU   rV   r@   rW   rX   rY   rZ   r[   r\   r+   r-   r.   r)      sd    )

rQ   rT   rd   rU   rV   rP   rR   rS   r@   rW   r[   rX   rY   set_default_headersrZ   set_default_query_extra_kwargsint | NotGivenMapping[str, Any]r   rQ   rT   rd   rU   rV   rP   rR   rS   r@   rW   r[   rX   rY   r{   rZ   r|   r}   r'   c                  sF   t  j||||||	|
|||||||p(| j|p0| j|p8| jd|dS z~
        Create a new client instance re-using the same options given to the current client with optional overriding.
        )rP   rR   rS   )rQ   rT   rd   rU   rV   r@   rW   r[   rX   rY   r{   rZ   r|   r}   r(   copyrw   rx   ry   r*   rQ   rT   rd   rU   rV   rP   rR   rS   r@   rW   r[   rX   rY   r{   rZ   r|   r}   r+   r-   r.   r     s*    zAzureOpenAI.copyr&   c                 C  sH   | j d k	r| j S | j}|d k	rD| }|r2t|ts@td| |S d S NzOExpected `azure_ad_token_provider` argument to return a string but it returned )rx   ry   
isinstancer?   ru   r*   Zprovidertokenr-   r-   r.   _get_azure_ad_token2  s    
zAzureOpenAI._get_azure_ad_tokenr   r:   r'   c                 C  s   t |jr|jni }t|}||_|  }|d k	rR|dd kr~d| |d< n,| jtk	rv|dd kr~| j|d< ntd|S NAuthorizationBearer api-keyzUnable to handle authr   headersr   r   r>   rQ   rv   ru   r*   r:   r   rR   r-   r-   r.   _prepare_optionsA  s    
zAzureOpenAI._prepare_optionsr    tuple[httpx.URL, dict[str, str]]r;   extra_queryr'   c           
      C  s   i }|| j | jp|d}| jr6| jdkr6d| ji}n|  }|rPdd| i}| jd k	rt| j}|jdd }|j	|d}n| 
d	}|j	d
d}|j	|d}	|	|fS N)rj   Z
deploymentz<missing API key>r   r   r   rC   s	   /realtimerD   z	/realtimewss)scheme)paramsrw   r6   rQ   r   rV   rF   rG   rE   rH   rJ   rK   
r*   r;   r   Zauth_headersqueryr   r@   rL   Zrealtime_urlr<   r-   r-   r.   _configure_realtimeU  s$    

zAzureOpenAI._configure_realtimer/   r0   r1   r
   r   r   r)   r   r   Zwith_optionsr   r   r   r2   r-   r-   r+   r.   rN   Y   s   2008j:.rN   c                      s  e Zd Zedddddddddeedddddddddddddddd	d
ddddddddZedddddddddeeddddddddddddddd	d
ddddddddZeddddddddeeddddddddddddddd	d
ddddddddZdddddddddddeeddddddddddddddddd	d
dddddd fddZedddddddddededdddi ddddddddddd	ddddddddd fd d!ZeZ	dd"d#d$Z
ed%d%d&d'd(Zdd)d*d+d,d-Z  ZS ).AsyncAzureOpenAINF)rO   rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   r?   r5   z)str | Callable[[], Awaitable[str]] | Nonez AsyncAzureADTokenProvider | Noner]   r^   r9   r_   r`   zhttpx.AsyncClient | Nonera   r%   )rb   rO   rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   r'   c                C  s   d S rc   r-   )r*   rb   rO   rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   r-   r-   r.   r)   p  s    zAsyncAzureOpenAI.__init__)rO   rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   r'   c                C  s   d S rc   r-   )r*   rO   rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   r-   r-   r.   r)     s    )rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   )r@   rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   r'   c                C  s   d S rc   r-   )r*   r@   rP   rQ   rR   rS   rT   rd   rU   rV   rW   rX   rY   rZ   r[   r\   r-   r-   r.   r)     s    )rb   rO   rP   rQ   rR   rS   rT   rd   rU   r@   rV   rW   rX   rY   rZ   r[   r\   )rb   rO   rP   rQ   rR   rS   rT   rd   rU   r@   rV   rW   rX   rY   rZ   r[   r\   r'   c                  sd  |dkrt jd}|dkr(t jd}|dkrH|dkrH|dkrHtd|dkr\t jd}|dkrltd|dkr~d|i}n|d|i}|
dkr|dkrt jd}|dkrtd	|dk	r|d
 d| }
q|d
 d}
n|dk	rtd|dkrt}t j||||	|
|||||||d || _	|| _
|| _|rD|nd| _|rZt|nd| _dS )ab  Construct a new asynchronous azure openai client instance.

        This automatically infers the following arguments from their corresponding environment variables if they are not provided:
        - `api_key` from `AZURE_OPENAI_API_KEY`
        - `organization` from `OPENAI_ORG_ID`
        - `project` from `OPENAI_PROJECT_ID`
        - `azure_ad_token` from `AZURE_OPENAI_AD_TOKEN`
        - `api_version` from `OPENAI_API_VERSION`
        - `azure_endpoint` from `AZURE_OPENAI_ENDPOINT`

        Args:
            azure_endpoint: Your Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`

            azure_ad_token: Your Azure Active Directory token, https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id

            azure_ad_token_provider: A function that returns an Azure Active Directory token, will be invoked on every request.

            azure_deployment: A model deployment, if given with `azure_endpoint`, sets the base client URL to include `/deployments/{azure_deployment}`.
                Not supported with Assistants APIs.
        Nre   rf   rg   rh   ri   rj   rk   rl   rm   rn   ro   rp   rq   rr   )r*   rb   rO   rP   rQ   rR   rS   rT   rd   rU   r@   rV   rW   rX   rY   rZ   r[   r\   r+   r-   r.   r)     sd    )

rz   r~   r   r   r   c                  sF   t  j||||||	|
|||||||p(| j|p0| j|p8| jd|dS r   r   r   r+   r-   r.   r     s*    zAsyncAzureOpenAI.copyr&   c                   sf   | j d k	r| j S | j}|d k	rb| }t|r8|I d H }|rLttt|tsZtd| t|S d S r   )	rx   ry   inspectisawaitabler   r	   r   r?   ru   r   r-   r-   r.   r   K  s    


z$AsyncAzureOpenAI._get_azure_ad_tokenr   r   c                   s   t |jr|jni }t|}||_|  I d H }|d k	rX|dd krd| |d< n,| jtk	r||dd kr| j|d< ntd|S r   r   r   r-   r-   r.   r   \  s    
z!AsyncAzureOpenAI._prepare_optionsr   r   r   c           
        s   i }|| j | jp|d}| jr6| jdkr6d| ji}n |  I d H }|rVdd| i}| jd k	rt| j}|jdd }|j	|d}n| 
d	}|j	d
d}|j	|d}	|	|fS r   r   r   r-   r-   r.   r   p  s$    

z$AsyncAzureOpenAI._configure_realtimer   r-   r-   r+   r.   r   o  s   4228j:.r   )8
__future__r   rs   r   typingr   r   r   r   r   r   r	   r
   Ztyping_extensionsr   r   rF   _typesr   r   r   r   r   _utilsr   r   Z_clientr   r   _compatr   Z_modelsr   Z
_streamingr   r   _exceptionsr   Z_base_clientr   r   setr=   r?   ZAzureADTokenProviderZAsyncAzureADTokenProviderZClientZAsyncClientr   r    joinrv   r$   r3   rN   r   r-   r-   r-   r.   <module>   sF   (%   