
    i9H                        d dl Z d dlmZ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mZmZ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mZ ddlmZmZmZ 	 	 	 ddee   dede dedede!dedee"   de"de#fdZ$ G d de      Z%y)    N)AnyOptionalUnion)!update_headers_with_filtered_beta)Logging)AsyncHTTPHandlerHTTPHandler_get_httpx_clientget_async_httpx_client)ModelResponse)CustomStreamWrapper   )
BaseAWSLLMCredentials)BedrockError_get_all_bedrock_regions   )AWSEventStreamDecoderMockResponseIterator	make_callclientapi_baseheadersdatamodelmessageslogging_obj	json_modefake_streamstream_chunk_sizec
                 
   | 
t               } | j                  |||| |      }
|
j                  dk7  r.t        |
j                  t	        |
j                                     |r[t        j                         j                  ||
t        j                         d|i d||t        j                  
      }t        ||      }n.t        ||      }|j                  |
j                  |		            }|j                  |dd
d|i       |S )N)r   r   streamr      status_codemessageT )
r   responsemodel_responser"   r   optional_paramsapi_keyr   r   encoding)r)   r   )r   r   )
chunk_sizezfirst stream response receivedcomplete_input_dict)inputr+   original_responseadditional_args)r
   postr%   r   strreadlitellmAmazonConverseConfig_transform_responser   r,   r   r   
iter_bytes	post_call)r   r   r   r   r   r   r   r   r   r    r(   r)   completion_streamdecoders                 {/Users/manta/Documents/Projects/TheRoad-I1/.venv/lib/python3.12/site-packages/litellm/llms/bedrock/chat/converse_handler.pymake_sync_callr=      s.    ~"${{  H s" ,,c(--/6J
 	
  ((*>>"002#%% ? 
 	 "6)Y"
 (eyI#..+<=

 :.5	       c                       e Zd Zd fdZdi dddddfdededed	ed
eee	e
j                  f      dedededee   dedee   dee   dedefdZdi ddfdededed	ed
eee	e
j                  f      dedededededee   dee   deeef   fdZ	 	 	 	 ddededee   ded	edededed
eee	e
j                  f      dedee   deeeef      dee   fdZ xZS )BedrockConverseLLMreturnNc                 "    t         |           y )N)super__init__)self	__class__s    r<   rD   zBedrockConverseLLM.__init__Y   s    r>   F   r   r   r   r)   timeoutr*   litellm_paramscredentialsr   r   r   r+   r    c                   K   t        j                         j                  |||	|
|       d {   }t        j                  |      }| j                  ||
j                  d      xs d|||||      }|j                  |d||t        |j                        d       t        ||t        |j                        |||||||
       d {   }t        ||d	|
      }|S 7 7 w)Nr   r   r*   rI   r   aws_region_name	us-west-2rJ   rM   extra_headersendpoint_urlr   r   r+   r'   r.   r   r   r/   r+   r1   )
r   r   r   r   r   r   r   r   r   r    bedrockr:   r   custom_llm_providerr   )r5   r6   _async_transform_requestjsondumpsget_request_headersgetpre_calldictr   r   r   )rE   r   r   r   r)   rH   r,   r   r"   r*   rI   rJ   	logger_fnr   r   r   r   r+   r    request_datar   preppedr:   streaming_responses                           r<   async_streamingz"BedrockConverseLLM.async_streaming\   s    * %99;TT+) U 
 
 zz,'**#*../@AP[!! + 
 	'+$0 	 	
 #,)##/#
 
 1/ )#	
 "!a
<
s"   ,C#CBC#	C!
C#!C#r   r   c                   K   t        j                         j                  |||	|
|       d {   }t        j                  |      }| j                  ||
j                  d      xs d|||||      }|j                  |d|||j                  d       t        |j                        }|t        |t              sdi }|:t        |t              st        |t              rt        j                  |      }||d<   t!        |t         j"                  j$                  	      }n|}	 |j'                  ||||
       d {   }|j)                          t        j                         j7                  |||t        |t8              r|nd|d|||	|
      S 7 a7 V# t        j*                  $ r<}|j,                  j.                  }t1        ||j,                  j2                        d }~wt        j4                  $ r t1        dd      w xY ww)NrL   rM   rN   rO   r'   rR   rS   rH   )paramsllm_providerurlr   r   r   r$     Timeout error occurred.F
r   r(   r)   r"   r   r+   r   r   r*   r,   )r5   r6   rW   rX   rY   rZ   r[   r\   r   r]   
isinstancer   floatinthttpxTimeoutr   LlmProvidersBEDROCKr2   raise_for_statusHTTPStatusErrorr(   r%   r   textTimeoutExceptionr7   bool)rE   r   r   r   r)   rH   r,   r   r"   r*   rI   rJ   r^   r   r   r+   r_   r   r`   _paramsr(   err
error_codes                          r<   async_completionz#BedrockConverseLLM.async_completion   s    $ %99;TT+) U 
 
 zz,'**#*../@AP[!! + 
 	'+$"?? 	 	
 w'>F4D!EG"gu-GS1I#mmG4G%,	"+W-A-A-I-IF F	S#[['	 )  H %%' ++-AA)'565#+ B 
 	
s
X $$ 	R11J:s||?P?PQQ%% 	S38QRR	SsG   ,HFC3H#F ;F<F A HF H'7G$HHcustom_prompt_dictacompletionrP   c                 	   |j                  dd       }|j                  dd      }|j                  dd       }|j                  dd      }|j                  dd      }|| j                  |      }n|}|}d	D ]#  }|j                  |      s|t	        |      d  } n d }|j                  d
d      d   }|t               v rd
|v r|}|j                  d
d      d   }|}dD ](  }|j                  |      s|j                  |dd      } n | j                  |      }|	d|vr||d<   t        j                         j                  |||d      }| j                  |||      }|j                  dd       }|j                  dd       }|j                  dd       }|j                  dd       } |j                  dd       }!|j                  dd       }"|j                  dd       }#|j                  dd       }$|j                  dd       }%|j                  dd       }&|j                  dd        ||d<   | j                  |||||!|"| |$|%|&
      }'| j                  ||#|      \  }(})||du r|s|( d | d!}(|) d | d!})n|( d | d"}(|) d | d"})d#d$i}*|d#d$i|}*t        |*d%&      }*|	rYt        |t               rd }|du r#| j#                  |||)||||d|||*|
||||'||'      S | j%                  |||)||||||||*|
||'|(      S t        j                         j'                  |||||)      }+t)        j*                  |+      },| j-                  |'|||)|,|*|*      }-|j/                  |d|,|)|-j0                  d+,       |t        |t2              rJi }.|
:t        |
t4              st        |
t6              rt9        j:                  |
      }
|
|.d-<   t=        |.      }n|}|H|du rDt?        |t        |t               r|nd |)|-j0                  |,||||||.
      }/tA        |/|d|/      }0|0S 	 |jC                  |)|-j0                  |,|0      }1|1jE                          t        j                         jS                  ||1|t        |tT              r|nd|d|,|||4
      S # t8        jF                  $ r<}2|2jH                  jJ                  }3tM        |3|2jH                  jN                  1      d }2~2wt8        jP                  $ r tM        d2d31      w xY w)5Nr"   r    rG   model_idr   Fr   )r~   )zbedrock/converse/zbedrock/z	converse//r   r   )znova-2/znova/r'   rM   rT   )r   r   r"   rV   )r*   r   r~   aws_secret_access_keyaws_access_key_idaws_session_tokenaws_role_nameaws_session_nameaws_profile_nameaws_bedrock_runtime_endpointaws_web_identity_tokenaws_sts_endpointaws_external_id)
