
    j$8                         d dl mZmZmZ d dlmZmZ d dlmZ d dl	m
Z
mZmZ ddlmZ ddlmZmZ  e       Z G d	 d
e      Zy)    )AnyCallableOptional)AsyncAzureOpenAIAzureOpenAI)prompt_factory)CustomStreamWrapperModelResponseTextCompletionResponse   )OpenAITextCompletionConfig   )AzureOpenAIErrorBaseAzureLLMc                   >    e Zd Zd fdZd Z	 	 	 ddedededee   ded	ed
edee   dee	   de	de
dee   fdZddi fdee   d	ededededededededee   defdZddi fdedee   d	ededededee   defdZddi fdedee   d	ededededee   defdZ xZS )AzureTextCompletionNc                 "    t         |           y )N)super__init__)self	__class__s    ~/Users/manta/Documents/Projects/TheRoad-I1/backend/.venv/lib/python3.12/site-packages/litellm/llms/azure/completion/handler.pyr   zAzureTextCompletion.__init__   s        c                 4    ddi}|||d<   |S |d| |d<   |S )Nzcontent-typezapplication/jsonzapi-keyBearer Authorization )r   api_keyazure_ad_tokenheaderss       r   validate_environmentz(AzureTextCompletion.validate_environment   sG    .
 !(GI  ')00@'AGO$r   modelmessagesmodel_responser   api_baseapi_versionapi_typer   azure_ad_token_providerprint_verboseacompletionr    c                 l   	 ||t        dd      |j                  dd      }t        ||d      }|(d|v r$| j                  |||||||||	||	      }d |d
|}n||d
|}|du rJ|j	                  dd      r| j                  ||||||||||
      S | j                  ||||||||||||      S d|v r!|d   du r| j                  |||||||||	      S |j                  ||||d|||d       t        |t              st        dd      | j                  |||||d|      }t        |t              st        dd       |j                  j                  j                  d i |d|i}|j!                         }|j#                         }|j%                  ||||||d       t&        j)                  t+        d i ||      S # t         $ r}|d }~wt,        $ rT}t/        |dd      }t/        |dd       }t/        |dd       }||rt/        |dd       }t        |t1        |      |      d }~ww xY w)!N  zMissing model or messagesstatus_codemessagemax_retriesr   
azure_text)r#   r"   custom_llm_providerzgateway.ai.cloudflare.com)r   r&   r%   r"   clientr0   timeoutr   r(   r*   litellm_params)r"   promptTstreamF)
logging_objr%   datar"   r   r&   r   r4   r3   r5   )r%   r9   r$   r   r&   r"   r   r4   r3   r8   r0   r5   )	r8   r%   r9   r"   r   r&   r   r4   r3   )r   r   )r    r&   r%   complete_input_dictinputr   additional_argsmax retries must be an int)r   r%   r&   r3   r5   	_is_asyncr"     .azure_client is not an instance of AzureOpenAIr4   )r    r&   r%   )r<   r   original_responser=   response_objectmodel_response_objectr.   r    responser.   r/   r    r   )r   popr   ,_init_azure_client_for_cloudflare_ai_gatewaygetasync_streamingr*   	streamingpre_call
isinstanceintget_azure_openai_clientr   completionswith_raw_responsecreateparse
model_dump	post_callopenai_text_completion_config%convert_to_chat_model_response_objectr   	Exceptiongetattrstr)r   r"   r#   r$   r   r%   r&   r'   r   r(   r)   r4   r8   optional_paramsr5   	logger_fnr*   r    r3   r0   r6   r9   azure_clientraw_responserF   stringified_responseer.   error_headerserror_responses                                 r   
completionzAzureTextCompletion.completion   sN   *R	} 0& #-H  *--mQ?K#!LF #(Cx(OJJ# +%! +##1,C +#1 K  "&K?K #$ & d""&&x7//$/!)!# '$/'5 '%'5 0    ++!)!'5 '$/#'5 '%$/$/'5 ,   _,1Jd1R~~ +%# +#1#! & 
 
 $$ # (/.<$ (3$,/3% %  "+s3*$'1M   $;;#% +!#1#  <   ",<*$' P 
  Q|77IIPP   $+  (--/'/':':'<$%% #&:#*'2$,%	 & 	 2WW(>(VAU(V.< X    	G 	!!]C8K#Ay$7M$Q
D9N$ '	4 H"'Q 	s8   BG G 2$G C0G 	H3GH3AH..H3r9   r4   r8   r0   r5   c           	        K   d }	 | j                  ||||d||      }t        |t              st        dd      |j	                  |d   |j
                  dd|j
                   i|j                  j                  d|d	
        |j                  j                  j                  di |d|i d {   }|j                         }t        j                  |j                         |      S 7 9# t        $ r}|d }~wt        $ rT}t!        |dd      }t!        |dd       }t!        |dd       }||rt!        |dd       }t        |t#        |      |      d }~ww xY ww)NTr&   r%   r   r"   r?   r3   r5   r@   3azure_client is not an instance of AsyncAzureOpenAIr-   r6   r   r   r    r%   r*   r:   r;   r4   rC   r.   r    rF   rG   r   )rP   rN   r   r   rM   r   	_base_url_uri_referencerQ   rR   rS   rT   rW   rX   rU   rY   rZ   r[   )r   r   r&   r"   r%   r9   r4   r$   r8   r0   r   r3   r5   rF   r^   r_   ra   r.   rb   rc   s                       r   r*   zAzureTextCompletion.acompletion   s     /	  77'!- 8 L l,<=& #Q    8n$,, /7<;O;O:P1QR , 6 6 E E#'+/	! ! 	 "S!9!9!K!K!R!R "" '" L $))+H0VV ( 3 3 5&4 W  	   	G 	!!]C8K#Ay$7M$Q
D9N$ '	4 H"'Q 	sH   EB+C- 1C+28C- *E+C- -	E6C88EAEEEc           	         |j                  dd      }t        |t              st        dd      | j	                  ||||d|	|
      }t        |t
              st        dd	      |j                  |d
   |j                  dd|j                   i|j                  j                  d|d        |j                  j                  j                  di |d|i}|j                         }t        ||d|      }|S )Nr0   r   r,   r>   r-   Frf   r@   rA   r6   r   r   Trh   r;   r4   r1   completion_streamr"   r2   r8   r   )rH   rN   rO   r   rP   r   rM   r   ri   rj   rQ   rR   rS   rT   r	   )r   r8   r%   r   r&   r9   r"   r4   r   r3   r5   r0   r^   r_   rF   streamwrappers                   r   rL   zAzureTextCompletion.streaming  s0    hh}a0+s+")E  33#) 4 
 ,4"H  	x. ((+w|7K7K6L-MN(22AA#'+	 	 		
 I|//AAHH 

#
  %%'+& ,#	
 r   c           	      r  K   	 | j                  ||||d|	|
      }t        |t              st        dd      |j	                  |d   |j
                  dd|j
                   i|j                  j                  d|d	
        |j                  j                  j                  di |d|i d {   }|j                         }t        ||d|      }|S 7 %# t        $ rT}t        |dd      }t        |dd       }t        |dd       }||rt        |dd       }t        |t        |      |      d }~ww xY ww)NTrf   r@   rg   r-   r6   r   r   rh   r;   r4   r1   rl   r.   r    rF   rG   r   )rP   rN   r   r   rM   r   ri   rj   rQ   rR   rS   rT   r	   rY   rZ   r[   )r   r8   r%   r   r&   r9   r"   r4   r   r3   r5   r^   r_   rF   rn   ra   r.   rb   rc   s                      r   rK   z#AzureTextCompletion.async_streaming?  s|    /	77'!- 8 L l,<=& #Q 
   8n$,, /7<;O;O:P1QR , 6 6 E E#'+/	! ! 	 "S!9!9!K!K!R!R "" '" L $))+H/"*$0'	M !   	!!]C8K#Ay$7M$Q
D9N$ '	4 H"'Q 	s<   D7B+C /C0$C D7C 	D4 AD//D44D7)returnN)FNN)__name__
__module____qualname__r   r!   r[   listr
   r   r   booldictrd   r   rO   r*   rL   rK   __classcell__)r   s   @r   r   r      s   6 ""&'gg g &	g
 #g g g g !g "*(!3g  g" #g$ $%gh )-!?#? ? 	?
 ? ? ? &? ? ? !? ?T )-!7 7 #	7
 7 7 7 7 !7 7D )-!< < #	<
 < < < < !< <r   r   N)typingr   r   r   openair   r   3litellm.litellm_core_utils.prompt_templates.factoryr   litellm.utilsr	   r
   r    openai.completion.transformationr   common_utilsr   r   rW   r   r   r   r   <module>r~      s4    * * 0 N T T J 9 : < m, mr   