@@ -71,7 +71,7 @@ def __init__(self, region: str = "cn-beijing", timeout: int = 50) -> None:
7171 "TOOL_LLM_SHIELD_URL" ,
7272 f"https://{ self .region } .sdk.access.llm-shield.omini-shield.com" ,
7373 )
74- self .api_key = getenv ("TOOL_LLM_SHIELD_API_KEY" )
74+ self .api_key = getenv ("TOOL_LLM_SHIELD_API_KEY" , allow_false_values = True )
7575
7676 self .category_map = {
7777 101 : "Model Misuse" ,
@@ -101,18 +101,6 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]:
101101 logger .error ("LLM Shield app ID not configured" )
102102 return None
103103
104- ak = os .getenv ("VOLCENGINE_ACCESS_KEY" )
105- sk = os .getenv ("VOLCENGINE_SECRET_KEY" )
106- session_token = ""
107- if not (ak and sk ):
108- logger .debug ("Get AK/SK from environment variables failed." )
109- credential = get_credential_from_vefaas_iam ()
110- ak = credential .access_key_id
111- sk = credential .secret_access_key
112- session_token = credential .session_token
113- else :
114- logger .debug ("Successfully get AK/SK from environment variables." )
115-
116104 body = {
117105 "Message" : {
118106 "Role" : role ,
@@ -124,25 +112,46 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]:
124112
125113 body_json = json .dumps (body ).encode ("utf-8" )
126114
127- header = {"X-Security-Token" : session_token }
128- # Add x-api-key header if API key is provided
129- if self .api_key :
130- header ["x-api-key" ] = self .api_key
131115 path = "/v2/moderate"
132116 action = "Moderate"
133117 version = "2025-08-31"
134118
135- signed_header = request_sign (
136- header , ak , sk , self .region , self .url , path , action , body_json
137- )
138-
139- signed_header . update (
140- {
119+ # Check if using API key authentication
120+ logger . debug ( f"API key value: { self .api_key } , type: { type ( self .api_key ) } " )
121+ if self . api_key and self . api_key != "" :
122+ logger . debug ( "Using API key authentication (no AK/SK signature)" )
123+ # Use API key authentication only - match curl command headers exactly
124+ signed_header = {
141125 "Content-Type" : "application/json" ,
142- "X-Top-Service" : "llmshield" ,
143- "X-Top-Region" : self .region ,
126+ "x-api-key" : self .api_key ,
144127 }
145- )
128+ else :
129+ logger .debug ("Using AK/SK signature authentication" )
130+ # Use AK/SK signature authentication
131+ ak = os .getenv ("VOLCENGINE_ACCESS_KEY" )
132+ sk = os .getenv ("VOLCENGINE_SECRET_KEY" )
133+ session_token = ""
134+ if not (ak and sk ):
135+ logger .debug ("Get AK/SK from environment variables failed." )
136+ credential = get_credential_from_vefaas_iam ()
137+ ak = credential .access_key_id
138+ sk = credential .secret_access_key
139+ session_token = credential .session_token
140+ else :
141+ logger .debug ("Successfully get AK/SK from environment variables." )
142+
143+ header = {"X-Security-Token" : session_token }
144+ signed_header = request_sign (
145+ header , ak , sk , self .region , self .url , path , action , body_json
146+ )
147+
148+ signed_header .update (
149+ {
150+ "Content-Type" : "application/json" ,
151+ "X-Top-Service" : "llmshield" ,
152+ "X-Top-Region" : self .region ,
153+ }
154+ )
146155
147156 try :
148157 response = requests .post (
@@ -158,7 +167,6 @@ def _request_llm_shield(self, message: str, role: str) -> Optional[str]:
158167 f"LLM Shield HTTP error: { response .status_code } - { response .text } "
159168 )
160169 return None
161-
162170 response = response .json ()
163171 except requests .exceptions .Timeout :
164172 logger .error ("LLM Shield request timeout" )
0 commit comments