mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Litellm OSS Staging (#29161)
* Cato Networks guardrail, based on Aim (#26597) * Aim was acquired by Cato Networks, creating Cato Networks guardrail based on Aim * Add more tests * Move test so they are reached by codecov coverage * base URL trailing slashes * Support Lemonade runtime context metadata (#28135) * Support Lemonade runtime context metadata * Add provider hook for runtime model metadata * Address provider model info review feedback Keep the runtime model info hook duck-typed instead of extending the base model-info class, and avoid importing ModelInfoBase from Ollama common utilities to reduce CodeQL cyclic-import noise. Co-authored-by: openhands <openhands@all-hands.dev> * Fix CI after staging rebase Relax the Ollama runtime metadata return annotation to match the provider-hook dict response and update the Google Interactions OpenAPI status expectation for the current live spec. Co-authored-by: openhands <openhands@all-hands.dev> * Normalize Lemonade runtime model metadata * Avoid leaking Ollama metadata auth * Avoid leaking Lemonade metadata auth --------- Co-authored-by: Graham Neubig <398875+neubig@users.noreply.github.com> Co-authored-by: openhands <openhands@all-hands.dev> * fix(cato): address guardrail review feedback Use proxy-authenticated user identity, forward moderation hook return values, and ensure streaming sender tasks are cancelled and awaited on exit. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path - clone of #28010 (#28846) * fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path Fixes #26083 vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the NON_GEMINI route. Per owtaylor's plan on #26083: add the google/gemma- prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up and should_use_openai_handler routes it to the OpenAI-compatible /endpoints/openapi/chat/completions URL. No gemma-detection exclusion needed (the "gemma/" check uses a slash, which google/gemma-... doesn't match). No OpenAIGPTConfig subclass needed — works with the base handler. * fix(vertex_ai): mark gemma-4-26b-a4b-it-maas as vision-capable (empirically verified) * fix(vertex_ai): address greptile feedback — provider category, canonical URL, sync backup * test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS Addresses oss-pr-review-agent-shin feedback on PR #28010: supports_function_calling, supports_tool_choice, and supports_vision were marked true but had no tests proving the payloads actually reached the OpenAI-compatible endpoint. Added: - test_gemma_maas_supports_function_calling — verifies the utility returns True when the model_cost entry carries supports_function_calling=true - test_gemma_maas_supports_vision — same for supports_vision - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice appear in the JSON body POSTed to /endpoints/openapi/chat/completions - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts survive transformation and reach the global endpoint URL * fix: Delete uv.lock * test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS Addresses oss-pr-review-agent-shin feedback on PR #28010: P1 (patch target): Added a comment explaining why patching litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct — get_async_httpx_client() (defined in http_handler.py) instantiates AsyncHTTPHandler within that module's scope, so the definition-site patch intercepts it. Without the mock the test raises AuthenticationError, confirming it never silently passes. P2 (partner-provider regression guard): Added test_gemma_routes_through_openai_handler() which calls VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's routing to VertexPartnerProvider.llama ever changes the URL-shape tests below it become a real regression guard rather than an unanchored unit test. Also added: - test_gemma_maas_supports_function_calling / supports_vision — capability flag checks via patch.dict(litellm.model_cost) - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice forwarded in the request body - test_vertex_ai_gemma_vision_passthrough — image_url part survives transformation to the global endpoint Added: - test_gemma_maas_supports_function_calling — verifies the utility returns True when the model_cost entry carries supports_function_calling=true - test_gemma_maas_supports_vision — same for supports_vision - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice appear in the JSON body POSTed to /endpoints/openapi/chat/completions - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts survive transformation and reach the global endpoint URL * fix: proper patch for unit tests --------- Co-authored-by: Iana <iana@Shivakumars-MacBook-Pro.local> * fix(cato): guardrail all completion choices on output When n > 1, only choices[0] was analyzed and redacted. Iterate every Choices entry so block and anonymize actions apply to all completions. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix review * fix(cato_networks): harden output anonymize handling and restructure nested UI routes Guard against empty redacted_output and empty all_redacted_messages from Cato. Restructure nested admin UI HTML exports to index.html so extensionless routes work. Co-authored-by: Cursor <cursoragent@cursor.com> * Fix mypy * fix(cato): guard missing policy_drill_down and all_redacted_messages keys * fix(cato): avoid KeyError bypassing block action on missing analysis_result * fix(cato): preserve non-text message fields during anonymize Rebuild redacted messages from the original messages, overwriting only content, so tool_calls, tool_call_id, name and multimodal fields survive the anonymize action. * fix(cato): preserve trailing messages when fewer redacted messages returned Avoid silently truncating the conversation in _anonymize_request when Cato returns fewer redacted messages than were sent, and isolate the no-api-key config test from a pre-existing CATO_API_KEY environment variable. * fix(cato,model-info): preserve stream block signal on sender teardown; forward api_key in dynamic model-info lookup Suppress ConnectionClosed (alongside CancelledError) when tearing down the Cato streaming sender task so a backend ConnectionClosed cannot mask the original StreamingCallbackError (e.g. a guardrail block) raised by the receive loop. Thread api_key through get_model_info -> _get_model_info_helper so an explicit key reaches a provider's dynamic get_model_info for a caller-supplied api_base. Previously only api_base was forwarded, so authenticated Ollama and Lemonade servers at a custom base could only be queried unauthenticated. * fix(cato): surface mid-stream forwarding errors instead of blocking on recv If the upstream LLM stream errors mid-flight, the sender task dies before sending the terminal done frame, so the consumer would block on websocket.recv() until Cato closes the connection. Race recv against the sender task and raise the stored sender exception promptly as a StreamingCallbackError. * fix(cato): drop spoofable end_user_id from guardrail user identity Only the key/JWT-bound user_email is a trusted identity. end_user_id is resolved from caller-supplied request fields (OpenAI user param, headers, metadata), so an authenticated caller with no bound user_email could set it to another user's email and have LiteLLM forward x-cato-user-email for that victim, poisoning Cato audit and policy attribution. Forward only user_email and omit the header otherwise. * fix(cato): harden output anonymize path against missing content key * fix(cato): fall back to original message when redacted content key is missing * refactor(model-info): drop unused api_key from cached model-info helper _cached_get_model_info_helper is only called by the cost-tracking hot path, which never authenticates, so the api_key parameter was never populated. Keeping it in the lru_cache key offered no benefit and risked fragmenting the high-RPS cache and retaining credential strings per entry. * fix(cato): preserve None content on tool-call-only choices in output hook * fix(ollama): respect static-model guard in OllamaConfig.get_model_info Delegate to OllamaModelInfo.get_model_info so statically-priced Ollama models short-circuit before the /api/show network call instead of hitting the server unconditionally. * fix(lemonade,ollama): treat empty api_key as unset to avoid leaking server creds An empty-string api_key was treated as an explicit key, so it passed the guard meant to keep server-side credentials off caller-supplied bases and then fell back through the env/global key chain. A caller could point api_base at a server they control and send api_key="" to receive the configured provider key in the Authorization header. Gate the credential fallback on the api_key being truthy instead of merely not-None. * fix(cato): inspect and redact Responses-API input, not just messages The guardrail only read data["messages"], so /v1/responses requests, which carry their text in data["input"], reached Cato as an empty message list and bypassed inspection entirely. Send build_inspection_messages(data) so both shapes are analyzed, and write anonymized results back with apply_redacted_messages_back when the request used input. * perf(utils): keep api_key out of get_model_info lru_cache key * fix(cato): propagate ssl_verify to streaming WebSocket connection The streaming hook applied ssl_verify only to the HTTP handler; the websockets.connect() call used default verification, so a custom Cato instance behind TLS with a self-signed cert worked for non-streaming calls but failed every streaming request. Resolve the ssl_verify setting into the connect() ssl argument, mirroring the HTTP handler. * refactor(utils): rename shadowing local in _get_model_info_helper * fix(cato): flatten multimodal chat content before inspection Chat Completions requests whose message content is a multimodal parts array were posted to Cato as the raw OpenAI parts, so text inside content: [{"type":"text", ...}] reached the model without Cato ever inspecting the string. Flatten each message's list content to plain text while keeping the list 1:1 with the request so the index-based redaction write-back stays valid; Responses-API input requests still go through build_inspection_messages. * test(lemonade): clear get_model_info cache around api_base test * fix(cato): inspect and redact Responses-API input even when messages present _inspection_messages returned early once messages was non-empty, so a /v1/responses caller could place benign text in messages and disallowed text in input and have only messages reach Cato while the model used input. Inspect both fields and write anonymize redactions back to input as well as the index-aligned messages. * test(log_db_metrics): assert table_name event_metadata contract log_db_metrics now emits minimal event_metadata via _safe_db_event_metadata (table_name only, function_name/function_kwargs/function_args dropped as redundant with call_type and unsafe to stamp on a span). The success-path test still asserted function_name membership and crashed with TypeError on the None metadata returned when no table_name is passed. Pass a table_name and assert the surfaced contract instead. * fix(cato): inspect and redact completion prompt and Responses-API instructions The Cato guardrail only inspected chat messages and the Responses-API input field, so blocked text placed in the legacy /v1/completions prompt or the /v1/responses instructions field reached the model without ever being sent to Cato. Both fields are now appended as synthetic inspection messages, and the anonymize path slices Cato's redactions back to the field they came from. * fix(cato): serialize non-str/bytes websocket chunks before forwarding * fix(cato): inspect tool descriptions and tool-call arguments * fix(cato): map redacted output by assistant index; restore get_model_info.cache_info * fix(cato): block output even when detection_message is null/empty A block_action returned by Cato on the output hook whose detection_message was null or empty was let through to the caller: the truthiness guard on detection_message skipped the HTTPException and the unblocked response was returned. Raise the HTTPException directly in _handle_block_action_on_output so the output path blocks unconditionally, mirroring the input path. * fix(cato): inspect and redact nested tool param and legacy function descriptions Tool/function parameter descriptions and the legacy functions[] array are forwarded to the model but were not seen by Cato, so blocked text hidden there bypassed inspection and anonymization. Recursively walk every description string in tools[].function and functions[] schemas for both the analyze payload and the anonymize write-back. * fix(cato): traverse schema descriptions iteratively to satisfy recursive detector The nested walk() generator recursed over tool/function JSON schemas with no depth bound, which the recursive_detector code-quality gate rejects. Replace it with an explicit-stack DFS that yields the same (container, key) refs in the same pre-order, so schema description redaction is unchanged. * fix(cato): inspect and redact response_format JSON schema descriptions response_format json_schema descriptions are forwarded to the model, so blocked text hidden in nested schema descriptions could bypass Cato inspection and redaction. Extend the schema-description walk to cover response_format alongside tools and legacy functions. * fix(cato): skip output rewrite when Cato returns no redaction Return None from call_cato_guardrail_on_output on monitor/no-action so the post-call hook only mutates the message when there is an actual redaction, instead of redundantly re-writing the original content. * refactor(utils): resolve explicit api_key model info without the cache Move the model-info build into a non-cached _build_model_info helper and drop api_key from the lru-cached _cached_get_model_info signature. Both cached helpers now take the same (model, provider, api_base) key and never forward api_key, while explicit per-caller keys are resolved through the builder directly instead of reaching into the cache wrapper's __wrapped__. * fix(cato): inspect and redact non-description schema string values Tool, function and response_format JSON schemas forward more than just description text to the model. enum, const, default, examples and title values are sent verbatim, so blocked content hidden in any of them bypassed Cato inspection and redaction. Walk those schema string values alongside descriptions on both the inspection and anonymize paths. * fix(model-info): surface swallowed dynamic model-info errors The provider-specific get_model_info dispatch falls back to the static cost map when a provider's dynamic lookup raises, which is intentional graceful degradation. Previously the exception was discarded with a bare debug line, so a real failure (e.g. a provider whose get_model_info signature does not accept api_key) was invisible. Log the exception at warning level with the model and provider context so the fallback is diagnosable. * fix(cato): inspect and redact Responses API output in post-call hook The post-call success hook only handled ModelResponse, so /v1/responses (which returns a ResponsesAPIResponse) bypassed the Cato output guardrail. Extract and inspect/redact every output_text content block and function-call arguments string, blocking on a block action, so generated text cannot escape inspection by using the Responses API. * chore: reset _experimental/out folder * chore(ui): remove orphaned prebuilt dashboard chunk files The _experimental/out manifests are byte-identical to the base branch, so the served dashboard already matches base. 436 unreferenced Next.js chunk files had accumulated in the directory and are not loaded by any manifest; removing them restores the committed UI artifacts to the base build and drops the artifact churn from this PR's diff. * fix(guardrails,ollama): forward ssl_verify to Cato init and raise_for_status on /api/show --------- Co-authored-by: Alex Yaroslavsky <trexinc@gmail.com> Co-authored-by: Graham Neubig <neubig@gmail.com> Co-authored-by: Graham Neubig <398875+neubig@users.noreply.github.com> Co-authored-by: openhands <openhands@all-hands.dev> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Piotr Placzko <piotr@icep-design.com> Co-authored-by: Iana <iana@Shivakumars-MacBook-Pro.local> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Sameer Kankute
GitHub
parent
68952a55d7
commit
e8fcb01215
26 changed files with 4937 additions and 108 deletions
|
|
@ -868,6 +868,7 @@ openai_text_completion_compatible_providers: List = (
|
|||
_openai_like_providers: List = [
|
||||
"predibase",
|
||||
"databricks",
|
||||
"lemonade",
|
||||
"watsonx",
|
||||
] # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk
|
||||
# well supported replicate llms
|
||||
|
|
|
|||
|
|
@ -87,6 +87,7 @@ class ExceptionCheckers:
|
|||
"is longer than the model's context length",
|
||||
"input tokens exceed the configured limit",
|
||||
"`inputs` tokens + `max_new_tokens` must be",
|
||||
"exceeds the available context size", # llama.cpp/Lemonade
|
||||
"exceeds the maximum number of tokens allowed", # Gemini
|
||||
]
|
||||
for substring in known_exception_substrings:
|
||||
|
|
@ -891,12 +892,14 @@ def exception_type( # type: ignore # noqa: PLR0915
|
|||
response=getattr(original_exception, "response", None),
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "model's maximum context limit" in error_str:
|
||||
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
|
||||
exception_mapping_worked = True
|
||||
raise ContextWindowExceededError(
|
||||
message=f"{custom_llm_provider.capitalize()}Exception: Context Window Error - {error_str}",
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=getattr(original_exception, "response", None),
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
elif "token_quota_reached" in error_str:
|
||||
exception_mapping_worked = True
|
||||
|
|
|
|||
|
|
@ -3,10 +3,12 @@ Translate from OpenAI's `/v1/chat/completions` to Lemonade's `/v1/chat/completio
|
|||
"""
|
||||
|
||||
from typing import Any, List, Optional, Tuple, Union
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -18,6 +20,8 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig
|
|||
|
||||
|
||||
class LemonadeChatConfig(OpenAILikeChatConfig):
|
||||
_DEFAULT_API_KEY = "lemonade"
|
||||
|
||||
repeat_penalty: Optional[float] = None
|
||||
functions: Optional[list] = None
|
||||
logit_bias: Optional[dict] = None
|
||||
|
|
@ -68,7 +72,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
|
|||
This method queries the Lemonade /models endpoint to retrieve the list of available models.
|
||||
|
||||
Args:
|
||||
api_key: Optional API key (Lemonade doesn't require authentication)
|
||||
api_key: Optional API key for authenticated Lemonade servers
|
||||
api_base: Optional API base URL (defaults to LEMONADE_API_BASE env var or http://localhost:8000)
|
||||
|
||||
Returns:
|
||||
|
|
@ -87,6 +91,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
|
|||
try:
|
||||
response = litellm.module_level_client.get(
|
||||
url=f"{api_base}/models",
|
||||
headers=self._get_auth_headers(api_key),
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
|
|
@ -101,19 +106,131 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
|
|||
model_list = response.json().get("data", [])
|
||||
return ["lemonade/" + model["id"] for model in model_list]
|
||||
|
||||
@staticmethod
|
||||
def _get_positive_int(value: Any) -> Optional[int]:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, int) and value > 0:
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
parsed = int(value)
|
||||
except ValueError:
|
||||
return None
|
||||
if parsed > 0:
|
||||
return parsed
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_provider_specific_entry(model_info: dict) -> dict:
|
||||
provider_specific_entry = model_info.get("provider_specific_entry")
|
||||
if not isinstance(provider_specific_entry, dict):
|
||||
provider_specific_entry = {}
|
||||
else:
|
||||
provider_specific_entry = provider_specific_entry.copy()
|
||||
|
||||
for key in ("recipe_options", "context_window", "max_context_window"):
|
||||
if key in model_info:
|
||||
provider_specific_entry[key] = model_info[key]
|
||||
|
||||
return provider_specific_entry
|
||||
|
||||
def _get_context_window(self, model_info: dict) -> Optional[int]:
|
||||
provider_specific_entry = self._get_provider_specific_entry(model_info)
|
||||
recipe_options = provider_specific_entry.get("recipe_options")
|
||||
if not isinstance(recipe_options, dict):
|
||||
recipe_options = {}
|
||||
|
||||
for value in (
|
||||
recipe_options.get("ctx_size"),
|
||||
model_info.get("max_input_tokens"),
|
||||
provider_specific_entry.get("context_window"),
|
||||
provider_specific_entry.get("max_context_window"),
|
||||
):
|
||||
parsed = self._get_positive_int(value)
|
||||
if parsed is not None:
|
||||
return parsed
|
||||
return None
|
||||
|
||||
def _get_default_model_info(self, model: str) -> dict:
|
||||
return {
|
||||
"key": "lemonade/" + model,
|
||||
"litellm_provider": "lemonade",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
"max_tokens": None,
|
||||
"max_input_tokens": None,
|
||||
"max_output_tokens": None,
|
||||
}
|
||||
|
||||
def get_model_info(
|
||||
self,
|
||||
model: str,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
) -> Any:
|
||||
if model.startswith("lemonade/"):
|
||||
model = model.split("/", 1)[1]
|
||||
|
||||
api_base, api_key = self._get_openai_compatible_provider_info(
|
||||
api_base=api_base, api_key=api_key
|
||||
)
|
||||
encoded_model = quote(model, safe="")
|
||||
|
||||
try:
|
||||
response = litellm.module_level_client.get(
|
||||
url=f"{api_base}/models/{encoded_model}",
|
||||
headers=self._get_auth_headers(api_key),
|
||||
)
|
||||
response.raise_for_status()
|
||||
model_info = response.json()
|
||||
except Exception:
|
||||
verbose_logger.debug("LemonadeError: Could not get model info.")
|
||||
return self._get_default_model_info(model)
|
||||
|
||||
max_input_tokens = self._get_context_window(model_info)
|
||||
max_output_tokens = self._get_positive_int(model_info.get("max_output_tokens"))
|
||||
max_tokens = self._get_positive_int(model_info.get("max_tokens"))
|
||||
provider_specific_entry = self._get_provider_specific_entry(model_info)
|
||||
|
||||
model_info_response = self._get_default_model_info(model)
|
||||
model_info_response.update(
|
||||
{
|
||||
"max_tokens": max_tokens or max_output_tokens,
|
||||
"max_input_tokens": max_input_tokens,
|
||||
"max_output_tokens": max_output_tokens,
|
||||
}
|
||||
)
|
||||
if provider_specific_entry:
|
||||
model_info_response["provider_specific_entry"] = provider_specific_entry
|
||||
return model_info_response
|
||||
|
||||
def _get_openai_compatible_provider_info(
|
||||
self, api_base: Optional[str], api_key: Optional[str]
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
# lemonade is openai compatible, we just need to set this to custom_openai and have the api_base be lemonade's endpoint
|
||||
passed_api_base = api_base
|
||||
api_base = (
|
||||
api_base
|
||||
or get_secret_str("LEMONADE_API_BASE")
|
||||
or "http://localhost:8000/api/v1"
|
||||
) # type: ignore
|
||||
# Lemonade doesn't check the key
|
||||
key = "lemonade"
|
||||
key = self._DEFAULT_API_KEY
|
||||
if passed_api_base is None or api_key:
|
||||
key = (
|
||||
api_key
|
||||
or litellm.lemonade_key
|
||||
or get_secret_str("LEMONADE_API_KEY")
|
||||
or self._DEFAULT_API_KEY
|
||||
)
|
||||
return api_base, key
|
||||
|
||||
def _get_auth_headers(self, api_key: Optional[str]) -> dict:
|
||||
if api_key is None or api_key == self._DEFAULT_API_KEY:
|
||||
return {}
|
||||
return {"Authorization": f"Bearer {api_key}"}
|
||||
|
||||
def transform_response(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from typing import List, Optional, Union
|
||||
from typing import Any, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -65,7 +65,8 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
|
||||
return (
|
||||
os.environ.get("OLLAMA_API_KEY")
|
||||
api_key
|
||||
or os.environ.get("OLLAMA_API_KEY")
|
||||
or litellm.api_key
|
||||
or litellm.openai_key
|
||||
or get_secret_str("OLLAMA_API_KEY")
|
||||
|
|
@ -78,13 +79,31 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
# env var OLLAMA_API_BASE or default
|
||||
return api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434"
|
||||
|
||||
@classmethod
|
||||
def get_server_api_base(cls, api_base: Optional[str] = None) -> str:
|
||||
api_base = cls.get_api_base(api_base).rstrip("/")
|
||||
for suffix in (
|
||||
"/api/generate",
|
||||
"/api/chat",
|
||||
"/api/embed",
|
||||
"/api/embeddings",
|
||||
"/api/show",
|
||||
"/api/tags",
|
||||
):
|
||||
if api_base.endswith(suffix):
|
||||
return api_base[: -len(suffix)]
|
||||
return api_base
|
||||
|
||||
def get_models(self, api_key=None, api_base: Optional[str] = None) -> List[str]:
|
||||
"""
|
||||
List all models available on the Ollama server via /api/tags endpoint.
|
||||
"""
|
||||
|
||||
base = self.get_api_base(api_base)
|
||||
api_key = self.get_api_key()
|
||||
passed_api_base = api_base
|
||||
base = self.get_server_api_base(api_base)
|
||||
api_key = (
|
||||
self.get_api_key(api_key) if passed_api_base is None or api_key else None
|
||||
)
|
||||
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
|
||||
names: set[str] = set()
|
||||
|
|
@ -126,6 +145,103 @@ class OllamaModelInfo(BaseLLMModelInfo):
|
|||
result = sorted(names)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _strip_ollama_model_prefix(model: str) -> str:
|
||||
if model.startswith("ollama/") or model.startswith("ollama_chat/"):
|
||||
return model.split("/", 1)[1]
|
||||
return model
|
||||
|
||||
@staticmethod
|
||||
def _is_static_ollama_model(model: str) -> bool:
|
||||
from litellm import model_cost
|
||||
|
||||
stripped_model = OllamaModelInfo._strip_ollama_model_prefix(model)
|
||||
potential_model_names = {
|
||||
model,
|
||||
stripped_model,
|
||||
"ollama/" + stripped_model,
|
||||
"ollama_chat/" + stripped_model,
|
||||
}
|
||||
model_cost_keys = {key.lower() for key in model_cost}
|
||||
return any(name.lower() in model_cost_keys for name in potential_model_names)
|
||||
|
||||
@staticmethod
|
||||
def _supports_function_calling(ollama_model_info: dict) -> bool:
|
||||
_template: str = str(ollama_model_info.get("template", "") or "")
|
||||
return "tools" in _template.lower()
|
||||
|
||||
@staticmethod
|
||||
def _get_max_tokens(ollama_model_info: dict) -> Optional[int]:
|
||||
_model_info: dict = ollama_model_info.get("model_info", {})
|
||||
|
||||
for key, value in _model_info.items():
|
||||
if "context_length" in key:
|
||||
return value
|
||||
return None
|
||||
|
||||
def get_runtime_model_info(
|
||||
self,
|
||||
model: str,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
) -> dict[str, Any]:
|
||||
from litellm import module_level_client
|
||||
|
||||
model = self._strip_ollama_model_prefix(model)
|
||||
passed_api_base = api_base
|
||||
api_base = self.get_server_api_base(api_base)
|
||||
api_key = (
|
||||
self.get_api_key(api_key) if passed_api_base is None or api_key else None
|
||||
)
|
||||
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
|
||||
try:
|
||||
response = module_level_client.post(
|
||||
url=f"{api_base}/api/show",
|
||||
json={"name": model},
|
||||
headers=headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except Exception:
|
||||
verbose_logger.debug("OllamaError: Could not get model info.")
|
||||
return {
|
||||
"key": model,
|
||||
"litellm_provider": "ollama",
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
"max_tokens": None,
|
||||
"max_input_tokens": None,
|
||||
"max_output_tokens": None,
|
||||
}
|
||||
|
||||
model_info = response.json()
|
||||
max_tokens = self._get_max_tokens(model_info)
|
||||
|
||||
return {
|
||||
"key": model,
|
||||
"litellm_provider": "ollama",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": self._supports_function_calling(model_info),
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
"max_tokens": max_tokens,
|
||||
"max_input_tokens": max_tokens,
|
||||
"max_output_tokens": max_tokens,
|
||||
}
|
||||
|
||||
def get_model_info(
|
||||
self,
|
||||
model: str,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
) -> Optional[dict[str, Any]]:
|
||||
if self._is_static_ollama_model(model):
|
||||
return None
|
||||
return self.get_runtime_model_info(
|
||||
model=model, api_base=api_base, api_key=api_key
|
||||
)
|
||||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional,
|
|||
from httpx._models import Headers, Response
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_str_from_messages,
|
||||
)
|
||||
|
|
@ -17,19 +17,17 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
)
|
||||
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock
|
||||
from litellm.types.utils import (
|
||||
Delta,
|
||||
GenericStreamingChunk,
|
||||
ModelInfoBase,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
ProviderField,
|
||||
StreamingChoices,
|
||||
)
|
||||
|
||||
from ..common_utils import OllamaError, _convert_image
|
||||
from ..common_utils import OllamaError, OllamaModelInfo, _convert_image
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
|
@ -224,59 +222,18 @@ class OllamaConfig(BaseConfig):
|
|||
)
|
||||
|
||||
def get_model_info(
|
||||
self, model: str, api_base: Optional[str] = None
|
||||
) -> ModelInfoBase:
|
||||
self,
|
||||
model: str,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
curl http://localhost:11434/api/show -d '{
|
||||
"name": "mistral"
|
||||
}'
|
||||
"""
|
||||
if model.startswith("ollama/") or model.startswith("ollama_chat/"):
|
||||
model = model.split("/", 1)[1]
|
||||
api_base = (
|
||||
api_base or get_secret_str("OLLAMA_API_BASE") or "http://localhost:11434"
|
||||
)
|
||||
api_key = self.get_api_key()
|
||||
headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
|
||||
try:
|
||||
response = litellm.module_level_client.post(
|
||||
url=f"{api_base}/api/show",
|
||||
json={"name": model},
|
||||
headers=headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
"OllamaError: Could not get model info for %s from %s. Error: %s",
|
||||
model,
|
||||
api_base,
|
||||
e,
|
||||
)
|
||||
return ModelInfoBase(
|
||||
key=model,
|
||||
litellm_provider="ollama",
|
||||
mode="chat",
|
||||
input_cost_per_token=0.0,
|
||||
output_cost_per_token=0.0,
|
||||
max_tokens=None,
|
||||
max_input_tokens=None,
|
||||
max_output_tokens=None,
|
||||
)
|
||||
|
||||
model_info = response.json()
|
||||
|
||||
_max_tokens: Optional[int] = self._get_max_tokens(model_info)
|
||||
|
||||
return ModelInfoBase(
|
||||
key=model,
|
||||
litellm_provider="ollama",
|
||||
mode="chat",
|
||||
supports_function_calling=self._supports_function_calling(model_info),
|
||||
input_cost_per_token=0.0,
|
||||
output_cost_per_token=0.0,
|
||||
max_tokens=_max_tokens,
|
||||
max_input_tokens=_max_tokens,
|
||||
max_output_tokens=_max_tokens,
|
||||
return OllamaModelInfo().get_model_info(
|
||||
model=model, api_base=api_base, api_key=api_key
|
||||
)
|
||||
|
||||
def get_error_class(
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ class PartnerModelPrefixes(str, Enum):
|
|||
MINIMAX_PREFIX = "minimaxai/"
|
||||
MOONSHOT_PREFIX = "moonshotai/"
|
||||
ZAI_PREFIX = "zai-org/"
|
||||
GEMMA_MAAS_PREFIX = "google/gemma-"
|
||||
|
||||
|
||||
class VertexAIPartnerModels(VertexBase):
|
||||
|
|
@ -68,6 +69,7 @@ class VertexAIPartnerModels(VertexBase):
|
|||
or model.startswith(PartnerModelPrefixes.MINIMAX_PREFIX)
|
||||
or model.startswith(PartnerModelPrefixes.MOONSHOT_PREFIX)
|
||||
or model.startswith(PartnerModelPrefixes.ZAI_PREFIX)
|
||||
or model.startswith(PartnerModelPrefixes.GEMMA_MAAS_PREFIX)
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
|
@ -82,6 +84,7 @@ class VertexAIPartnerModels(VertexBase):
|
|||
PartnerModelPrefixes.MINIMAX_PREFIX,
|
||||
PartnerModelPrefixes.MOONSHOT_PREFIX,
|
||||
PartnerModelPrefixes.ZAI_PREFIX,
|
||||
PartnerModelPrefixes.GEMMA_MAAS_PREFIX,
|
||||
]
|
||||
if any(provider in model for provider in OPENAI_LIKE_VERTEX_PROVIDERS):
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -34959,6 +34959,22 @@
|
|||
"us-central1"
|
||||
]
|
||||
},
|
||||
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-openai_models",
|
||||
"max_input_tokens": 256000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/maas/google/gemma-4-26b-a4b-it",
|
||||
"supported_regions": [
|
||||
"global"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"vertex_ai/openai/gpt-oss-120b-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-openai_models",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,37 @@
|
|||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .cato_networks import CatoNetworksGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
|
||||
import litellm
|
||||
from litellm.proxy.guardrails.guardrail_hooks.cato_networks import (
|
||||
CatoNetworksGuardrail,
|
||||
)
|
||||
|
||||
_cato_callback = CatoNetworksGuardrail(
|
||||
api_base=litellm_params.api_base,
|
||||
api_key=litellm_params.api_key,
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
ssl_verify=getattr(litellm_params, "ssl_verify", None),
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(_cato_callback)
|
||||
|
||||
return _cato_callback
|
||||
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.CATO_NETWORKS.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.CATO_NETWORKS.value: CatoNetworksGuardrail,
|
||||
}
|
||||
|
|
@ -0,0 +1,635 @@
|
|||
# +-------------------------------------------------------------+
|
||||
#
|
||||
# Use Cato Networks Guardrails for your LLM calls
|
||||
# https://www.catonetworks.com/
|
||||
#
|
||||
# +-------------------------------------------------------------+
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import os
|
||||
import ssl
|
||||
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Type, Union
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import BaseModel
|
||||
from websockets.asyncio.client import ClientConnection, connect
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
from litellm import DualCache
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._version import version as litellm_version
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
get_ssl_configuration,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails._content_utils import (
|
||||
apply_redacted_messages_back,
|
||||
build_inspection_messages,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CallTypesLiteral,
|
||||
Choices,
|
||||
EmbeddingResponse,
|
||||
ImageResponse,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
ResponsesAPIResponse,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
|
||||
class CatoNetworksGuardrailMissingSecrets(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class CatoNetworksGuardrail(CustomGuardrail):
|
||||
def __init__(
|
||||
self, api_key: Optional[str] = None, api_base: Optional[str] = None, **kwargs
|
||||
):
|
||||
ssl_verify = kwargs.pop("ssl_verify", None)
|
||||
self.async_handler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback,
|
||||
params={"ssl_verify": ssl_verify} if ssl_verify is not None else None,
|
||||
)
|
||||
self.api_key = api_key or os.environ.get("CATO_API_KEY")
|
||||
if not self.api_key:
|
||||
msg = (
|
||||
"Couldn't get Cato Networks api key, either set the `CATO_API_KEY` in the environment or "
|
||||
"pass it as a parameter to the guardrail in the config file"
|
||||
)
|
||||
raise CatoNetworksGuardrailMissingSecrets(msg)
|
||||
self.api_base = (
|
||||
api_base
|
||||
or os.environ.get("CATO_API_BASE")
|
||||
or "https://api.aisec.catonetworks.com"
|
||||
)
|
||||
self.api_base = self.api_base.rstrip("/")
|
||||
self.ws_api_base = self.api_base.replace("http://", "ws://").replace(
|
||||
"https://", "wss://"
|
||||
)
|
||||
self._ws_connect_ssl_kwargs = self._build_ws_ssl_kwargs(
|
||||
ssl_verify, self.ws_api_base
|
||||
)
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@staticmethod
|
||||
def _build_ws_ssl_kwargs(
|
||||
ssl_verify: Optional[Union[bool, str]], ws_api_base: str
|
||||
) -> dict:
|
||||
"""Resolve the ``ssl`` argument for ``websockets.connect``. Mirrors the
|
||||
``ssl_verify`` handling applied to the HTTP handler so a custom Cato instance
|
||||
behind TLS honours the same verification settings for streaming."""
|
||||
if ssl_verify is None or not ws_api_base.startswith("wss://"):
|
||||
return {}
|
||||
ssl_config = get_ssl_configuration(ssl_verify)
|
||||
if ssl_config is False:
|
||||
ssl_config = ssl.create_default_context()
|
||||
ssl_config.check_hostname = False
|
||||
ssl_config.verify_mode = ssl.CERT_NONE
|
||||
return {"ssl": ssl_config}
|
||||
|
||||
@staticmethod
|
||||
def _resolve_cato_user_email(user_api_key_dict: UserAPIKeyAuth) -> Optional[str]:
|
||||
"""Only the key/JWT-bound user email is trusted. ``end_user_id`` is derived from
|
||||
caller-supplied request fields (OpenAI ``user``, headers, metadata) and is spoofable,
|
||||
so it must never be forwarded as the Cato user identity."""
|
||||
return user_api_key_dict.user_email
|
||||
|
||||
@staticmethod
|
||||
async def _cancel_background_task(task: asyncio.Task) -> None:
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||||
await task
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> Union[Exception, str, dict, None]:
|
||||
verbose_proxy_logger.debug("Inside Cato Pre-Call Hook")
|
||||
return await self.call_cato_guardrail(
|
||||
data,
|
||||
hook="pre_call",
|
||||
key_alias=user_api_key_dict.key_alias,
|
||||
user_email=self._resolve_cato_user_email(user_api_key_dict),
|
||||
)
|
||||
|
||||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> Union[Exception, str, dict, None]:
|
||||
verbose_proxy_logger.debug("Inside Cato Moderation Hook")
|
||||
return await self.call_cato_guardrail(
|
||||
data,
|
||||
hook="moderation",
|
||||
key_alias=user_api_key_dict.key_alias,
|
||||
user_email=self._resolve_cato_user_email(user_api_key_dict),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _inspection_messages(cls, data: dict) -> list:
|
||||
"""Flatten multimodal list ``content`` into plain text so Cato inspects
|
||||
every text fragment. Chat ``messages`` stay 1:1 with the request so
|
||||
redacted results map back by index, and every other field the proxy
|
||||
forwards to the model (Responses-API ``input``/``instructions``, legacy
|
||||
completion ``prompt`` and tool/function/``response_format`` schema strings)
|
||||
is appended as synthetic messages so blocked text cannot bypass inspection
|
||||
by hiding in one of them."""
|
||||
flattened = []
|
||||
for message in data.get("messages") or []:
|
||||
if isinstance(message, dict) and isinstance(message.get("content"), list):
|
||||
parts = build_inspection_messages({"messages": [message]})
|
||||
flattened.append(
|
||||
{**message, "content": parts[0]["content"] if parts else ""}
|
||||
)
|
||||
else:
|
||||
flattened.append(message)
|
||||
for _field, messages in cls._extra_inspection_sources(data):
|
||||
flattened.extend(messages)
|
||||
return flattened
|
||||
|
||||
@staticmethod
|
||||
def _prompt_inspection_messages(prompt: Any) -> list:
|
||||
"""Synthetic user messages for a legacy completion ``prompt`` (a string
|
||||
or a list of string prompts)."""
|
||||
if isinstance(prompt, str):
|
||||
return [{"role": "user", "content": prompt}] if prompt else []
|
||||
if isinstance(prompt, list):
|
||||
return [
|
||||
{"role": "user", "content": part}
|
||||
for part in prompt
|
||||
if isinstance(part, str) and part
|
||||
]
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _iter_schema_string_refs(data: dict):
|
||||
"""Yield ``(container, key)`` for every non-empty schema string the proxy
|
||||
forwards to the model inside tool/function and structured-output schemas:
|
||||
each ``tools[].function`` and legacy ``functions[]`` entry plus the
|
||||
``response_format`` JSON schema, walked recursively for the free-text and
|
||||
value strings a caller could hide blocked text in (``description``,
|
||||
``title``, ``const``, ``default`` and every ``enum``/``examples`` item).
|
||||
Blocked text in any of them must be inspected and redacted like any other
|
||||
prompt."""
|
||||
scalar_keys = ("description", "title", "const", "default")
|
||||
list_keys = ("enum", "examples")
|
||||
|
||||
stack: list = []
|
||||
for tool in data.get("tools") or []:
|
||||
if isinstance(tool, dict) and isinstance(tool.get("function"), dict):
|
||||
stack.append(tool["function"])
|
||||
for function in data.get("functions") or []:
|
||||
if isinstance(function, dict):
|
||||
stack.append(function)
|
||||
response_format = data.get("response_format")
|
||||
if isinstance(response_format, dict):
|
||||
stack.append(response_format)
|
||||
stack.reverse()
|
||||
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
if isinstance(node, dict):
|
||||
for key in scalar_keys:
|
||||
value = node.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
yield node, key
|
||||
for key in list_keys:
|
||||
items = node.get(key)
|
||||
if isinstance(items, list):
|
||||
for idx, item in enumerate(items):
|
||||
if isinstance(item, str) and item:
|
||||
yield items, idx
|
||||
stack.extend(reversed(list(node.values())))
|
||||
elif isinstance(node, list):
|
||||
stack.extend(reversed(node))
|
||||
|
||||
@classmethod
|
||||
def _extra_inspection_sources(cls, data: dict) -> list:
|
||||
"""Text the proxy forwards to the model outside chat ``messages``:
|
||||
Responses-API ``input`` and ``instructions``, legacy completion
|
||||
``prompt`` and tool/function/``response_format`` schema strings. Returned
|
||||
as ``(field, messages)`` in a fixed order so the anonymize path can slice
|
||||
redactions back to the field they came from."""
|
||||
sources: list = []
|
||||
input_messages = build_inspection_messages({"input": data.get("input")})
|
||||
if input_messages:
|
||||
sources.append(("input", input_messages))
|
||||
instructions = data.get("instructions")
|
||||
if isinstance(instructions, str) and instructions:
|
||||
sources.append(
|
||||
("instructions", [{"role": "system", "content": instructions}])
|
||||
)
|
||||
prompt_messages = cls._prompt_inspection_messages(data.get("prompt"))
|
||||
if prompt_messages:
|
||||
sources.append(("prompt", prompt_messages))
|
||||
schema_strings = [
|
||||
{"role": "system", "content": container[key]}
|
||||
for container, key in cls._iter_schema_string_refs(data)
|
||||
]
|
||||
if schema_strings:
|
||||
sources.append(("schema_strings", schema_strings))
|
||||
return sources
|
||||
|
||||
async def call_cato_guardrail(
|
||||
self,
|
||||
data: dict,
|
||||
hook: str,
|
||||
key_alias: Optional[str],
|
||||
user_email: Optional[str] = None,
|
||||
) -> dict:
|
||||
call_id = data.get("litellm_call_id")
|
||||
headers = self._build_cato_headers(
|
||||
hook=hook,
|
||||
key_alias=key_alias,
|
||||
user_email=user_email,
|
||||
litellm_call_id=call_id,
|
||||
)
|
||||
response = await self.async_handler.post(
|
||||
f"{self.api_base}/fw/v1/analyze",
|
||||
headers=headers,
|
||||
json={"messages": self._inspection_messages(data)},
|
||||
)
|
||||
response.raise_for_status()
|
||||
res = response.json()
|
||||
required_action = res.get("required_action")
|
||||
action_type = required_action and required_action.get("action_type", None)
|
||||
if action_type is None:
|
||||
verbose_proxy_logger.debug("Cato: No required action specified")
|
||||
return data
|
||||
if action_type == "monitor_action":
|
||||
verbose_proxy_logger.info("Cato: monitor action")
|
||||
elif action_type == "block_action":
|
||||
self._handle_block_action(res.get("analysis_result", {}), required_action)
|
||||
elif action_type == "anonymize_action":
|
||||
return self._anonymize_request(res, data)
|
||||
else:
|
||||
verbose_proxy_logger.error(f"Cato: {action_type} action")
|
||||
return data
|
||||
|
||||
def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None:
|
||||
detection_message = required_action.get("detection_message", None)
|
||||
verbose_proxy_logger.info(
|
||||
"Cato: Violation detected enabled policies: {policies}".format(
|
||||
policies=list(analysis_result.get("policy_drill_down", {}).keys()),
|
||||
),
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=detection_message)
|
||||
|
||||
def _anonymize_request(self, res: Any, data: dict) -> dict:
|
||||
verbose_proxy_logger.info("Cato: anonymize action")
|
||||
redacted_chat = res.get("redacted_chat")
|
||||
if not redacted_chat:
|
||||
return data
|
||||
redacted_messages = redacted_chat.get("all_redacted_messages") or []
|
||||
original_messages = data.get("messages")
|
||||
offset = 0
|
||||
if original_messages:
|
||||
data["messages"] = [
|
||||
(
|
||||
{**original, "content": redacted_messages[idx]["content"]}
|
||||
if idx < len(redacted_messages)
|
||||
and redacted_messages[idx].get("content") is not None
|
||||
else original
|
||||
)
|
||||
for idx, original in enumerate(original_messages)
|
||||
]
|
||||
offset = len(original_messages)
|
||||
for field, messages in self._extra_inspection_sources(data):
|
||||
redacted_slice = redacted_messages[offset : offset + len(messages)]
|
||||
offset += len(messages)
|
||||
if redacted_slice:
|
||||
self._apply_extra_redaction(data, field, redacted_slice)
|
||||
return data
|
||||
|
||||
@classmethod
|
||||
def _apply_extra_redaction(cls, data: dict, field: str, redacted: list) -> None:
|
||||
if field == "input":
|
||||
input_only = {"input": data["input"]}
|
||||
apply_redacted_messages_back(input_only, redacted)
|
||||
data["input"] = input_only["input"]
|
||||
elif field == "instructions":
|
||||
if redacted[0].get("content") is not None:
|
||||
data["instructions"] = redacted[0]["content"]
|
||||
elif field == "prompt":
|
||||
cls._apply_prompt_redaction(data, redacted)
|
||||
elif field == "schema_strings":
|
||||
cls._apply_schema_string_redaction(data, redacted)
|
||||
|
||||
@classmethod
|
||||
def _apply_schema_string_redaction(cls, data: dict, redacted: list) -> None:
|
||||
redactions = iter(redacted)
|
||||
for container, key in cls._iter_schema_string_refs(data):
|
||||
replacement = next(redactions, None)
|
||||
if replacement is not None and replacement.get("content") is not None:
|
||||
container[key] = replacement["content"]
|
||||
|
||||
@staticmethod
|
||||
def _apply_prompt_redaction(data: dict, redacted: list) -> None:
|
||||
contents = [m.get("content") for m in redacted if isinstance(m, dict)]
|
||||
prompt = data.get("prompt")
|
||||
if isinstance(prompt, str):
|
||||
if contents and contents[0] is not None:
|
||||
data["prompt"] = contents[0]
|
||||
return
|
||||
if isinstance(prompt, list):
|
||||
new_prompt = list(prompt)
|
||||
redactions = iter(contents)
|
||||
for idx, part in enumerate(new_prompt):
|
||||
if isinstance(part, str) and part:
|
||||
replacement = next(redactions, None)
|
||||
if replacement is not None:
|
||||
new_prompt[idx] = replacement
|
||||
data["prompt"] = new_prompt
|
||||
|
||||
async def call_cato_guardrail_on_output(
|
||||
self,
|
||||
request_data: dict,
|
||||
output: str,
|
||||
hook: str,
|
||||
key_alias: Optional[str],
|
||||
user_email: Optional[str] = None,
|
||||
) -> Optional[dict]:
|
||||
call_id = request_data.get("litellm_call_id")
|
||||
inspection_messages = self._inspection_messages(request_data)
|
||||
assistant_index = len(inspection_messages)
|
||||
response = await self.async_handler.post(
|
||||
f"{self.api_base}/fw/v1/analyze",
|
||||
headers=self._build_cato_headers(
|
||||
hook=hook,
|
||||
key_alias=key_alias,
|
||||
user_email=user_email,
|
||||
litellm_call_id=call_id,
|
||||
),
|
||||
json={
|
||||
"messages": inspection_messages
|
||||
+ [{"role": "assistant", "content": output}]
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
res = response.json()
|
||||
required_action = res.get("required_action")
|
||||
action_type = required_action and required_action.get("action_type", None)
|
||||
if action_type and action_type == "block_action":
|
||||
self._handle_block_action_on_output(
|
||||
res.get("analysis_result", {}), required_action
|
||||
)
|
||||
redacted_chat = res.get("redacted_chat", None)
|
||||
|
||||
if action_type and action_type == "anonymize_action" and redacted_chat:
|
||||
all_redacted = redacted_chat.get("all_redacted_messages") or []
|
||||
if assistant_index < len(all_redacted):
|
||||
redacted_output = all_redacted[assistant_index].get("content")
|
||||
if redacted_output is not None:
|
||||
return {"redacted_output": redacted_output}
|
||||
return None
|
||||
|
||||
def _handle_block_action_on_output(
|
||||
self, analysis_result: Any, required_action: Any
|
||||
) -> None:
|
||||
detection_message = required_action.get("detection_message", None)
|
||||
verbose_proxy_logger.info(
|
||||
"Cato: detected: {detected}, enabled policies: {policies}".format(
|
||||
detected=True,
|
||||
policies=list(analysis_result.get("policy_drill_down", {}).keys()),
|
||||
),
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=detection_message)
|
||||
|
||||
def _build_cato_headers(
|
||||
self,
|
||||
*,
|
||||
hook: str,
|
||||
key_alias: Optional[str],
|
||||
user_email: Optional[str],
|
||||
litellm_call_id: Optional[str],
|
||||
):
|
||||
"""
|
||||
A helper function to build the http headers that are required by Cato guardrails.
|
||||
"""
|
||||
return (
|
||||
{
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
# Used by Cato Networks to apply only the guardrails that should be applied in a specific request phase.
|
||||
"x-cato-litellm-hook": hook,
|
||||
# Used by Cato Networks to track LiteLLM version and provide backward compatibility.
|
||||
"x-cato-litellm-version": litellm_version,
|
||||
}
|
||||
# Used by Cato Networks to track together single call input and output
|
||||
| ({"x-cato-call-id": litellm_call_id} if litellm_call_id else {})
|
||||
# Used by Cato Networks to track guardrails violations by user.
|
||||
| ({"x-cato-user-email": user_email} if user_email else {})
|
||||
| (
|
||||
{
|
||||
# Used by Cato Networks apply only the guardrails that are associated with the key alias.
|
||||
"x-cato-gateway-key-alias": key_alias,
|
||||
}
|
||||
if key_alias
|
||||
else {}
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _output_fragments(message: Any) -> list:
|
||||
"""Assistant text the proxy returns to the caller: ``content`` plus every
|
||||
``tool_calls[].function.arguments`` string, each tagged with where a
|
||||
redaction must be written back. ``content`` is only included when present
|
||||
so a tool-call-only choice keeps its ``None`` content (the text-vs-tool-call
|
||||
signal downstream consumers rely on) while its arguments are still inspected."""
|
||||
fragments: list = []
|
||||
if message.content is not None:
|
||||
fragments.append((("content", None), message.content))
|
||||
for idx, tool_call in enumerate(message.tool_calls or []):
|
||||
function = getattr(tool_call, "function", None)
|
||||
arguments = getattr(function, "arguments", None)
|
||||
if isinstance(arguments, str) and arguments:
|
||||
fragments.append((("tool_call", idx), arguments))
|
||||
return fragments
|
||||
|
||||
@staticmethod
|
||||
def _apply_output_fragment(message: Any, target: tuple, redacted: str) -> None:
|
||||
kind, idx = target
|
||||
if kind == "content":
|
||||
message.content = redacted
|
||||
else:
|
||||
message.tool_calls[idx].function.arguments = redacted
|
||||
|
||||
@staticmethod
|
||||
def _responses_output_field(item: Any, key: str) -> Any:
|
||||
return item.get(key) if isinstance(item, dict) else getattr(item, key, None)
|
||||
|
||||
@classmethod
|
||||
def _responses_output_fragments(cls, response: ResponsesAPIResponse) -> list:
|
||||
"""Assistant text the Responses API returns to the caller: every
|
||||
``output_text`` content block plus every function-call ``arguments``
|
||||
string, each paired with the ``(container, key)`` a Cato redaction is
|
||||
written back to. Output items and their content may be pydantic objects
|
||||
or plain dicts, so both access patterns are handled."""
|
||||
fragments: list = []
|
||||
for item in response.output or []:
|
||||
item_type = cls._responses_output_field(item, "type")
|
||||
if item_type == "function_call":
|
||||
arguments = cls._responses_output_field(item, "arguments")
|
||||
if isinstance(arguments, str) and arguments:
|
||||
fragments.append((item, "arguments", arguments))
|
||||
elif item_type == "message":
|
||||
for content in cls._responses_output_field(item, "content") or []:
|
||||
if cls._responses_output_field(content, "type") != "output_text":
|
||||
continue
|
||||
text = cls._responses_output_field(content, "text")
|
||||
if isinstance(text, str) and text:
|
||||
fragments.append((content, "text", text))
|
||||
return fragments
|
||||
|
||||
@staticmethod
|
||||
def _apply_responses_output_fragment(
|
||||
container: Any, key: str, redacted: str
|
||||
) -> None:
|
||||
if isinstance(container, dict):
|
||||
container[key] = redacted
|
||||
else:
|
||||
setattr(container, key, redacted)
|
||||
|
||||
async def _inspect_output_text(
|
||||
self,
|
||||
data: dict,
|
||||
text: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
user_email: Optional[str],
|
||||
) -> Optional[str]:
|
||||
"""Run the Cato output guardrail on a single assistant text fragment.
|
||||
Raises on a block action and returns the redacted replacement, or
|
||||
``None`` when the fragment must be left unchanged."""
|
||||
cato_output_guardrail_result = await self.call_cato_guardrail_on_output(
|
||||
data,
|
||||
text,
|
||||
hook="output",
|
||||
key_alias=user_api_key_dict.key_alias,
|
||||
user_email=user_email,
|
||||
)
|
||||
if cato_output_guardrail_result:
|
||||
return cato_output_guardrail_result.get("redacted_output")
|
||||
return None
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse],
|
||||
) -> Any:
|
||||
user_email = self._resolve_cato_user_email(user_api_key_dict)
|
||||
if isinstance(response, ModelResponse) and response.choices:
|
||||
for choice in response.choices:
|
||||
if not isinstance(choice, Choices):
|
||||
continue
|
||||
for target, text in self._output_fragments(choice.message):
|
||||
redacted_output = await self._inspect_output_text(
|
||||
data, text, user_api_key_dict, user_email
|
||||
)
|
||||
if redacted_output is not None:
|
||||
self._apply_output_fragment(
|
||||
choice.message, target, redacted_output
|
||||
)
|
||||
elif isinstance(response, ResponsesAPIResponse):
|
||||
for container, key, text in self._responses_output_fragments(response):
|
||||
redacted_output = await self._inspect_output_text(
|
||||
data, text, user_api_key_dict, user_email
|
||||
)
|
||||
if redacted_output is not None:
|
||||
self._apply_responses_output_fragment(
|
||||
container, key, redacted_output
|
||||
)
|
||||
return response
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response,
|
||||
request_data: dict,
|
||||
) -> AsyncGenerator[ModelResponseStream, None]:
|
||||
from litellm.proxy.proxy_server import StreamingCallbackError
|
||||
|
||||
user_email = self._resolve_cato_user_email(user_api_key_dict)
|
||||
call_id = request_data.get("litellm_call_id")
|
||||
async with connect(
|
||||
f"{self.ws_api_base}/fw/v1/analyze/stream",
|
||||
additional_headers=self._build_cato_headers(
|
||||
hook="output",
|
||||
key_alias=user_api_key_dict.key_alias,
|
||||
user_email=user_email,
|
||||
litellm_call_id=call_id,
|
||||
),
|
||||
**self._ws_connect_ssl_kwargs,
|
||||
) as websocket:
|
||||
sender = asyncio.create_task(
|
||||
self.forward_the_stream_to_cato(websocket, response)
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
raw_message = await self._await_cato_message(websocket, sender)
|
||||
result = json.loads(raw_message)
|
||||
if verified_chunk := result.get("verified_chunk"):
|
||||
yield ModelResponseStream.model_validate(verified_chunk)
|
||||
continue
|
||||
if result.get("done"):
|
||||
return
|
||||
if blocking_message := result.get("blocking_message"):
|
||||
raise StreamingCallbackError(blocking_message)
|
||||
verbose_proxy_logger.error(
|
||||
f"Unknown message received from Cato: {result}"
|
||||
)
|
||||
return
|
||||
finally:
|
||||
await self._cancel_background_task(sender)
|
||||
|
||||
async def _await_cato_message(
|
||||
self, websocket: ClientConnection, sender: asyncio.Task
|
||||
) -> Any:
|
||||
"""Wait for the next Cato message, surfacing a dead forwarding task instead of blocking."""
|
||||
from litellm.proxy.proxy_server import StreamingCallbackError
|
||||
|
||||
recv_task = asyncio.ensure_future(websocket.recv())
|
||||
pending = {recv_task, sender} if not sender.done() else {recv_task}
|
||||
await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
|
||||
if sender.done() and (sender_exc := sender.exception()) is not None:
|
||||
await self._cancel_background_task(recv_task)
|
||||
raise StreamingCallbackError(
|
||||
"Cato guardrail upstream stream failed"
|
||||
) from sender_exc
|
||||
try:
|
||||
return await recv_task
|
||||
except ConnectionClosed as exc:
|
||||
raise StreamingCallbackError(
|
||||
"Cato guardrail connection closed unexpectedly"
|
||||
) from exc
|
||||
|
||||
async def forward_the_stream_to_cato(
|
||||
self,
|
||||
websocket: ClientConnection,
|
||||
response_iter: AsyncGenerator[Any, None],
|
||||
) -> None:
|
||||
async for chunk in response_iter:
|
||||
if isinstance(chunk, BaseModel):
|
||||
chunk = chunk.model_dump_json()
|
||||
elif not isinstance(chunk, (str, bytes)):
|
||||
chunk = json.dumps(chunk)
|
||||
await websocket.send(chunk)
|
||||
await websocket.send(json.dumps({"done": True}))
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.cato_networks import (
|
||||
CatoNetworksGuardrailConfigModel,
|
||||
)
|
||||
|
||||
return CatoNetworksGuardrailConfigModel
|
||||
|
|
@ -67,6 +67,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
HIDE_SECRETS = "hide-secrets"
|
||||
HIDDENLAYER = "hiddenlayer"
|
||||
AIM = "aim"
|
||||
CATO_NETWORKS = "cato_networks"
|
||||
PANGEA = "pangea"
|
||||
CROWDSTRIKE_AIDR = "crowdstrike_aidr"
|
||||
LASSO = "lasso"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,20 @@
|
|||
from typing import Optional
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class CatoNetworksGuardrailConfigModel(GuardrailConfigModel):
|
||||
api_key: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The API key for the Cato Networks guardrail. If not provided, the `CATO_API_KEY` environment variable is checked.",
|
||||
)
|
||||
api_base: Optional[str] = Field(
|
||||
default=None,
|
||||
description="The API base for the Cato Networks guardrail. Default is https://api.aisec.catonetworks.com. Also checks if the `CATO_API_BASE` environment variable is set.",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Cato Networks Guardrail"
|
||||
109
litellm/utils.py
109
litellm/utils.py
|
|
@ -5443,7 +5443,7 @@ def _invalidate_model_cost_lowercase_map() -> None:
|
|||
_model_cost_mutation_generation += 1
|
||||
|
||||
# Clear LRU caches that depend on model_cost data
|
||||
get_model_info.cache_clear()
|
||||
_cached_get_model_info.cache_clear()
|
||||
_cached_get_model_info_helper.cache_clear()
|
||||
|
||||
|
||||
|
|
@ -5680,7 +5680,9 @@ def _cached_get_model_info_helper(
|
|||
Speed Optimization to hit high RPS
|
||||
"""
|
||||
return _get_model_info_helper(
|
||||
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -5720,6 +5722,7 @@ def _get_model_info_helper( # noqa: PLR0915
|
|||
model: str,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
) -> ModelInfoBase:
|
||||
"""
|
||||
Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's
|
||||
|
|
@ -5754,6 +5757,31 @@ def _get_model_info_helper( # noqa: PLR0915
|
|||
split_model = potential_model_names["split_model"]
|
||||
custom_llm_provider = potential_model_names["custom_llm_provider"]
|
||||
#########################
|
||||
provider_config: Optional[BaseLLMModelInfo] = None
|
||||
if custom_llm_provider and custom_llm_provider in LlmProvidersSet:
|
||||
provider_config = ProviderConfigManager.get_provider_model_info(
|
||||
model=model, provider=LlmProviders(custom_llm_provider)
|
||||
)
|
||||
if provider_config is not None:
|
||||
provider_get_model_info = getattr(provider_config, "get_model_info", None)
|
||||
if callable(provider_get_model_info):
|
||||
try:
|
||||
provider_model_info = provider_get_model_info(
|
||||
model=model,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
if provider_model_info is not None:
|
||||
return provider_model_info
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"Could not get dynamic model info for model=%s, provider=%s; "
|
||||
"falling back to the static cost map: %s",
|
||||
model,
|
||||
custom_llm_provider,
|
||||
e,
|
||||
)
|
||||
|
||||
if custom_llm_provider == "huggingface":
|
||||
max_tokens = _get_max_position_embeddings(model_name=model)
|
||||
return ModelInfoBase(
|
||||
|
|
@ -5774,10 +5802,6 @@ def _get_model_info_helper( # noqa: PLR0915
|
|||
supports_computer_use=None,
|
||||
supports_pdf_input=None,
|
||||
)
|
||||
elif (
|
||||
custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat"
|
||||
) and not _is_potential_model_name_in_model_cost(potential_model_names):
|
||||
return litellm.OllamaConfig().get_model_info(model, api_base=api_base)
|
||||
else:
|
||||
"""
|
||||
Check if: (in order of specificity)
|
||||
|
|
@ -6064,11 +6088,53 @@ def _get_model_info_helper( # noqa: PLR0915
|
|||
)
|
||||
|
||||
|
||||
def _build_model_info(
|
||||
model: str,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
) -> ModelInfo:
|
||||
supported_openai_params = litellm.get_supported_openai_params(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
|
||||
_model_info = _get_model_info_helper(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
provider_info = get_provider_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
if provider_info:
|
||||
for key, value in provider_info.items():
|
||||
if value is not None:
|
||||
_model_info[key] = value # type: ignore
|
||||
|
||||
# if verbose_logger.isEnabledFor(logging.DEBUG):
|
||||
# verbose_logger.debug(f"model_info: {_model_info}")
|
||||
|
||||
return ModelInfo(**_model_info, supported_openai_params=supported_openai_params)
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _cached_get_model_info(
|
||||
model: str,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
) -> ModelInfo:
|
||||
return _build_model_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider, api_base=api_base
|
||||
)
|
||||
|
||||
|
||||
def get_model_info(
|
||||
model: str,
|
||||
custom_llm_provider: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
api_key: Optional[str] = None,
|
||||
) -> ModelInfo:
|
||||
"""
|
||||
Get a dict for the maximum tokens (context window), input_cost_per_token, output_cost_per_token for a given model.
|
||||
|
|
@ -6140,32 +6206,15 @@ def get_model_info(
|
|||
"supported_openai_params": ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]
|
||||
}
|
||||
"""
|
||||
supported_openai_params = litellm.get_supported_openai_params(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
# api_key is a per-caller credential, not part of the model identity, so it is
|
||||
# kept out of the cache key; explicit keys are resolved without the cache.
|
||||
if api_key is not None:
|
||||
return _build_model_info(model, custom_llm_provider, api_base, api_key)
|
||||
return _cached_get_model_info(model, custom_llm_provider, api_base)
|
||||
|
||||
_model_info = _get_model_info_helper(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
)
|
||||
|
||||
provider_info = get_provider_info(
|
||||
model=model, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
if provider_info:
|
||||
for key, value in provider_info.items():
|
||||
if value is not None:
|
||||
_model_info[key] = value # type: ignore
|
||||
|
||||
# if verbose_logger.isEnabledFor(logging.DEBUG):
|
||||
# verbose_logger.debug(f"model_info: {_model_info}")
|
||||
|
||||
returned_model_info = ModelInfo(
|
||||
**_model_info, supported_openai_params=supported_openai_params
|
||||
)
|
||||
|
||||
return returned_model_info
|
||||
get_model_info.cache_clear = _cached_get_model_info.cache_clear # type: ignore[attr-defined]
|
||||
get_model_info.cache_info = _cached_get_model_info.cache_info # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def json_schema_type(python_type_name: str):
|
||||
|
|
|
|||
|
|
@ -34843,6 +34843,22 @@
|
|||
"us-central1"
|
||||
]
|
||||
},
|
||||
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-openai_models",
|
||||
"max_input_tokens": 256000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6e-07,
|
||||
"source": "https://cloud.google.com/vertex-ai/generative-ai/docs/maas/google/gemma-4-26b-a4b-it",
|
||||
"supported_regions": [
|
||||
"global"
|
||||
],
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"vertex_ai/openai/gpt-oss-120b-maas": {
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "vertex_ai-openai_models",
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import (
|
|||
exception_type,
|
||||
extract_and_raise_litellm_exception,
|
||||
)
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
|
||||
# Test cases for is_error_str_context_window_exceeded
|
||||
# Tuple format: (error_message, expected_result)
|
||||
|
|
@ -41,6 +42,10 @@ context_window_test_cases = [
|
|||
"`inputs` tokens + `max_new_tokens` must be <= 4096",
|
||||
True,
|
||||
),
|
||||
(
|
||||
"request (67311 tokens) exceeds the available context size (65536 tokens), try increasing it",
|
||||
True,
|
||||
),
|
||||
# Gemini 2.5/3 format
|
||||
(
|
||||
"The input token count exceeds the maximum number of tokens allowed 1048576.",
|
||||
|
|
@ -182,7 +187,6 @@ class TestExceptionCheckers:
|
|||
]
|
||||
|
||||
for error_str in positive_cases:
|
||||
print("testing positive case=", error_str)
|
||||
result = ExceptionCheckers.is_azure_content_policy_violation_error(
|
||||
error_str
|
||||
)
|
||||
|
|
@ -255,6 +259,33 @@ def test_gemini_context_window_error_mapping(
|
|||
)
|
||||
|
||||
|
||||
def test_lemonade_context_window_error_mapping():
|
||||
"""Lemonade's llama.cpp backend should map context overflows to LiteLLM's standard error."""
|
||||
|
||||
model = "lemonade/Qwen3.6-35B-A3B-GGUF"
|
||||
error_message = (
|
||||
'{"error":{"code":"context_length_exceeded","message":"request '
|
||||
"(80010 tokens) exceeds the available context size (65536 tokens), "
|
||||
'try increasing it","status_code":400,"type":"invalid_request_error"}}'
|
||||
)
|
||||
original_exception = OpenAIError(
|
||||
status_code=400,
|
||||
message=error_message,
|
||||
headers={},
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.ContextWindowExceededError) as excinfo:
|
||||
exception_type(
|
||||
model=model,
|
||||
original_exception=original_exception,
|
||||
custom_llm_provider="lemonade",
|
||||
)
|
||||
|
||||
assert excinfo.value.status_code == 400
|
||||
assert excinfo.value.llm_provider == "lemonade"
|
||||
assert excinfo.value.model == model
|
||||
|
||||
|
||||
# Test cases for Vertex AI RateLimitError mapping
|
||||
# As per https://github.com/BerriAI/litellm/issues/16189
|
||||
vertex_rate_limit_test_cases = [
|
||||
|
|
|
|||
|
|
@ -1,17 +1,14 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.llms.lemonade.chat.transformation import LemonadeChatConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
import httpx
|
||||
|
||||
|
||||
def test_lemonade_config_initialization():
|
||||
|
|
@ -28,8 +25,11 @@ def test_lemonade_config_initialization():
|
|||
assert config.repeat_penalty == 1.1
|
||||
|
||||
|
||||
def test_get_openai_compatible_provider_info():
|
||||
def test_get_openai_compatible_provider_info(monkeypatch):
|
||||
"""Test the provider info method returns correct API base and key"""
|
||||
monkeypatch.delenv("LEMONADE_API_KEY", raising=False)
|
||||
monkeypatch.setattr(litellm, "lemonade_key", None)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
api_base, key = config._get_openai_compatible_provider_info(
|
||||
|
|
@ -40,8 +40,11 @@ def test_get_openai_compatible_provider_info():
|
|||
assert key == "lemonade"
|
||||
|
||||
|
||||
def test_get_openai_compatible_provider_info_with_custom_base():
|
||||
def test_get_openai_compatible_provider_info_with_custom_base(monkeypatch):
|
||||
"""Test the provider info method with custom API base"""
|
||||
monkeypatch.delenv("LEMONADE_API_KEY", raising=False)
|
||||
monkeypatch.setattr(litellm, "lemonade_key", None)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
custom_api_base = "https://custom.lemonade.ai/v1"
|
||||
|
|
@ -53,6 +56,335 @@ def test_get_openai_compatible_provider_info_with_custom_base():
|
|||
assert key == "lemonade"
|
||||
|
||||
|
||||
def test_get_openai_compatible_provider_info_with_api_key_env(monkeypatch):
|
||||
"""Test the provider info method reads Lemonade's API key from the environment."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "test-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", None)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
api_base, key = config._get_openai_compatible_provider_info(
|
||||
api_base=None, api_key=None
|
||||
)
|
||||
|
||||
assert api_base == "http://localhost:8000/api/v1"
|
||||
assert key == "test-key"
|
||||
|
||||
|
||||
def test_get_openai_compatible_provider_info_skips_env_key_for_custom_base(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Test that caller-supplied bases do not receive server-side Lemonade keys."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
api_base, key = config._get_openai_compatible_provider_info(
|
||||
api_base="https://attacker.example/v1", api_key=None
|
||||
)
|
||||
|
||||
assert api_base == "https://attacker.example/v1"
|
||||
assert key == "lemonade"
|
||||
assert config._get_auth_headers(key) == {}
|
||||
|
||||
|
||||
def test_get_openai_compatible_provider_info_uses_explicit_key_for_custom_base(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Test that explicitly supplied Lemonade keys are sent to supplied bases."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
api_base, key = config._get_openai_compatible_provider_info(
|
||||
api_base="https://lemonade.example/v1", api_key="explicit-lemonade-key"
|
||||
)
|
||||
|
||||
assert api_base == "https://lemonade.example/v1"
|
||||
assert key == "explicit-lemonade-key"
|
||||
assert config._get_auth_headers(key) == {
|
||||
"Authorization": "Bearer explicit-lemonade-key"
|
||||
}
|
||||
|
||||
|
||||
def test_get_openai_compatible_provider_info_empty_key_does_not_leak_to_custom_base(
|
||||
monkeypatch,
|
||||
):
|
||||
"""An empty explicit key must not fall back to server-side Lemonade creds for a custom base."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
api_base, key = config._get_openai_compatible_provider_info(
|
||||
api_base="https://attacker.example/v1", api_key=""
|
||||
)
|
||||
|
||||
assert api_base == "https://attacker.example/v1"
|
||||
assert key == "lemonade"
|
||||
assert config._get_auth_headers(key) == {}
|
||||
|
||||
|
||||
def test_get_openai_compatible_provider_info_ignores_global_api_key(monkeypatch):
|
||||
"""Test that Lemonade discovery does not send unrelated global API keys."""
|
||||
monkeypatch.delenv("LEMONADE_API_KEY", raising=False)
|
||||
monkeypatch.setattr(litellm, "lemonade_key", None)
|
||||
monkeypatch.setattr(litellm, "api_key", "global-openai-key")
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
api_base, key = config._get_openai_compatible_provider_info(
|
||||
api_base="http://lemonade.test/v1", api_key=None
|
||||
)
|
||||
|
||||
assert api_base == "http://lemonade.test/v1"
|
||||
assert key == "lemonade"
|
||||
assert config._get_auth_headers(key) == {}
|
||||
|
||||
|
||||
def test_get_models_does_not_leak_lemonade_key_to_custom_base(monkeypatch):
|
||||
"""Test Lemonade discovery does not send server-side keys to supplied bases."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
|
||||
config = LemonadeChatConfig()
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {"data": []}
|
||||
|
||||
with patch.object(
|
||||
litellm.module_level_client, "get", return_value=response
|
||||
) as mock_get:
|
||||
models = config.get_models(api_base="https://attacker.example/v1")
|
||||
|
||||
assert models == []
|
||||
assert mock_get.call_args.kwargs["headers"] == {}
|
||||
|
||||
|
||||
def test_get_model_info_uses_loaded_context_size():
|
||||
"""Test that Lemonade model info prefers the effective loaded ctx_size."""
|
||||
config = LemonadeChatConfig()
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"id": "Qwen3.6-35B-A3B-GGUF",
|
||||
"recipe_options": {"ctx_size": 65536},
|
||||
"max_context_window": 262144,
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
litellm.module_level_client, "get", return_value=response
|
||||
) as mock_get:
|
||||
model_info = config.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
api_base="http://lemonade.test/v1",
|
||||
)
|
||||
|
||||
assert model_info["key"] == "lemonade/Qwen3.6-35B-A3B-GGUF"
|
||||
assert model_info["litellm_provider"] == "lemonade"
|
||||
assert model_info["max_input_tokens"] == 65536
|
||||
assert model_info["provider_specific_entry"] == {
|
||||
"recipe_options": {"ctx_size": 65536},
|
||||
"max_context_window": 262144,
|
||||
}
|
||||
assert "supports_function_calling" not in model_info
|
||||
assert "supports_response_schema" not in model_info
|
||||
assert "supports_tool_choice" not in model_info
|
||||
assert mock_get.call_args.kwargs["headers"] == {}
|
||||
|
||||
|
||||
def test_get_model_info_falls_back_when_server_unavailable():
|
||||
"""Test that Lemonade metadata lookup failures return safe defaults."""
|
||||
config = LemonadeChatConfig()
|
||||
|
||||
with patch.object(
|
||||
litellm.module_level_client, "get", side_effect=Exception("boom")
|
||||
):
|
||||
model_info = config.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
api_base="http://lemonade.test/v1",
|
||||
)
|
||||
|
||||
assert model_info["key"] == "lemonade/Qwen3.6-35B-A3B-GGUF"
|
||||
assert model_info["litellm_provider"] == "lemonade"
|
||||
assert model_info["mode"] == "chat"
|
||||
assert model_info["input_cost_per_token"] == 0.0
|
||||
assert model_info["output_cost_per_token"] == 0.0
|
||||
assert model_info["max_tokens"] is None
|
||||
assert model_info["max_input_tokens"] is None
|
||||
assert model_info["max_output_tokens"] is None
|
||||
assert "supports_function_calling" not in model_info
|
||||
assert "supports_response_schema" not in model_info
|
||||
assert "supports_tool_choice" not in model_info
|
||||
|
||||
|
||||
def test_get_model_info_reads_context_from_provider_specific_entry():
|
||||
"""Test that Lemonade model info uses provider-specific runtime metadata."""
|
||||
config = LemonadeChatConfig()
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"id": "Qwen3.6-35B-A3B-GGUF",
|
||||
"provider_specific_entry": {
|
||||
"recipe_options": {"ctx_size": "32768"},
|
||||
"max_context_window": 262144,
|
||||
},
|
||||
}
|
||||
|
||||
with patch.object(litellm.module_level_client, "get", return_value=response):
|
||||
model_info = config.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
api_base="http://lemonade.test/v1",
|
||||
)
|
||||
|
||||
assert model_info["max_input_tokens"] == 32768
|
||||
assert model_info["provider_specific_entry"] == {
|
||||
"recipe_options": {"ctx_size": "32768"},
|
||||
"max_context_window": 262144,
|
||||
}
|
||||
|
||||
|
||||
def test_get_model_info_sends_lemonade_api_key_for_configured_base(monkeypatch):
|
||||
"""Test that Lemonade model info uses auth for configured servers."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "test-key")
|
||||
monkeypatch.setenv("LEMONADE_API_BASE", "http://lemonade.test/v1")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", None)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"id": "Qwen3.6-35B-A3B-GGUF",
|
||||
"recipe_options": {"ctx_size": 65536},
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
litellm.module_level_client, "get", return_value=response
|
||||
) as mock_get:
|
||||
config.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
)
|
||||
|
||||
assert mock_get.call_args.kwargs["headers"] == {"Authorization": "Bearer test-key"}
|
||||
|
||||
|
||||
def test_get_model_info_sends_explicit_lemonade_api_key_for_custom_base(monkeypatch):
|
||||
"""Test that Lemonade model info sends explicitly supplied auth to supplied bases."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", None)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
config = LemonadeChatConfig()
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"id": "Qwen3.6-35B-A3B-GGUF",
|
||||
"recipe_options": {"ctx_size": 65536},
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
litellm.module_level_client, "get", return_value=response
|
||||
) as mock_get:
|
||||
config.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
api_base="http://lemonade.test/v1",
|
||||
api_key="explicit-test-key",
|
||||
)
|
||||
|
||||
assert mock_get.call_args.kwargs["headers"] == {
|
||||
"Authorization": "Bearer explicit-test-key"
|
||||
}
|
||||
|
||||
|
||||
def test_litellm_get_model_info_does_not_leak_lemonade_key_to_custom_base(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Test top-level model info does not send server-side keys to supplied bases."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"id": "Qwen3.6-35B-A3B-GGUF",
|
||||
"max_input_tokens": 65536,
|
||||
"max_context_window": 262144,
|
||||
}
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
with patch.object(
|
||||
litellm.module_level_client, "get", return_value=response
|
||||
) as mock_get:
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
api_base="https://attacker.example/v1",
|
||||
)
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert model_info["max_input_tokens"] == 65536
|
||||
assert mock_get.call_args.kwargs["headers"] == {}
|
||||
|
||||
|
||||
def test_litellm_get_model_info_forwards_explicit_lemonade_key_to_custom_base(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Top-level model info must forward an explicit api_key to the supplied base."""
|
||||
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"id": "Qwen3.6-35B-A3B-GGUF",
|
||||
"max_input_tokens": 65536,
|
||||
}
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
with patch.object(
|
||||
litellm.module_level_client, "get", return_value=response
|
||||
) as mock_get:
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
api_base="https://lemonade.example/v1",
|
||||
api_key="explicit-lemonade-key",
|
||||
)
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert model_info["max_input_tokens"] == 65536
|
||||
assert mock_get.call_args.kwargs["headers"] == {
|
||||
"Authorization": "Bearer explicit-lemonade-key"
|
||||
}
|
||||
|
||||
|
||||
def test_litellm_get_model_info_uses_lemonade_api_base():
|
||||
"""Test that LiteLLM model info is wired to Lemonade's model metadata API."""
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"id": "Qwen3.6-35B-A3B-GGUF",
|
||||
"max_input_tokens": 65536,
|
||||
"max_context_window": 262144,
|
||||
}
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
with patch.object(litellm.module_level_client, "get", return_value=response):
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
model="lemonade/Qwen3.6-35B-A3B-GGUF",
|
||||
api_base="http://lemonade.test/v1",
|
||||
)
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert model_info["max_input_tokens"] == 65536
|
||||
assert response.raise_for_status.called
|
||||
assert response.json.called
|
||||
|
||||
|
||||
def test_transform_response():
|
||||
"""Test the response transformation adds lemonade prefix to model name"""
|
||||
config = LemonadeChatConfig()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
import os
|
||||
import sys
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -23,6 +22,7 @@ if "httpx" not in sys.modules:
|
|||
sys.modules["httpx"] = httpx_mod
|
||||
|
||||
import httpx
|
||||
import litellm
|
||||
|
||||
from litellm.llms.ollama.common_utils import OllamaModelInfo
|
||||
|
||||
|
|
@ -105,6 +105,68 @@ class TestOllamaModelInfo:
|
|||
"Authorization": "Bearer test_api_key"
|
||||
}
|
||||
|
||||
def test_get_models_does_not_leak_server_key_to_provided_api_base(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""Model discovery should not send server-side keys to caller-supplied bases."""
|
||||
call_headers = []
|
||||
|
||||
def mock_get(url, headers):
|
||||
call_headers.append(headers)
|
||||
return DummyResponse({"models": []}, status_code=200)
|
||||
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
|
||||
monkeypatch.setattr(litellm, "openai_key", "global-openai-key")
|
||||
monkeypatch.setattr(httpx, "get", mock_get)
|
||||
|
||||
info = OllamaModelInfo()
|
||||
models = info.get_models(api_base="https://attacker.example")
|
||||
|
||||
assert models == []
|
||||
assert call_headers[0] == {}
|
||||
|
||||
def test_get_models_uses_explicit_api_key_for_provided_api_base(self, monkeypatch):
|
||||
"""Model discovery should send an explicitly supplied key to the provided base."""
|
||||
call_headers = []
|
||||
|
||||
def mock_get(url, headers):
|
||||
call_headers.append(headers)
|
||||
return DummyResponse({"models": []}, status_code=200)
|
||||
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key")
|
||||
monkeypatch.setattr(httpx, "get", mock_get)
|
||||
|
||||
info = OllamaModelInfo()
|
||||
models = info.get_models(
|
||||
api_base="https://ollama.example",
|
||||
api_key="explicit-api-key",
|
||||
)
|
||||
|
||||
assert models == []
|
||||
assert call_headers[0] == {"Authorization": "Bearer explicit-api-key"}
|
||||
|
||||
def test_get_models_empty_key_does_not_leak_to_provided_api_base(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""An empty explicit key must not fall back to server-side creds for a custom base."""
|
||||
call_headers = []
|
||||
|
||||
def mock_get(url, headers):
|
||||
call_headers.append(headers)
|
||||
return DummyResponse({"models": []}, status_code=200)
|
||||
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
|
||||
monkeypatch.setattr(litellm, "openai_key", "global-openai-key")
|
||||
monkeypatch.setattr(httpx, "get", mock_get)
|
||||
|
||||
info = OllamaModelInfo()
|
||||
models = info.get_models(api_base="https://attacker.example", api_key="")
|
||||
|
||||
assert models == []
|
||||
assert call_headers[0] == {}
|
||||
|
||||
def test_get_models_from_list_response(self, monkeypatch):
|
||||
"""
|
||||
When the /api/tags endpoint returns a list of dicts,
|
||||
|
|
@ -190,7 +252,7 @@ class TestOllamaGetModelInfo:
|
|||
|
||||
config = OllamaConfig()
|
||||
result = config.get_model_info(
|
||||
"llama3", api_base="http://my-remote-server:11434"
|
||||
"my-custom-model", api_base="http://my-remote-server:11434"
|
||||
)
|
||||
|
||||
assert captured_urls[0] == "http://my-remote-server:11434/api/show"
|
||||
|
|
@ -200,6 +262,181 @@ class TestOllamaGetModelInfo:
|
|||
"""When no api_base is passed, should fall back to OLLAMA_API_BASE env var."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
captured_urls = []
|
||||
captured_headers = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_urls.append(url)
|
||||
captured_headers.append(headers)
|
||||
return DummyResponse({"template": "", "model_info": {}}, status_code=200)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", "http://env-server:11434")
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "env-api-key")
|
||||
|
||||
config = OllamaConfig()
|
||||
config.get_model_info("my-custom-model")
|
||||
|
||||
assert captured_urls[0] == "http://env-server:11434/api/show"
|
||||
assert captured_headers[0] == {"Authorization": "Bearer env-api-key"}
|
||||
|
||||
def test_get_model_info_uses_explicit_api_key_for_provided_api_base(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""When api_key is explicit, model info should send it to the provided api_base."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
captured_headers = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_headers.append(headers)
|
||||
return DummyResponse({"template": "", "model_info": {}}, status_code=200)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
|
||||
config = OllamaConfig()
|
||||
config.get_model_info(
|
||||
"my-custom-model",
|
||||
api_base="http://my-remote-server:11434",
|
||||
api_key="explicit-api-key",
|
||||
)
|
||||
|
||||
assert captured_headers[0] == {"Authorization": "Bearer explicit-api-key"}
|
||||
|
||||
def test_get_model_info_empty_key_does_not_leak_to_provided_api_base(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""An empty explicit key must not fall back to server-side creds for a custom base."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
captured_headers = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_headers.append(headers)
|
||||
return DummyResponse({"template": "", "model_info": {}}, status_code=200)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
|
||||
monkeypatch.setattr(litellm, "openai_key", "global-openai-key")
|
||||
|
||||
config = OllamaConfig()
|
||||
config.get_model_info(
|
||||
"my-custom-model",
|
||||
api_base="https://attacker.example",
|
||||
api_key="",
|
||||
)
|
||||
|
||||
assert captured_headers[0] == {}
|
||||
|
||||
def test_litellm_get_model_info_does_not_leak_server_key_to_provided_api_base(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""Global model info should not send server-side keys to caller-supplied bases."""
|
||||
captured_headers = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_headers.append(headers)
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} tools {{ .Prompt }}",
|
||||
"model_info": {"llama.context_length": 32768},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key")
|
||||
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
|
||||
monkeypatch.setattr(litellm, "openai_key", "global-openai-key")
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
"ollama/unknown-model",
|
||||
api_base="https://attacker.example",
|
||||
)
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert model_info["max_input_tokens"] == 32768
|
||||
assert captured_headers[0] == {}
|
||||
|
||||
def test_litellm_get_model_info_forwards_explicit_api_key_to_provided_base(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""An explicit api_key passed to litellm.get_model_info must reach the provided base."""
|
||||
captured_headers = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_headers.append(headers)
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} tools {{ .Prompt }}",
|
||||
"model_info": {"llama.context_length": 32768},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key")
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
"ollama/unknown-model",
|
||||
api_base="https://ollama.example",
|
||||
api_key="explicit-api-key",
|
||||
)
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert model_info["max_input_tokens"] == 32768
|
||||
assert captured_headers[0] == {"Authorization": "Bearer explicit-api-key"}
|
||||
|
||||
def test_litellm_get_model_info_does_not_cache_on_api_key(self, monkeypatch):
|
||||
"""Regression: api_key must not be part of the get_model_info cache key.
|
||||
|
||||
Distinct api_keys for the same (model, api_base) must not each create their
|
||||
own cache entry (which would churn the shared LRU cache), and every explicit
|
||||
key must still reach the backend rather than be served from a result cached
|
||||
with a different key.
|
||||
"""
|
||||
from litellm.utils import _cached_get_model_info
|
||||
|
||||
captured_headers = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_headers.append(headers)
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} tools {{ .Prompt }}",
|
||||
"model_info": {"llama.context_length": 32768},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
litellm.get_model_info.cache_clear()
|
||||
try:
|
||||
for api_key in ("key-one", "key-two", "key-three"):
|
||||
litellm.get_model_info(
|
||||
"ollama/unknown-model",
|
||||
api_base="https://ollama.example",
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
assert _cached_get_model_info.cache_info().currsize <= 1
|
||||
assert captured_headers == [
|
||||
{"Authorization": "Bearer key-one"},
|
||||
{"Authorization": "Bearer key-two"},
|
||||
{"Authorization": "Bearer key-three"},
|
||||
]
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
def test_get_model_info_normalizes_generate_api_base(self, monkeypatch):
|
||||
"""When completion passes the final generate URL, model info should use the server base."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
captured_urls = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
|
|
@ -207,12 +444,13 @@ class TestOllamaGetModelInfo:
|
|||
return DummyResponse({"template": "", "model_info": {}}, status_code=200)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
monkeypatch.setenv("OLLAMA_API_BASE", "http://env-server:11434")
|
||||
|
||||
config = OllamaConfig()
|
||||
config.get_model_info("llama3")
|
||||
config.get_model_info(
|
||||
"my-custom-model", api_base="http://localhost:11434/api/generate"
|
||||
)
|
||||
|
||||
assert captured_urls[0] == "http://env-server:11434/api/show"
|
||||
assert captured_urls[0] == "http://localhost:11434/api/show"
|
||||
|
||||
def test_get_model_info_graceful_fallback_on_connection_error(self, monkeypatch):
|
||||
"""When the Ollama server is unreachable, should return defaults instead of raising."""
|
||||
|
|
@ -225,14 +463,42 @@ class TestOllamaGetModelInfo:
|
|||
monkeypatch.delenv("OLLAMA_API_BASE", raising=False)
|
||||
|
||||
config = OllamaConfig()
|
||||
result = config.get_model_info("llama3", api_base="http://unreachable:11434")
|
||||
result = config.get_model_info(
|
||||
"my-custom-model", api_base="http://unreachable:11434"
|
||||
)
|
||||
|
||||
assert result["key"] == "llama3"
|
||||
assert result["key"] == "my-custom-model"
|
||||
assert result["litellm_provider"] == "ollama"
|
||||
assert result["input_cost_per_token"] == 0.0
|
||||
assert result["output_cost_per_token"] == 0.0
|
||||
assert result["max_tokens"] is None
|
||||
|
||||
def test_get_model_info_graceful_fallback_on_http_error_status(self, monkeypatch):
|
||||
"""A non-2xx /api/show response must fall back to defaults, not parse the error body."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} tools {{ .Prompt }}",
|
||||
"model_info": {"llama.context_length": 8192},
|
||||
},
|
||||
status_code=404,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
|
||||
config = OllamaConfig()
|
||||
result = config.get_model_info(
|
||||
"my-custom-model", api_base="http://localhost:11434"
|
||||
)
|
||||
|
||||
assert result["key"] == "my-custom-model"
|
||||
assert result["litellm_provider"] == "ollama"
|
||||
assert result["max_tokens"] is None
|
||||
assert result["max_input_tokens"] is None
|
||||
assert "supports_function_calling" not in result
|
||||
|
||||
def test_get_model_info_strips_ollama_prefix(self, monkeypatch):
|
||||
"""Should strip 'ollama/' or 'ollama_chat/' prefix from model name."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
|
@ -246,11 +512,72 @@ class TestOllamaGetModelInfo:
|
|||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
|
||||
config = OllamaConfig()
|
||||
config.get_model_info("ollama/llama3", api_base="http://localhost:11434")
|
||||
assert captured_json[0]["name"] == "llama3"
|
||||
config.get_model_info(
|
||||
"ollama/my-custom-model", api_base="http://localhost:11434"
|
||||
)
|
||||
assert captured_json[0]["name"] == "my-custom-model"
|
||||
|
||||
config.get_model_info("ollama_chat/llama3", api_base="http://localhost:11434")
|
||||
assert captured_json[1]["name"] == "llama3"
|
||||
config.get_model_info(
|
||||
"ollama_chat/my-custom-model", api_base="http://localhost:11434"
|
||||
)
|
||||
assert captured_json[1]["name"] == "my-custom-model"
|
||||
|
||||
def test_get_model_info_skips_network_for_static_model(self, monkeypatch):
|
||||
"""Statically-priced models must not trigger an /api/show network call."""
|
||||
from litellm.llms.ollama.completion.transformation import OllamaConfig
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
raise AssertionError("Static Ollama model should not query /api/show")
|
||||
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
|
||||
config = OllamaConfig()
|
||||
assert config.get_model_info("ollama/llama2") is None
|
||||
|
||||
def test_litellm_get_model_info_uses_provider_hook_for_unknown_model(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""Unmapped Ollama models should use the provider-level dynamic hook."""
|
||||
captured_json = []
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
captured_json.append(json)
|
||||
return DummyResponse(
|
||||
{
|
||||
"template": "{{ .System }} tools {{ .Prompt }}",
|
||||
"model_info": {"llama.context_length": 32768},
|
||||
},
|
||||
status_code=200,
|
||||
)
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
try:
|
||||
model_info = litellm.get_model_info(
|
||||
"ollama/unknown-model", api_base="http://localhost:11434"
|
||||
)
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert model_info["max_input_tokens"] == 32768
|
||||
assert model_info["supports_function_calling"] is True
|
||||
assert captured_json[0]["name"] == "unknown-model"
|
||||
|
||||
def test_litellm_get_model_info_keeps_static_map_for_known_model(self, monkeypatch):
|
||||
"""Mapped Ollama models should keep using the static model map."""
|
||||
|
||||
def mock_post(url, json, headers=None):
|
||||
raise AssertionError("Static Ollama model should not query /api/show")
|
||||
|
||||
litellm.get_model_info.cache_clear()
|
||||
monkeypatch.setattr("litellm.module_level_client.post", mock_post)
|
||||
try:
|
||||
model_info = litellm.get_model_info("ollama/llama2")
|
||||
finally:
|
||||
litellm.get_model_info.cache_clear()
|
||||
|
||||
assert model_info["key"] == "ollama/llama2"
|
||||
assert model_info["litellm_provider"] == "ollama"
|
||||
|
||||
|
||||
class TestOllamaAuthHeaders:
|
||||
|
|
|
|||
|
|
@ -1326,6 +1326,63 @@ def test_vertex_ai_zai_is_partner_model():
|
|||
assert VertexAIPartnerModels.is_vertex_partner_model("zai-org/glm-4.7-maas")
|
||||
|
||||
|
||||
def test_vertex_ai_gemma_maas_is_partner_model():
|
||||
"""
|
||||
Ensure Gemma MaaS models are detected as Vertex AI partner models so they
|
||||
route through the OpenAI-compatible /endpoints/openapi path (not the
|
||||
legacy non-gemini path or the vertex_ai/gemma/ predict-endpoint handler).
|
||||
"""
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
|
||||
VertexAIPartnerModels,
|
||||
)
|
||||
|
||||
assert VertexAIPartnerModels.is_vertex_partner_model(
|
||||
"google/gemma-4-26b-a4b-it-maas"
|
||||
)
|
||||
|
||||
|
||||
def test_vertex_ai_gemma_maas_uses_openai_handler():
|
||||
"""
|
||||
Ensure Gemma MaaS partner models re-use the OpenAI-format handler.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
|
||||
VertexAIPartnerModels,
|
||||
)
|
||||
|
||||
assert VertexAIPartnerModels.should_use_openai_handler(
|
||||
"google/gemma-4-26b-a4b-it-maas"
|
||||
)
|
||||
|
||||
|
||||
def test_vertex_ai_gemma_maas_routes_to_partner_models():
|
||||
"""
|
||||
Regression guard for owtaylor's worry that Gemma MaaS could be misrouted as
|
||||
a gemma model. get_vertex_ai_model_route must return PARTNER_MODELS, never
|
||||
GEMMA, MODEL_GARDEN, or NON_GEMINI.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
VertexAIModelRoute,
|
||||
get_vertex_ai_model_route,
|
||||
)
|
||||
|
||||
route = get_vertex_ai_model_route("google/gemma-4-26b-a4b-it-maas")
|
||||
assert route == VertexAIModelRoute.PARTNER_MODELS
|
||||
|
||||
|
||||
def test_vertex_ai_google_gemini_not_detected_as_gemma_maas():
|
||||
"""
|
||||
Negative: adding the "google/gemma-" prefix must not widen detection to
|
||||
other google/* models like google/gemini-* (which should keep flowing
|
||||
through the gemini route, not partner_models).
|
||||
"""
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import (
|
||||
VertexAIPartnerModels,
|
||||
)
|
||||
|
||||
assert not VertexAIPartnerModels.is_vertex_partner_model("google/gemini-1.5-pro")
|
||||
assert not VertexAIPartnerModels.should_use_openai_handler("google/gemini-1.5-pro")
|
||||
|
||||
|
||||
def test_build_vertex_schema_empty_properties():
|
||||
"""
|
||||
Test _build_vertex_schema handles empty properties objects correctly.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,441 @@
|
|||
"""
|
||||
Tests for Vertex AI Gemma MaaS models that route through the partner-models
|
||||
OpenAI-compatible path (https://aiplatform.googleapis.com/.../endpoints/openapi).
|
||||
|
||||
These tests verify that:
|
||||
1. The correct global URL is constructed (https://aiplatform.googleapis.com)
|
||||
2. get_vertex_region resolves to "global" when model_cost says so
|
||||
3. acompletion() goes through the OpenAI-compatible handler and hits
|
||||
/endpoints/openapi/chat/completions
|
||||
4. Function-calling payloads (tools + tool_choice) pass through unchanged
|
||||
5. Vision/image_url payloads pass through unchanged
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import litellm
|
||||
from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
from litellm.types.llms.vertex_ai import VertexPartnerProvider
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model-cost entry used by all tests that need the model to be known
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_GEMMA_MODEL_COST_ENTRY = {
|
||||
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
|
||||
"litellm_provider": "vertex_ai-openai_models",
|
||||
"max_input_tokens": 256000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"output_cost_per_token": 6e-07,
|
||||
"supported_regions": ["global"],
|
||||
"supports_function_calling": True,
|
||||
"supports_tool_choice": True,
|
||||
"supports_vision": True,
|
||||
}
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_litellm_http_client_cache():
|
||||
"""Ensure each test gets a fresh async HTTP client mock."""
|
||||
from litellm import in_memory_llm_clients_cache
|
||||
|
||||
in_memory_llm_clients_cache.flush_cache()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clean_vertex_env():
|
||||
"""Clear Google/Vertex AI environment variables before each test to prevent test isolation issues."""
|
||||
saved_env = {}
|
||||
env_vars_to_clear = [
|
||||
"GOOGLE_APPLICATION_CREDENTIALS",
|
||||
"GOOGLE_CLOUD_PROJECT",
|
||||
"VERTEXAI_PROJECT",
|
||||
"VERTEX_PROJECT",
|
||||
"VERTEX_LOCATION",
|
||||
"VERTEX_AI_PROJECT",
|
||||
]
|
||||
for var in env_vars_to_clear:
|
||||
if var in os.environ:
|
||||
saved_env[var] = os.environ[var]
|
||||
del os.environ[var]
|
||||
|
||||
yield
|
||||
|
||||
for var, value in saved_env.items():
|
||||
os.environ[var] = value
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests: region and URL construction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestVertexBaseGetVertexRegionGemma:
|
||||
"""Test the get_vertex_region method for Gemma MaaS via model_cost lookup."""
|
||||
|
||||
def test_global_model_no_user_region_returns_global(self):
|
||||
vertex_base = VertexBase()
|
||||
|
||||
with patch.dict(
|
||||
litellm.model_cost,
|
||||
{
|
||||
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
|
||||
"supported_regions": ["global"]
|
||||
}
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
result = vertex_base.get_vertex_region(
|
||||
vertex_region=None,
|
||||
model="google/gemma-4-26b-a4b-it-maas",
|
||||
)
|
||||
assert result == "global"
|
||||
|
||||
def test_global_model_with_unsupported_user_region_overrides(self):
|
||||
vertex_base = VertexBase()
|
||||
|
||||
with patch.dict(
|
||||
litellm.model_cost,
|
||||
{
|
||||
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
|
||||
"supported_regions": ["global"]
|
||||
}
|
||||
},
|
||||
clear=False,
|
||||
):
|
||||
result = vertex_base.get_vertex_region(
|
||||
vertex_region="us-central1",
|
||||
model="google/gemma-4-26b-a4b-it-maas",
|
||||
)
|
||||
assert result == "global"
|
||||
|
||||
|
||||
class TestCreateVertexURLGemma:
|
||||
"""Test that create_vertex_url produces the expected OpenAI-compatible URL.
|
||||
|
||||
Gemma MaaS models reach this code path via should_use_openai_handler(), which
|
||||
selects VertexPartnerProvider.llama for all OpenAI-compatible partners including
|
||||
Gemma. test_gemma_routes_through_openai_handler() guards that mapping so the
|
||||
URL-format tests below are meaningful regression guards for the Gemma path.
|
||||
"""
|
||||
|
||||
def test_gemma_routes_through_openai_handler(self):
|
||||
"""Gemma MaaS must be routed through the OpenAI-compatible handler.
|
||||
|
||||
This is what causes VertexPartnerProvider.llama to be selected downstream,
|
||||
which in turn generates the /endpoints/openapi URL shape. If this mapping
|
||||
ever changes, the URL-shape tests below become misleading.
|
||||
"""
|
||||
assert VertexAIPartnerModels.should_use_openai_handler(
|
||||
"google/gemma-4-26b-a4b-it-maas"
|
||||
), "Gemma MaaS must use the OpenAI-compatible handler (VertexPartnerProvider.llama path)"
|
||||
|
||||
def test_global_location_url_format(self):
|
||||
# VertexPartnerProvider.llama is correct: Gemma MaaS reaches create_vertex_url
|
||||
# via should_use_openai_handler() → partner = VertexPartnerProvider.llama.
|
||||
# See test_gemma_routes_through_openai_handler for the routing guard.
|
||||
url = VertexBase.create_vertex_url(
|
||||
vertex_location="global",
|
||||
vertex_project="test-project",
|
||||
partner=VertexPartnerProvider.llama,
|
||||
stream=False,
|
||||
model="google/gemma-4-26b-a4b-it-maas",
|
||||
)
|
||||
|
||||
assert url.startswith("https://aiplatform.googleapis.com")
|
||||
assert "global-aiplatform.googleapis.com" not in url
|
||||
assert "/locations/global/" in url
|
||||
assert url.endswith("/endpoints/openapi/chat/completions")
|
||||
|
||||
def test_regional_location_url_format(self):
|
||||
url = VertexBase.create_vertex_url(
|
||||
vertex_location="us-central1",
|
||||
vertex_project="test-project",
|
||||
partner=VertexPartnerProvider.llama,
|
||||
stream=False,
|
||||
model="google/gemma-4-26b-a4b-it-maas",
|
||||
)
|
||||
|
||||
assert url.startswith("https://us-central1-aiplatform.googleapis.com")
|
||||
assert "/locations/us-central1/" in url
|
||||
assert url.endswith("/endpoints/openapi/chat/completions")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Capability-flag tests: verify get_model_info surfaces the advertised flags
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_gemma_maas_supports_function_calling():
|
||||
"""supports_function_calling=true in model_cost must be surfaced by the utility."""
|
||||
with patch.dict(litellm.model_cost, _GEMMA_MODEL_COST_ENTRY, clear=False):
|
||||
assert (
|
||||
litellm.utils.supports_function_calling(
|
||||
model="vertex_ai/google/gemma-4-26b-a4b-it-maas"
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_gemma_maas_supports_vision():
|
||||
"""supports_vision=true in model_cost must be surfaced by the utility."""
|
||||
with patch.dict(litellm.model_cost, _GEMMA_MODEL_COST_ENTRY, clear=False):
|
||||
assert (
|
||||
litellm.utils.supports_vision(
|
||||
model="vertex_ai/google/gemma-4-26b-a4b-it-maas"
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Integration tests: verify payloads reach the global OpenAI endpoint
|
||||
#
|
||||
# Patch target note (P1): AsyncHTTPHandler is patched at its *definition* site
|
||||
# (litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler). This works
|
||||
# correctly because the client is created by get_async_httpx_client(), which is
|
||||
# also defined in http_handler.py and calls AsyncHTTPHandler(...) using the
|
||||
# module-local name — so the patch intercepts instantiation there.
|
||||
# llm_http_handler.py only imports the class for type annotations; it never
|
||||
# instantiates it directly. Confirmed: without the mock the test raises
|
||||
# AuthenticationError, proving the assertion would never silently pass against
|
||||
# an un-mocked real call.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_MOCK_RESPONSE_JSON = {
|
||||
"id": "chatcmpl-gemma-test",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "google/gemma-4-26b-a4b-it-maas",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I help you today?",
|
||||
},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_gemma_global_endpoint_url():
|
||||
"""
|
||||
End-to-end: acompletion on vertex_ai/google/gemma-4-26b-a4b-it-maas should
|
||||
POST to the global endpoints/openapi/chat/completions URL.
|
||||
"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {}
|
||||
mock_response.json.return_value = _MOCK_RESPONSE_JSON
|
||||
|
||||
mock_vertexai = MagicMock()
|
||||
mock_vertexai.preview = MagicMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler"
|
||||
) as mock_http_handler,
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token",
|
||||
return_value=("fake-token", "test-project"),
|
||||
),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{"vertexai": mock_vertexai, "vertexai.preview": mock_vertexai.preview},
|
||||
),
|
||||
patch.dict(
|
||||
litellm.model_cost,
|
||||
{
|
||||
"vertex_ai/google/gemma-4-26b-a4b-it-maas": {
|
||||
"supported_regions": ["global"]
|
||||
}
|
||||
},
|
||||
clear=False,
|
||||
),
|
||||
):
|
||||
mock_http_handler.return_value.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
response = await litellm.acompletion(
|
||||
model="vertex_ai/google/gemma-4-26b-a4b-it-maas",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
vertex_ai_project="test-project",
|
||||
)
|
||||
|
||||
mock_http_handler.return_value.post.assert_called_once()
|
||||
|
||||
call_args = mock_http_handler.return_value.post.call_args
|
||||
called_url = call_args.kwargs["url"]
|
||||
|
||||
assert called_url.startswith("https://aiplatform.googleapis.com")
|
||||
assert "global-aiplatform.googleapis.com" not in called_url
|
||||
assert "/locations/global/" in called_url
|
||||
assert "/endpoints/openapi/chat/completions" in called_url
|
||||
|
||||
assert response.model == "google/gemma-4-26b-a4b-it-maas"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_gemma_function_calling_passthrough():
|
||||
"""
|
||||
Tools and tool_choice defined in the acompletion call must appear in the
|
||||
JSON body POSTed to the global endpoints/openapi/chat/completions URL.
|
||||
|
||||
This confirms that supports_function_calling=true is backed by real
|
||||
pass-through behaviour and that callers gating on get_model_info won't
|
||||
silently send unsupported requests.
|
||||
"""
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Return the current weather for a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {}
|
||||
mock_response.json.return_value = _MOCK_RESPONSE_JSON
|
||||
|
||||
mock_vertexai = MagicMock()
|
||||
mock_vertexai.preview = MagicMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler"
|
||||
) as mock_http_handler,
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token",
|
||||
return_value=("fake-token", "test-project"),
|
||||
),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{"vertexai": mock_vertexai, "vertexai.preview": mock_vertexai.preview},
|
||||
),
|
||||
patch.dict(litellm.model_cost, _GEMMA_MODEL_COST_ENTRY, clear=False),
|
||||
):
|
||||
mock_http_handler.return_value.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
await litellm.acompletion(
|
||||
model="vertex_ai/google/gemma-4-26b-a4b-it-maas",
|
||||
messages=[{"role": "user", "content": "What's the weather in Paris?"}],
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
vertex_ai_project="test-project",
|
||||
)
|
||||
|
||||
mock_http_handler.return_value.post.assert_called_once()
|
||||
call_args = mock_http_handler.return_value.post.call_args
|
||||
|
||||
# Must route to the global OpenAI-compatible endpoint
|
||||
called_url = call_args.kwargs["url"]
|
||||
assert called_url.startswith("https://aiplatform.googleapis.com"), called_url
|
||||
assert "/endpoints/openapi/chat/completions" in called_url, called_url
|
||||
|
||||
# Tools and tool_choice must be forwarded in the request body
|
||||
body = json.loads(call_args.kwargs["data"])
|
||||
assert "tools" in body, f"'tools' key missing from request body: {body}"
|
||||
assert body["tools"][0]["function"]["name"] == "get_weather"
|
||||
assert "tool_choice" in body, f"'tool_choice' missing from request body: {body}"
|
||||
assert body["tool_choice"] == "auto"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_gemma_vision_passthrough():
|
||||
"""
|
||||
An image_url content part must survive transformation and appear in the
|
||||
JSON body POSTed to the global endpoints/openapi/chat/completions URL.
|
||||
|
||||
This confirms that supports_vision=true is backed by real pass-through
|
||||
behaviour and that callers gating on get_model_info won't silently send
|
||||
unsupported multimodal requests.
|
||||
"""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe this image."},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
},
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {}
|
||||
mock_response.json.return_value = _MOCK_RESPONSE_JSON
|
||||
|
||||
mock_vertexai = MagicMock()
|
||||
mock_vertexai.preview = MagicMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler"
|
||||
) as mock_http_handler,
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_ai_partner_models.main.VertexAIPartnerModels._ensure_access_token",
|
||||
return_value=("fake-token", "test-project"),
|
||||
),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{"vertexai": mock_vertexai, "vertexai.preview": mock_vertexai.preview},
|
||||
),
|
||||
patch.dict(litellm.model_cost, _GEMMA_MODEL_COST_ENTRY, clear=False),
|
||||
):
|
||||
mock_http_handler.return_value.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
await litellm.acompletion(
|
||||
model="vertex_ai/google/gemma-4-26b-a4b-it-maas",
|
||||
messages=messages,
|
||||
vertex_ai_project="test-project",
|
||||
)
|
||||
|
||||
mock_http_handler.return_value.post.assert_called_once()
|
||||
call_args = mock_http_handler.return_value.post.call_args
|
||||
|
||||
# Must still route to the global OpenAI-compatible endpoint
|
||||
called_url = call_args.kwargs["url"]
|
||||
assert called_url.startswith("https://aiplatform.googleapis.com"), called_url
|
||||
assert "/endpoints/openapi/chat/completions" in called_url, called_url
|
||||
|
||||
# The image_url content part must be present in the forwarded body
|
||||
body = json.loads(call_args.kwargs["data"])
|
||||
user_msg = next(m for m in body["messages"] if m["role"] == "user")
|
||||
content = user_msg["content"]
|
||||
assert isinstance(content, list), f"Expected list content, got: {content}"
|
||||
image_parts = [p for p in content if p.get("type") == "image_url"]
|
||||
assert image_parts, f"No image_url part in forwarded message content: {content}"
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -15,9 +15,11 @@ import pytest
|
|||
sys.path.insert(0, str(Path(__file__).parent))
|
||||
|
||||
import litellm.proxy.guardrails.guardrail_hooks.aim.aim as _aim_module
|
||||
import litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks as _cato_networks_module
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM
|
||||
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
|
||||
from litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks import CatoNetworksGuardrail
|
||||
|
||||
|
||||
class TestBaseAWSLLMSSLVerify:
|
||||
|
|
@ -144,6 +146,48 @@ class TestAimGuardrailSSLVerify:
|
|||
assert mock_get_client.called
|
||||
|
||||
|
||||
class TestCatoNetworksGuardrailSSLVerify:
|
||||
"""Test SSL verification parameter handling in CatoNetworksGuardrail."""
|
||||
|
||||
def test_init_accepts_ssl_verify(self):
|
||||
"""Test that CatoNetworksGuardrail.__init__ accepts and uses ssl_verify parameter."""
|
||||
mock_handler = Mock()
|
||||
|
||||
# Use patch.object on the actual module reference for reliable patching
|
||||
# across different import orders / CI environments
|
||||
with patch.object(
|
||||
_cato_networks_module, "get_async_httpx_client", return_value=mock_handler
|
||||
) as mock_get_client:
|
||||
# Initialize with ssl_verify
|
||||
cert_path = "/path/to/cato_cert.pem"
|
||||
CatoNetworksGuardrail(
|
||||
api_key="test_key",
|
||||
api_base="https://test.catonetworks.api",
|
||||
ssl_verify=cert_path,
|
||||
)
|
||||
|
||||
# Verify get_async_httpx_client was called with ssl_verify in params
|
||||
assert mock_get_client.called
|
||||
call_kwargs = mock_get_client.call_args[1]
|
||||
assert "params" in call_kwargs
|
||||
assert call_kwargs["params"] is not None
|
||||
assert call_kwargs["params"]["ssl_verify"] == cert_path
|
||||
|
||||
def test_init_without_ssl_verify(self):
|
||||
"""Test that CatoNetworksGuardrail works without ssl_verify parameter."""
|
||||
mock_handler = Mock()
|
||||
|
||||
# Use patch.object on the actual module reference for reliable patching
|
||||
with patch.object(
|
||||
_cato_networks_module, "get_async_httpx_client", return_value=mock_handler
|
||||
) as mock_get_client:
|
||||
# Initialize without ssl_verify
|
||||
CatoNetworksGuardrail(api_key="test_key", api_base="https://test.catonetworks.api")
|
||||
|
||||
# Should still work, just without custom SSL
|
||||
assert mock_get_client.called
|
||||
|
||||
|
||||
class TestHTTPHandlerSSLVerify:
|
||||
"""Test SSL verification parameter handling in HTTP handlers."""
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,4 @@
|
|||
<svg width="143" height="71" viewBox="0 0 143 71" fill="none" xmlns="http://www.w3.org/2000/svg">
|
||||
<path d="M57.0307 7.1665H55.4564L38.9813 48.4353H44.813L47.6069 40.8693H64.4145L67.2084 48.4353H73.0401L57.0307 7.1665ZM62.5741 35.8106H49.4695L56.2546 18.8815L62.5741 35.8106ZM69.8914 7.18869V12.2474H81.7543V48.4353H87.2312V12.2474H99.0942V7.18869H69.8914ZM14.2355 62.9681L11.7299 58.0203H10.089V64.7653H11.3751V59.884L13.8807 64.7653H15.5216V58.0203H14.2355V62.9681ZM26.9854 58.0203H26.5197V64.7875H31.2649V63.5449H27.8501V62.0362H30.7105V60.7937H27.8501V59.2406H31.2649V58.0203H27.8501H26.9854ZM42.2409 59.2406H44.17V64.7875H45.5004V59.2406H47.4295V58.0203H42.2409V59.2406ZM64.5032 62.5243L63.1506 58.0203H62.0641L60.7115 62.5243L59.8024 58.0203H58.4276L60.0685 64.7875H61.3102L62.6406 60.0393L63.8824 64.7875H65.1463L66.7871 58.0203H65.4123L64.5032 62.5243ZM81.71 58.2421C81.3109 58.0424 80.8674 57.9315 80.3574 57.9315C79.8474 57.9315 79.3818 58.0424 79.0048 58.2421C78.6057 58.4418 78.3174 58.7303 78.0957 59.1074C77.874 59.4846 77.7631 59.9284 77.7631 60.4165V62.3246C77.7631 62.8128 77.874 63.2565 78.0957 63.6337C78.3174 64.0109 78.6057 64.2993 79.0048 64.499C79.4039 64.6987 79.8474 64.8096 80.3574 64.8096C80.8674 64.8096 81.333 64.6987 81.71 64.499C82.1091 64.2993 82.3974 64.0109 82.6191 63.6337C82.8409 63.2565 82.9517 62.8128 82.9517 62.3246V60.4165C82.9517 59.9284 82.8409 59.4846 82.6191 59.1074C82.4196 58.7524 82.1091 58.464 81.71 58.2421ZM81.5548 62.3912C81.5548 62.6353 81.5104 62.8349 81.3996 63.0124C81.3109 63.1899 81.1557 63.3231 80.9783 63.434C80.8009 63.5228 80.5791 63.5893 80.3352 63.5893C80.0913 63.5893 79.8918 63.5449 79.6922 63.434C79.5148 63.3453 79.3596 63.1899 79.2709 63.0124C79.1822 62.8349 79.1157 62.6353 79.1157 62.3912V60.4165C79.1157 60.1724 79.16 59.9728 79.2709 59.7953C79.3596 59.6178 79.5148 59.4846 79.6922 59.3737C79.8696 59.285 80.0913 59.2184 80.3352 59.2184C80.5791 59.2184 80.7787 59.2628 80.9783 59.3737C81.1557 59.4624 81.3109 59.6178 81.3996 59.7953C81.4883 59.9728 81.5548 60.1724 81.5548 60.4165V62.3912ZM98.0742 61.8143C98.3624 61.6368 98.6063 61.4149 98.7616 61.1043C98.9168 60.7937 99.0055 60.4387 99.0055 60.0393C99.0055 59.6399 98.9168 59.2849 98.7616 58.9743C98.6063 58.6637 98.3624 58.4418 98.0742 58.2643C97.7859 58.0868 97.4311 58.0203 97.0542 58.0203H93.9277V64.7875H95.2581V62.0584H96.0785L97.675 64.7875H99.2937L97.6085 61.9918C97.7637 61.9474 97.919 61.8809 98.0742 61.8143ZM95.2803 59.2406H96.9433C97.0764 59.2406 97.2094 59.2628 97.2981 59.3293C97.3868 59.3959 97.4755 59.4846 97.542 59.6178C97.6085 59.7287 97.6307 59.884 97.6307 60.0393C97.6307 60.1946 97.6085 60.3278 97.542 60.4609C97.4755 60.5718 97.409 60.6828 97.2981 60.7493C97.1872 60.8159 97.0764 60.8381 96.9433 60.8381H95.2803V59.2406ZM115.614 58.0203H113.928L111.622 61.3706V58.0203H110.292V64.7875H111.622V63.3231L112.642 61.9696L114.327 64.7875H115.902L113.507 60.8159L115.614 58.0203ZM131.956 61.7034C131.756 61.4371 131.512 61.2374 131.202 61.1265C130.891 61.0156 130.536 60.9046 130.071 60.8159C130.049 60.8159 130.026 60.8159 130.004 60.7937C129.982 60.7937 129.96 60.7937 129.938 60.7715H129.849C129.539 60.7049 129.295 60.6606 129.117 60.594C128.94 60.5496 128.807 60.4609 128.674 60.3499C128.563 60.239 128.496 60.0837 128.496 59.884C128.496 59.6621 128.607 59.4624 128.829 59.3515C129.051 59.2184 129.361 59.1518 129.783 59.1518C130.049 59.1518 130.337 59.1962 130.647 59.3071C130.936 59.3959 131.224 59.5512 131.512 59.7287L132.066 58.6415C131.845 58.4862 131.601 58.3753 131.335 58.2643C131.091 58.1534 130.825 58.0868 130.559 58.0203C130.293 57.9537 130.026 57.9315 129.783 57.9315C129.228 57.9315 128.763 58.0203 128.363 58.1756C127.964 58.3309 127.676 58.5749 127.476 58.8856C127.277 59.1962 127.166 59.5734 127.166 60.0171C127.166 60.5053 127.277 60.8825 127.499 61.1709C127.72 61.4371 127.986 61.6368 128.297 61.7256C128.607 61.8365 129.006 61.9253 129.472 61.9918L129.583 62.014H129.627C129.893 62.0584 130.115 62.1028 130.293 62.1471C130.47 62.1915 130.603 62.2803 130.714 62.369C130.825 62.4799 130.869 62.6131 130.869 62.7906C130.869 63.0568 130.758 63.2565 130.514 63.3896C130.27 63.5228 129.938 63.6115 129.494 63.6115C129.117 63.6115 128.74 63.5449 128.386 63.434C128.031 63.3231 127.72 63.1456 127.432 62.9459L126.811 63.9887C127.033 64.1662 127.299 64.3215 127.565 64.4324C127.853 64.5656 128.164 64.6543 128.474 64.7209C128.807 64.7875 128.962 64.8096 128.962 64.8096H129.472C130.049 64.8096 130.536 64.7209 130.936 64.5656C131.335 64.4103 131.645 64.1662 131.867 63.8556C132.089 63.5449 132.177 63.1899 132.177 62.7462C132.266 62.3246 132.155 61.9696 131.956 61.7034Z" fill="#148964"/>
|
||||
<path d="M120.98 6.146C109.316 6.146 99.8037 15.4647 99.8037 27.4016C99.8037 39.3385 109.117 48.4354 120.758 48.4354C132.421 48.4354 141.934 39.1166 141.934 27.1797C141.934 15.3982 132.621 6.146 120.98 6.146ZM120.98 43.1104C112.243 43.1104 105.502 36.3876 105.502 27.1797C105.502 18.4379 112.11 11.471 120.758 11.471C129.494 11.471 136.235 18.1938 136.235 27.4016C136.213 36.2101 129.605 43.1104 120.98 43.1104ZM32.7062 38.3179C29.9566 41.291 26.0541 43.1104 21.6193 43.1104C12.8829 43.1104 6.1421 36.3876 6.1421 27.1797C6.1421 18.4379 12.7499 11.471 21.3976 11.471C25.921 11.471 29.9123 13.2682 32.7062 16.2857L36.5866 12.3807C32.7949 8.52006 27.4954 6.146 21.5972 6.146C9.9338 6.146 0.421295 15.4647 0.421295 27.4016C0.421295 39.3385 9.73424 48.4354 21.3754 48.4354C27.2958 48.4354 32.6618 46.0391 36.4979 42.1119L32.7062 38.3179Z" fill="#148964"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 5.4 KiB |
|
|
@ -301,6 +301,17 @@ const EditGuardrailForm: React.FC<EditGuardrailFormProps> = ({
|
|||
/>
|
||||
</Form.Item>
|
||||
);
|
||||
case "CatoNetworks":
|
||||
return (
|
||||
<Form.Item label="Cato Networks Configuration" name="config" tooltip="JSON configuration for Cato Networks">
|
||||
<Input.TextArea
|
||||
rows={4}
|
||||
placeholder={`{
|
||||
"api_key": "your_cato_api_key"
|
||||
}`}
|
||||
/>
|
||||
</Form.Item>
|
||||
);
|
||||
case "GuardrailsAI":
|
||||
return (
|
||||
<Form.Item label="Guardrails.ai Configuration" name="config" tooltip="JSON configuration for Guardrails.ai">
|
||||
|
|
|
|||
|
|
@ -228,6 +228,12 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
|
|||
mode: "pre_call",
|
||||
defaultOn: false,
|
||||
},
|
||||
cato_networks: {
|
||||
provider: "Cato Networks",
|
||||
guardrailNameSuggestion: "Cato Networks Guardrail",
|
||||
mode: "pre_call",
|
||||
defaultOn: false,
|
||||
},
|
||||
prompt_security: {
|
||||
provider: "PromptSecurity",
|
||||
guardrailNameSuggestion: "Prompt Security",
|
||||
|
|
|
|||
|
|
@ -325,6 +325,14 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [
|
|||
logo: `${ASSET_PREFIX}aim_security.jpeg`,
|
||||
tags: ["Security", "Threat Detection"],
|
||||
},
|
||||
{
|
||||
id: "cato_networks",
|
||||
name: "Cato Networks Guardrail",
|
||||
description: "Cato Networks guardrails for comprehensive AI threat detection and mitigation.",
|
||||
category: "partner",
|
||||
logo: `${ASSET_PREFIX}cato_networks.svg`,
|
||||
tags: ["Security", "Threat Detection"],
|
||||
},
|
||||
{
|
||||
id: "prompt_security",
|
||||
name: "Prompt Security",
|
||||
|
|
|
|||
|
|
@ -131,6 +131,7 @@ export const guardrailLogoMap: Record<string, string> = {
|
|||
"Lasso Guardrail": `${asset_logos_folder}lasso.png`,
|
||||
"Pangea Guardrail": `${asset_logos_folder}pangea.png`,
|
||||
"AIM Guardrail": `${asset_logos_folder}aim_security.jpeg`,
|
||||
"Cato Networks Guardrail": `${asset_logos_folder}cato_networks.svg`,
|
||||
"OpenAI Moderation": `${asset_logos_folder}openai_small.svg`,
|
||||
EnkryptAI: `${asset_logos_folder}enkrypt_ai.avif`,
|
||||
"Prompt Security": `${asset_logos_folder}prompt_security.png`,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue