mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix black issues
This commit is contained in:
Sameer Kankute
parent
367c48e815
commit
77df511559
3 changed files with 902 additions and 315 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -35,9 +35,13 @@ 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
|
||||
|
|
@ -104,7 +108,9 @@ 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":
|
||||
|
|
@ -169,8 +175,13 @@ class PredibaseConfig(BaseConfig):
|
|||
completion_response["generated_text"]
|
||||
)
|
||||
|
||||
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"])
|
||||
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"]
|
||||
)
|
||||
sum_logprob = 0
|
||||
for token in completion_response["details"]["tokens"]:
|
||||
if token["logprob"] is not None:
|
||||
|
|
@ -190,9 +201,14 @@ 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"]):
|
||||
for idx, item in enumerate(
|
||||
completion_response["details"]["best_of_sequences"]
|
||||
):
|
||||
sum_logprob = 0
|
||||
for token in item["tokens"]:
|
||||
if token["logprob"] is not None:
|
||||
|
|
@ -222,7 +238,11 @@ 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
|
||||
|
|
@ -312,7 +332,9 @@ 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`."
|
||||
|
|
@ -325,15 +347,21 @@ class PredibaseConfig(BaseConfig):
|
|||
base_url = os.getenv("PREDIBASE_API_BASE", "")
|
||||
|
||||
completion_url = f"{base_url}/{tenant_id}/deployments/v2/llms/{model}"
|
||||
should_stream = stream if stream is not None else optional_params.get("stream", False)
|
||||
should_stream = (
|
||||
stream if stream is not None else optional_params.get("stream", False)
|
||||
)
|
||||
if should_stream is True:
|
||||
completion_url += "/generate_stream"
|
||||
else:
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -98,7 +98,9 @@ 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
|
||||
|
|
@ -113,7 +115,9 @@ 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,
|
||||
|
|
@ -175,11 +179,16 @@ 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,
|
||||
},
|
||||
|
|
@ -203,8 +212,11 @@ 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 kwargs["litellm_params"]["metadata"]["standard_logging_guardrail_information"]
|
||||
and "standard_logging_guardrail_information"
|
||||
in kwargs["litellm_params"]["metadata"]
|
||||
and kwargs["litellm_params"]["metadata"][
|
||||
"standard_logging_guardrail_information"
|
||||
]
|
||||
):
|
||||
return kwargs, result
|
||||
|
||||
|
|
@ -240,7 +252,9 @@ 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"] = {
|
||||
|
|
@ -281,7 +295,11 @@ 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",
|
||||
|
|
@ -303,7 +321,9 @@ 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,
|
||||
|
|
@ -361,7 +381,9 @@ class XecGuardGuardrail(CustomGuardrail):
|
|||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": (f"XecGuard API unreachable (block_on_error=True): {exc}"),
|
||||
"error": (
|
||||
f"XecGuard API unreachable (block_on_error=True): {exc}"
|
||||
),
|
||||
"guardrail_name": self.guardrail_name or "xecguard",
|
||||
},
|
||||
) from exc
|
||||
|
|
@ -385,7 +407,9 @@ 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:
|
||||
|
|
@ -398,7 +422,9 @@ 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})
|
||||
|
|
@ -475,7 +501,9 @@ 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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue