
    :jN                        d 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 ddl	m
Z
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 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!  ed      Z" G d de      Z# e!de#       y)z@
AWS Bedrock provider implementation with manual SigV4 signing.
    N)datetimetimezone)DictListOptional)DEFAULT_QUESTION)ErrorReason)CheckResult	Condition)
get_logger)trim   )http_error_messagehttp_error_statushttp_getrequest   )AIBaseProvider)register_providerproviderc                       e Zd ZdZdee   f fdZdedededefdZ	dedefd	Z
ded
edefdZdedefdZdededeeef   dedeeef   f
dZdededededef
dZdedededededefdZdededededef
dZ	 d-dedededededededeeef   fdZd-dededeeef   dedef
dZdededefd Zd!ed"edefd#Zd.d$ed%ee   dee   fd&Zd
edefd'Zd/d$ed(ed)ed*edef
d+Zd0d$ed(ed)edee   fd,Z xZS )1BedrockProviderz6AWS Bedrock provider with manual SigV4 authentication.
conditionsc           	      ~    | j                  |dddddd      }t        |   |d   |d   |d	   |d
   |d   |fi | y )Nbedrockz/https://bedrock-runtime.us-east-1.amazonaws.comz/model/{model-id}/invoke/foundation-modelsz)anthropic.claude-3-7-sonnet-20250219-v1:0)namebase_urlcompletion_path
model_pathdefault_modelr   r   r   r    r!   )extractsuper__init__)selfr   kwargsconfig	__class__s       >/root/.openclaw/workspace/harvester/search/provider/bedrock.pyr$   zBedrockProvider.__init__!   sl    !M#=2!L	
 	6N:$%< ?#	
 	
    region
access_key
secret_keyreturnc                     t        |      xs d}t        |      }t        |      }|r|st        j                  d       yt        j                  d|      st        j                  d|dd  d       |||fS )	z#Parse and validate AWS credentials.	us-east-1z7AWS credentials are required: access_key and secret_key)NNNz^AKIA[0-9A-Z]{16}$z"Access key format may be invalid: N   ...)r   loggererrorrematchwarning)r%   r+   r,   r-   s       r)   _parse_credentialsz"BedrockProvider._parse_credentials8   sv     f,*%
*%
LLRS# xx-z:NN?
2A?OsSTz:--r*   c                     t        |      xs d}t        j                  d|      st        j	                  d| d       y|S )z"Validate and normalize AWS region.r0   z"^[a-z]{2}(-[a-z]+)*-[a-z]+-[0-9]+$zRegion format may be invalid: z, using us-east-1)r   r5   r6   r3   r7   )r%   r+   s     r)   _validate_regionz BedrockProvider._validate_regionI   s@    f, xx=vFNN;F8CTUVr*   model_idc                     d| d}| d| dS )z"Build Bedrock InvokeModel API URL.zhttps://bedrock-runtime..amazonaws.comz/model/z/invoke )r%   r+   r;   r   s       r)   _build_invoke_urlz!BedrockProvider._build_invoke_urlT   s#    -fX^D78*G44r*   c                     d| d}| dS )z+Build Bedrock ListFoundationModels API URL.zhttps://bedrock.r=   r   r>   )r%   r+   r   s      r)   _build_models_urlz!BedrockProvider._build_models_urlY   s    %fX^<-..r*   methodurlheaderspayloadc                    t         j                  j                  |      }t         j                  j                  |j                  xs dd      }|j
                  xs d}|rPt         j                  j                  |d      }|j                          t         j                  j                  |      }d}	d}
t        |j                         d       }g }|D ]?  \  }}|j                         }|j                  |       |	| d|j                          d	z  }	A d
j                  |      }
t        j                   |j#                  d            j%                         }| d	| d	| d	|	 d	|
 d	| }t&        j)                  d| d|        ||
fS )z"Build AWS SigV4 canonical request./)safe T)keep_blank_valuesc                 (    | d   j                         S )Nr   lower)xs    r)   <lambda>z:BedrockProvider._build_canonical_request.<locals>.<lambda>q   s    qtzz| r*   )key:
;utf-8zCanonical request built for  )urllibparseurlparsequotepathquery	parse_qslsort	urlencodesorteditemsrM   appendstripjoinhashlibsha256encode	hexdigestr3   debug)r%   rB   rC   rD   rE   parsedurir[   paramscanonical_headerssigned_headerssorted_headersheader_namesrP   value	key_lowerpayload_hashcontents                     r)   _build_canonical_requestz(BedrockProvider._build_canonical_request^   s}    &&s+ll  !3# >" \\++ET+JFKKMLL**62E   5KL( 	BJC		I	*I;ab!AA	B
 ,/ ~~gnnW&=>HHJ HBse2eWB/@.ANCSSUVbUcd3F81SEBC&&r*   	timestampservicecanonical_requestc                     d}|dd }| d| d| d}t        j                  |j                  d            j                         }| d| d| d| }	t        j                  d| d	|        |	|fS )