r   r   r   rM   r   r   r   r   r   r   )r   r   rM   Tz/model/z/converse-streamz	/conversezContent-Typezapplication/jsonbedrock_converse)r   provider)r   r   r   r)   r,   r   r*   r"   rI   r^   r   rH   r   r   r   rJ   r+   r    )r   r   r   r)   r,   r   r*   r"   rI   r^   r   rH   r   rJ   r+   rL   rO   rR   rS   rH   )
r   r   r   r   r   r   r   r   r   r    rU   rf   r$   rh   ri   rj   )+popr[   encode_model_id
startswithlensplitr   replacer5   r6   should_fake_stream_get_aws_region_nameget_credentialsget_runtime_endpointr   rk   r	   rb   rz   _transform_requestrX   rY   rZ   r\   r   r   rl   rm   rn   ro   r
   r=   r   r2   rr   rs   r(   r%   r   rt   ru   r7   rv   )4rE   r   r   r   r{   r)   r,   r   r*   r|   rH   rI   r^   rP   r   r+   r"   r    unencoded_model_idr   r   modelId_model_for_id	_strippedrp_region_from_model_potential_region_nova_prefixrM   r   r   r   r   r   r   r   r   r   r   rJ   rQ   proxy_endpoint_urlr   _datar   r`   rw   r:   ra   r(   rx   ry   s4                                                       r<   
completionzBedrockConverseLLM.completion   sy   & !$$Xt4+//0CTJ,00TB%))-?#''U;	)**4F*GG "M%ID''+ )#b') 4I E 15 )Q 7 : $<$>>3)CS%6"%OOC3A6	 ) 4''5$1$9$9,A$NM !5 **M*BG #.%_<5G 12224GG# )	 H 
 33+' 4 
 !0 3 34KT R+//0CTJ+//0CTJ'++OTB*../A4H*../A4H'6':':*D(
$ "1!4!45Mt!T*../A4H)--.?F-t4  		
 $(#7#7/"7/+--'#9-+ $8 $
 ,0+D+D)E+ ,E ,
((
 6T>;*^77);KLL$6#7wwiGW!X*^77)9EL$6#7wwiy!Q "#56$%'9K]KG 4&8
 &+.~++%/#1% +$3#1'##!' + +#&7% ,  * ((!+-!' /-#' )  ( ,,.AA+)! B 
 zz% **#+'+ + 
 	'+."?? 	 	
 >Z0@AG"gu-GS1I#mmG4G%,	"&w/FF&D. . )j.M +!'#'"3!  "5"3$-'	" &%	S{{&'	 # H %%' ++-AA)'565#+ B 
 	
 $$ 	R11J:s||?P?PQQ%% 	S38QRR	Ss   /R	 	S77S$S7)rA   N)NNNN)__name__
__module____qualname__rD   r3   listr   r   r   rl   rn   ro   r]   r   r   rv   rm   r   rb   LiteLLMLoggingObjectrz   r	   r   __classcell__)rF   s   @r<   r@   r@   X   sy     -1!$)!%!%'E"E" E" 	E"
 &E" %u}} 456E" E" E" !E" )*E"  !E"" D>#E"$ #%E"& 'E"( 
)E"h -1!%!V
V
 V
 	V

 &V
 %u}} 456V
 *V
 V
 V
 !V
 V
 )*V
  #!V
" 
}11	2#V
J (,AE!%!E
E
 E
 3-	E

 !E
 &E
 *E
 E
 E
 %u}} 456E
 E
  ~E
 /<=>E
  #!E
r>   r@   )FFrG   )&rX   typingr   r   r   rn   r5   &litellm.anthropic_beta_headers_managerr   *litellm.litellm_core_utils.litellm_loggingr   r   &litellm.llms.custom_httpx.http_handlerr   r	   r
   r   litellm.types.utilsr   litellm.utilsr   base_aws_llmr   r   common_utilsr   r   invoke_handlerr   r   r   r3   r]   r   rv   rm   r=   r@    r>   r<   <module>r      s     ' '   W  . - 2 A R R !&!<[!<< < 	<
 < < &< ~< < <~h
 h
r>   