Fix black

This commit is contained in:
Sameer Kankute 2026-04-27 09:31:47 +05:30
parent c014bfa683
commit 367c48e815
No known key found for this signature in database
3 changed files with 315 additions and 902 deletions

File diff suppressed because it is too large Load diff

View file

@ -35,13 +35,9 @@ class PredibaseConfig(BaseConfig):
best_of: Optional[int] = None
decoder_input_details: Optional[bool] = None
details: bool = True # enables returning logprobs + best of
max_new_tokens: int = (
DEFAULT_MAX_TOKENS # openai default - requests hang if max_new_tokens not given
)
max_new_tokens: int = DEFAULT_MAX_TOKENS # openai default - requests hang if max_new_tokens not given
repetition_penalty: Optional[float] = None
return_full_text: Optional[bool] = (
False # by default don't return the input as part of the output
)
return_full_text: Optional[bool] = False # by default don't return the input as part of the output
seed: Optional[int] = None
stop: Optional[List[str]] = None
temperature: Optional[float] = None
@ -108,9 +104,7 @@ class PredibaseConfig(BaseConfig):
optional_params["top_p"] = value
if param == "n":
optional_params["best_of"] = value
optional_params["do_sample"] = (
True # Need to sample if you want best of for hf inference endpoints
)
optional_params["do_sample"] = True # Need to sample if you want best of for hf inference endpoints
if param == "stream":
optional_params["stream"] = value
if param == "stop":
@ -176,9 +170,7 @@ class PredibaseConfig(BaseConfig):
)
if "details" in completion_response and "tokens" in completion_response["details"]:
model_response.choices[0].finish_reason = map_finish_reason(
completion_response["details"]["finish_reason"]
)
model_response.choices[0].finish_reason = map_finish_reason(completion_response["details"]["finish_reason"])
sum_logprob = 0
for token in completion_response["details"]["tokens"]:
if token["logprob"] is not None:
@ -198,10 +190,7 @@ class PredibaseConfig(BaseConfig):
best_of_value = 0
if best_of_value > 1:
if (
"details" in completion_response
and "best_of_sequences" in completion_response["details"]
):
if "details" in completion_response and "best_of_sequences" in completion_response["details"]:
choices_list = []
for idx, item in enumerate(completion_response["details"]["best_of_sequences"]):
sum_logprob = 0
@ -233,11 +222,7 @@ class PredibaseConfig(BaseConfig):
if output_text is not None and len(output_text) > 0:
completion_tokens = 0
try:
completion_tokens = len(
encoding.encode(
model_response["choices"][0]["message"].get("content", "")
)
)
completion_tokens = len(encoding.encode(model_response["choices"][0]["message"].get("content", "")))
except Exception:
# Keep usage calculation non-blocking if encoding fails.
pass
@ -327,9 +312,7 @@ class PredibaseConfig(BaseConfig):
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
tenant_id = litellm_params.get("predibase_tenant_id") or litellm_params.get(
"tenant_id"
)
tenant_id = litellm_params.get("predibase_tenant_id") or litellm_params.get("tenant_id")
if tenant_id is None:
raise ValueError(
"Missing Predibase Tenant ID - Required for making the request. Set dynamically (e.g. `completion(..tenant_id=<MY-ID>)`) or in env - `PREDIBASE_TENANT_ID`."
@ -349,12 +332,8 @@ class PredibaseConfig(BaseConfig):
completion_url += "/generate"
return completion_url
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, Headers]
) -> BaseLLMException:
return PredibaseError(
status_code=status_code, message=error_message, headers=headers
)
def get_error_class(self, error_message: str, status_code: int, headers: Union[dict, Headers]) -> BaseLLMException:
return PredibaseError(status_code=status_code, message=error_message, headers=headers)
def validate_environment(
self,

View file

@ -98,9 +98,7 @@ class XecGuardGuardrail(CustomGuardrail):
"the guardrail config."
)
self.api_base = (
api_base or os.environ.get("XECGUARD_API_BASE") or _DEFAULT_API_BASE
).rstrip("/")
self.api_base = (api_base or os.environ.get("XECGUARD_API_BASE") or _DEFAULT_API_BASE).rstrip("/")
self.xecguard_model = xecguard_model or _DEFAULT_MODEL
self.policy_names = policy_names
@ -115,9 +113,7 @@ class XecGuardGuardrail(CustomGuardrail):
else:
self.block_on_error = block_on_error
self.grounding_strictness = (
grounding_strictness or _DEFAULT_GROUNDING_STRICTNESS
)
self.grounding_strictness = grounding_strictness or _DEFAULT_GROUNDING_STRICTNESS
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback,
@ -179,16 +175,11 @@ class XecGuardGuardrail(CustomGuardrail):
messages=messages,
documents=documents,
)
if (
grounding_result is not None
and grounding_result.get("decision") == "UNSAFE"
):
if grounding_result is not None and grounding_result.get("decision") == "UNSAFE":
raise HTTPException(
status_code=400,
detail={
"error": self._format_grounding_block_message(
grounding_result
),
"error": self._format_grounding_block_message(grounding_result),
"guardrail_name": self.guardrail_name or "xecguard",
"xecguard_response": grounding_result,
},
@ -212,7 +203,7 @@ class XecGuardGuardrail(CustomGuardrail):
isinstance(kwargs, dict)
and "litellm_params" in kwargs
and "metadata" in kwargs["litellm_params"]
and "standard_logging_guardrail_information"in kwargs["litellm_params"]["metadata"]
and "standard_logging_guardrail_information" in kwargs["litellm_params"]["metadata"]
and kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
):
return kwargs, result
@ -249,9 +240,7 @@ class XecGuardGuardrail(CustomGuardrail):
return kwargs, result
guardrail_status: GuardrailStatus = (
"guardrail_intervened"
if scan_result.get("decision") == "UNSAFE"
else "success"
"guardrail_intervened" if scan_result.get("decision") == "UNSAFE" else "success"
)
end_time = datetime.now()
kwargs["standard_logging_object"]["guardrail_information"] = {
@ -292,11 +281,7 @@ class XecGuardGuardrail(CustomGuardrail):
asyncio.set_event_loop(loop)
if loop.is_running():
return kwargs, result
loop.run_until_complete(
self.async_logging_hook(
kwargs=kwargs, result=result, call_type=call_type
)
)
loop.run_until_complete(self.async_logging_hook(kwargs=kwargs, result=result, call_type=call_type))
except Exception as exc:
verbose_proxy_logger.debug(
"XecGuard sync logging_hook swallowed exception: %s",
@ -318,9 +303,7 @@ class XecGuardGuardrail(CustomGuardrail):
"model": self.xecguard_model,
"scan_type": scan_type,
"messages": messages,
"policy_names": (
self.policy_names if self.policy_names else _DEFAULT_POLICIES
),
"policy_names": (self.policy_names if self.policy_names else _DEFAULT_POLICIES),
}
return await self._post(
path=_SCAN_ENDPOINT,
@ -378,9 +361,7 @@ class XecGuardGuardrail(CustomGuardrail):
raise HTTPException(
status_code=400,
detail={
"error": (
f"XecGuard API unreachable " f"(block_on_error=True): {exc}"
),
"error": (f"XecGuard API unreachable (block_on_error=True): {exc}"),
"guardrail_name": self.guardrail_name or "xecguard",
},
) from exc
@ -404,9 +385,7 @@ class XecGuardGuardrail(CustomGuardrail):
the request data is incomplete.
"""
raw_messages = request_data.get("messages") or []
messages: List[dict] = [
self._normalize_message(m) for m in raw_messages if isinstance(m, dict)
]
messages: List[dict] = [self._normalize_message(m) for m in raw_messages if isinstance(m, dict)]
if input_type == "request":
if not messages:
@ -419,9 +398,7 @@ class XecGuardGuardrail(CustomGuardrail):
return messages
# input_type == "response"
assistant_text = self._extract_assistant_text_from_response(
request_data.get("response")
)
assistant_text = self._extract_assistant_text_from_response(request_data.get("response"))
if assistant_text is None:
return []
messages.append({"role": "assistant", "content": assistant_text})
@ -498,9 +475,7 @@ class XecGuardGuardrail(CustomGuardrail):
parts = [
item.get("text")
for item in content
if isinstance(item, dict)
and item.get("type") == "text"
and isinstance(item.get("text"), str)
if isinstance(item, dict) and item.get("type") == "text" and isinstance(item.get("text"), str)
]
joined = "\n".join(p for p in parts if p)
return joined or None
@ -563,10 +538,7 @@ class XecGuardGuardrail(CustomGuardrail):
if isinstance(candidate, str) and candidate:
rationale = candidate[:_RATIONALE_TRUNCATE_CHARS]
break
return (
f"Blocked by XecGuard: policies=[{policies}] "
f"trace_id={trace_id} rationale={rationale}"
)
return f"Blocked by XecGuard: policies=[{policies}] trace_id={trace_id} rationale={rationale}"
@staticmethod
def _format_grounding_block_message(result: dict) -> str:
@ -582,7 +554,4 @@ class XecGuardGuardrail(CustomGuardrail):
if isinstance(candidate, str):
rationale = candidate[:_RATIONALE_TRUNCATE_CHARS]
rules_str = ",".join(rules) if rules else "unknown"
return (
f"Blocked by XecGuard grounding: rules=[{rules_str}] "
f"trace_id={trace_id} rationale={rationale}"
)
return f"Blocked by XecGuard grounding: rules=[{rules_str}] trace_id={trace_id} rationale={rationale}"