z Create AWS SigV4 string to sign.zAWS4-HMAC-SHA256Nr1   rG   z/aws4_requestrT   rR   zString to sign created for z in )rd   re   rf   rg   r3   rh   )
r%   ru   r+   rv   rw   	algorithmdatescoperequest_hashrs   s
             r)   _create_string_to_signz&BedrockProvider._create_string_to_sign   s    &	!}&&7)=9 ~~&7&>&>w&GHRRT Kr)BugR~F27)4xHI~r*   string_to_signc                    |dd }t        j                  d| j                         |j                  d      t        j                        j                         }t        j                  ||j                  d      t        j                        j                         }t        j                  ||j                  d      t        j                        j                         }	t        j                  |	dt        j                        j                         }
t        j                  |
|j                  d      t        j                        j                         }t        j                  d       |S )zCalculate AWS SigV4 signature.Nr1   AWS4rT   s   aws4_requestzAWS SigV4 signature calculated)	hmacnewrf   rd   re   digestrg   r3   rh   )r%   r-   ru   r+   rv   r~   rz   k_datek_region	k_service	k_signing	signatures               r)   _calculate_signaturez$BedrockProvider._calculate_signature   s     !} D-446G8Lgnn]ddf88FFMM'$:GNNKRRTHHXw~~g'>OVVX	HHYHOOQ	 HHY(=(=g(FWaac	56r*   credential_scoperm   r   c                 L    d| d| d| d| }t         j                  d       |S )z(Generate AWS SigV4 authorization header.zAWS4-HMAC-SHA256 Credential=rG   z, SignedHeaders=z, Signature=zAuthorization header generated)r3   rh   )r%   r,   r   rm   r   auth_headers         r)   _generate_auth_headerz%BedrockProvider._generate_auth_header   sH    $Q'7&8 9+, -"% 	 	56r*   c                    t        j                  t        j                        }|j	                  d      }	t
        j                  j                  |      }
|
j                  }||	dd}| j                  ||||      \  }}| j                  |	|||      \  }}| j                  ||	|||      }| j                  ||||      }||d<   t        j                  d| d       |S )zSign AWS request with SigV4.z%Y%m%dT%H%M%SZapplication/json)Hostz
X-Amz-DateContent-TypeAuthorizationzRequest signed for z API)r   nowr   utcstrftimerV   rW   rX   netlocrt   r}   r   r   r3   rh   )r%   rB   rC   r+   rv   r,   r-   rE   r   ru   ri   hostrD   rw   rm   textr{   r   r   s                      r)   _sign_requestzBedrockProvider._sign_request   s    
 ll8<<(LL!12	 &&s+}}  yJ\] -1,I,I&RUW^`g,h)> 11)VWN_`e --j)VWVZ[	 00UNT]^#. *7)489r*   c                 f   | j                  d      }	 |dk(  r&t        ||| j                  d      |      }|rd|fS y|dk(  rB	 t        d||j	                  d	      ||
      5 }|j
                  |j                  fcddd       S dd| fS # 1 sw Y   nxY wy# t        j                  j                  $ r }t        |      t        |      fcY d}~S d}~wt        $ r}dt        |      fcY d}~S d}~ww xY w# t        $ r/}t        j                  d|        dt        |      fcY d}~S d}~ww xY w)z6Send HTTP request and return status code and response.   )defaultGETr   )rC   rD   retriestimeout   )  zRequest failedPOSTrT   )datarD   r   Nr     zUnsupported method: zRequest failed: )_get_timeoutr   _get_retriesr   rf   status_coder   requests
exceptions	HTTPErrorr   r   	Exceptionstrr3   r4   )r%   rB   rC   rD   rE   r   responsees           r)   _send_requestzBedrockProvider._send_request   sN   ##B#/	#WdFWFW`aFWFblst=(06!' $^^G4 ' ' C "'33X]]BC C 26(;;;C C C  **44 G,Q/1CA1FFF  'A;&'
  	LL+A3/0A;	s   )C8 C8 B $B<	B C8 BB C57CC5C8 C5C0*C5+C8 0C55C8 8	D0$D+%D0+D0coder   c                    |dk(  r{	 t        j                  |      }d|v sd|v sd|v r)t        j                  d       t	        j
                         S t        j                  d|dd  d	       t	        j
                         S | j                  ||      S # t         j                  $ r, t        j                  d
       t	        j
                         cY S w xY w)zParse AWS Bedrock API response.r   body
completionrs   zBedrock API call successfulzUnexpected response format: Nd   r2   z3Non-JSON response with 200 status, assuming success)	jsonloadsr3   rh   r
   successr7   JSONDecodeError_handle_error)r%   r   r   r   s       r)   _parse_responsezBedrockProvider._parse_response   s    3;-zz(+T>\T%9Y$=NLL!>?&..00NN%A(4C.AQQT#UV&..00 %%dH55 '' -RS"**,,-s   A	B /B <CCr   response_bodyc                    t              t        j                  d| ddd  d       |dk(  rdv rZt        fdd	D              r#t	        j
                  t        j                        S t	        j
                  t        j                        S d
v r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S t	        j
                  t        j                        S |dk(  rdv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S t	        j
                  t        j                        S |dk(  rdv rqdv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S t	        j
                  t        j                        S |dk(  rJdv r#t	        j
                  t        j                        S t	        j
                  t        j                        S |dk(  rqdv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S t	        j
                  t        j                        S |dk\  rqdv r#t	        j
                  t        j                        S dv r#t	        j
                  t        j                        S t	        j
                  t        j                        S t        j                  d |        t	        j
                  t        j                         S )!zHandle AWS Bedrock API errors.zHandling error:  - Nr   r2   r   ValidationExceptionc              3   B   K   | ]  }|j                         v   y wNrL   ).0keywordr   s     r)   	<genexpr>z0BedrockProvider._handle_error.<locals>.<genexpr>  s'       }2244s   )zmodel identifierzinvalid modelzmodel not foundModelNotReadyExceptionModelTimeoutExceptionModelErrorExceptionServiceQuotaExceededExceptioni  UnrecognizedClientExceptionInvalidSignatureExceptionTokenRefreshRequiredi  AccessDeniedExceptionznot authorized to performzdoes not have permissionUnauthorizedOperationRequestExpiredi  ResourceNotFoundExceptioni  ThrottlingExceptionTooManyRequestsExceptionr   InternalServerExceptionServiceUnavailableExceptionzUnknown error status: )r   r3   rh   anyr
   failr	   NO_MODELBAD_REQUESTSERVER_ERRORNO_QUOTAINVALID_KEYEXPIRED_KEY	NO_ACCESSRATE_LIMITEDr7   UNKNOWN)r%   r   r   s     `r)   r   zBedrockProvider._handle_error  sr   ]+'}Cds8K7LCPQ#$5 #[  '++K,@,@AA&++K,C,CDD)]:"''(<(<==(M9"''(@(@AA&-7"''(<(<==0MA"''(<(<=="''(?(?@@C,="''(?(?@@,="''(?(?@@'=8"''(?(?@@"''(?(?@@C&-7.-?&++K,A,ABB/=@&++K,A,ABB&++K,A,ABB(M9"''(?(?@@!]2"''(?(?@@"''(=(=>>C*m;"''(<(<=="''(<(<==C$5"''(@(@AA+}<"''(@(@AA"''(@(@AAC(M9"''(@(@AA.-?"''(@(@AA"''(@(@AANN3K=AB##K$7$788r*   token
additionalc                 *    | j                  ddi|      S )z<Get headers for AWS Bedrock API requests. Not used directly.r   r   )_merge_headers)r%   r   r   s      r)   _get_headerszBedrockProvider._get_headersR  s     ""N4F#GTTr*   c                    |j                         }d|v sd|v r!t        j                  dddt        dgd      S d|v sd	|v rt        j                  t        dd
dd      S d|v sd|v rt        j                  t        dd
d      S d|v sd|v rt        j                  t        dd
d      S d|v sd|v rt        j                  t        dd
d      S t        j                  d| d       t        j                  dddt        dgd      S )z3Build appropriate test payload based on model type.	anthropicclaudezbedrock-2023-05-31
   user)rolers   )anthropic_version
max_tokensmessagesamazontitang?)maxTokenCounttemperature)	inputTexttextGenerationConfigai21jurassic)prompt	maxTokensr   coherecommand)r   r   r   metallama)r   max_gen_lenr   zUnknown model type z, using Anthropic format)rM   r   dumpsr   r3   rh   )r%   r;   models      r)   _build_test_payloadz#BedrockProvider._build_test_payloadX  s.    %8u#4::)="$*0=M!N O  'U"2::!1>@QT,U  u_
e 3::)9[^_``)u"4::)9\_`aau_5 0::)9"]`abb LL.xj8PQR::)="$*0=M!N O r*   addressendpointr   c           	      $   | j                  |||      \  }}}t        |||g      s8t        j                  d       t	        j
                  t        j                        S | j                  |      }t        |      xs | j                  }| j                  ||      }	| j                  |      }
	 | j                  d|	|d|||
      }| j                  d|	||
      \  }}| j                  ||      }|j                   rt        j#                  d| d|        |S t        j%                  d|j&                   d| d	| d
       |S # t(        $ rE}t        j                  d|        t	        j
                  t        j*                        cY d}~S d}~ww xY w)z:Check if AWS credentials are valid by calling Bedrock API.z Invalid AWS credentials providedr   zbedrock-runtimerB   rC   r+   rv   r,   r-   rE   z6Bedrock credentials validated successfully for region z with model z'Bedrock credentials validation failed: z
 (region: z	, model: )zBedrock check failed: N)r8   allr3   r4   r
   r   r	   r   r:   r   _default_modelr?   r   r   r   r   	availableinfor7   reasonr   r   )r%   r   r   r   r   r+   r,   r-   r;   rC   rE   rD   r   r   resultr   s                   r)   checkzBedrockProvider.check  s    *.)@)@(TY)Z&
JFJ
34LL;<##K$;$;<< &&v. ;5$"5"5 $$VX6 **84	9(()%% ) G *.););FCRY)Z&K ))+}EFTU[T\\hiqhrst M	 =fmm_JW]V^^ghpgqqrs M 	9LL1!56##K$7$788	9s%   ,A)E *E 	F
:F
F
Fc           	         | j                  |||      \  }}}t        |||g      st        j                  d       g S | j	                  |      }| j                  |      }	 | j                  d||d||d      }| j                  d||      \  }	}
|	dk(  re	 t        j                  |
      }g }d|v r#|d   D ]  }d|v s|j                  |d           t        j                  d	t        |       d
|        |S t        j                  d|	 d|
dd  d       g S # t        j                  $ r$}t        j                  d|        g cY d}~S d}~ww xY w# t        $ r$}t        j                  d|        g cY d}~S d}~ww xY w)z'List available models from AWS Bedrock.z2Invalid AWS credentials provided for model listingr   r   rI   r  r   modelSummariesmodelIdz
Retrieved z models from Bedrock in region z!Failed to parse models response: NzFailed to list models: r   r2   zBedrock list_models failed: )r8   r  r3   r4   r:   rA   r   r   r   r   ra   r  lenr   r   )r%   r   r   r   r+   r,   r-   rC   rD   r   r   r   models
model_infor   s                  r)   inspectzBedrockProvider.inspect  s    *.)@)@(TY)Z&
JFJ
34LLMNI &&v. $$V,&	((!%% ) G "//sGDND(s{::h/DF (4/*./?*@ EJ(J6 &j.C DE KK*S[M9XY_X` ab!M 6tfC#?OsST	 ++ LL#DQC!HII  	LL7s;<I	sT    3E 'D <;D 8 E E,EEE EE 	F E;5F ;F )rI   r   )rI   rI   rI   )rI   rI   )__name__
__module____qualname____doc__r   r   r$   r   tupler8   r:   r?   rA   r   rt   r}   r   r   r   r   intr
   r   r   r   r   r   r	  r  __classcell__)r(   s   @r)   r   r      s   @
4	? 
.. .# .3 .SX ."	s 	s 	5 5s 5s 5
/ / /
$'s $' $'tCQTH~ $'`c $'hmnqsvnvhw $'L S 3 cf kn *-7:EHZ]	$
 
s 
\_ 
lo 
tw 
 rt   # -0 ;> LO ]` kn 	c3h DC c DcN UX bg @6C 63 6; 6&J9 J9S J9[ J9XU# U8D> UXVZ^ U'C 'C 'R293 29 29S 29c 29[f 29h5S 53 5s 5DQTI 5r*   r   r   )$r  rd   r   r   r5   urllib.parserV   r   r   typingr   r   r   r   constant.systemr   
core.enumsr	   core.modelsr
   r   tools.loggerr   tools.utilsr   clientr   r   r   r   baser   registryr   r3   r   r>   r*   r)   <module>r"     sg       	  ' ' '  , " . #  M M   '	J	Ln L^ )_ -r*   