From 1f7eeb274cad2dfb19b10184ede952bad6aa4b0e Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 21 Feb 2026 15:34:42 -0800 Subject: [PATCH] Agent Builder - improve rejected response detection based on agent response (#21850) * fix: feat: add litellm_system_prompt support * feat: support new 'litellm_agent' model provider * feat: ui/ - new agent builder ui * fix(anthropic/chat/transformation.py): normalize max_tokens if decimal * feat(agentbuilderview.tsx): run compliance datasets against litellm agent * feat: new response rejection detector * fix: multiple fixes * feat: add mcp tools support to agent builder create an agent with access to llm's + mcp servers --- .../guardrail_hooks/custom_code/__init__.py | 7 + .../custom_code/custom_code_guardrail.py | 18 +- .../custom_code/response_rejection_code.py | 76 ++++ .../policy_endpoints/endpoints.py | 285 +++++++++---- litellm/types/utils.py | 66 ++- .../test_response_rejection_guardrail_code.py | 78 ++++ .../src/components/networking.tsx | 16 +- .../playground/chat_ui/AgentBuilderView.tsx | 391 ++++++++++++++++-- .../playground/complianceUI/ComplianceUI.tsx | 194 ++++----- .../playground/llm_calls/fetch_agents.tsx | 12 + 10 files changed, 890 insertions(+), 253 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/custom_code/response_rejection_code.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py index 747b188fee..d166e66dba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/__init__.py @@ -2,6 +2,9 @@ This module allows users to write custom guardrail logic using Python-like code that runs in a sandboxed environment with access to LiteLLM-provided primitives. + +Pre-built custom code for common guardrails (e.g. response rejection detection) +is available in response_rejection_code.py. """ from typing import TYPE_CHECKING @@ -9,6 +12,8 @@ from typing import TYPE_CHECKING from litellm.types.guardrails import SupportedGuardrailIntegrations from .custom_code_guardrail import CustomCodeGuardrail +from .response_rejection_code import (DEFAULT_REJECTION_PHRASES, + RESPONSE_REJECTION_GUARDRAIL_CODE) if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams @@ -61,5 +66,7 @@ guardrail_class_registry = { __all__ = [ "CustomCodeGuardrail", + "DEFAULT_REJECTION_PHRASES", + "RESPONSE_REJECTION_GUARDRAIL_CODE", "initialize_guardrail", ] diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index b2ef495be8..c557a093c4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -26,6 +26,12 @@ Example custom code (async with HTTP): if response["success"] and response["body"].get("flagged"): return block("Content flagged by moderation API") return allow() + +Example: block when response rejects the user (input_type response only): + + Use RESPONSE_REJECTION_GUARDRAIL_CODE from .response_rejection_code — it + checks response texts for phrases like "That's not something I can help with" + and returns block() so the guardrail raises a block error. """ import asyncio @@ -35,18 +41,18 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import ( - CustomGuardrail, - log_guardrail_information, -) +from litellm.integrations.custom_guardrail import (CustomGuardrail, + log_guardrail_information) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.base import \ + GuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs from .primitives import get_custom_code_primitives if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.litellm_logging import \ + Logging as LiteLLMLoggingObj class CustomCodeGuardrailError(Exception): diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/response_rejection_code.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/response_rejection_code.py new file mode 100644 index 0000000000..012895dbe1 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/response_rejection_code.py @@ -0,0 +1,76 @@ +""" +Custom code for a response guardrail that blocks when the model response +indicates it is rejecting the user request (e.g. "That's not something I can help with"). + +Use this with the Custom Code Guardrail (custom_code) by setting litellm_params.custom_code +to RESPONSE_REJECTION_GUARDRAIL_CODE. The guardrail runs only on input_type "response" +and raises a block error if any response text matches known rejection phrases. +""" + +# Default phrases that indicate the model is refusing the user request (lowercase for case-insensitive match). +# Custom code guardrails can override by defining rejection_phrases in the code. +DEFAULT_REJECTION_PHRASES = [ + "that's not something i can help with", + "that is not something i can help with", + "i can't help with that", + "i cannot help with that", + "i'm not able to help", + "i am not able to help", + "i'm unable to help", + "i cannot assist", + "i can't assist", + "i'm not allowed to", + "i'm not permitted to", + "i won't be able to help", + "i'm sorry, i can't", + "i'm sorry, i cannot", + "as an ai, i can't", + "as an ai, i cannot", +] + +# Custom code string for the Custom Code Guardrail. Only runs on input_type "response". +# Uses primitives: allow(), block(), lower(), contains() +RESPONSE_REJECTION_GUARDRAIL_CODE = ''' +def apply_guardrail(inputs, request_data, input_type): + """Block responses that indicate the model rejected the user request.""" + if input_type != "response": + return allow() + + texts = inputs.get("texts") or [] + # All lowercase for case-insensitive matching (text is lowercased before check) + rejection_phrases = [ + "that's not something i can help with", + "that is not something i can help with", + "i can't help with that", + "i cannot help with that", + "i'm not able to help", + "i am not able to help", + "i'm unable to help", + "i cannot assist", + "i can't assist", + "i'm not allowed to", + "i'm not permitted to", + "i won't be able to help", + "i'm sorry, i can't", + "i'm sorry, i cannot", + "as an ai, i can't", + "as an ai, i cannot", + ] + + for text in texts: + if not text: + continue + text_lower = lower(text) + for phrase in rejection_phrases: + if contains(text_lower, phrase): + return block( + "Response indicates the model rejected the user request.", + detection_info={"matched_phrase": phrase, "input_type": "response"}, + ) + return allow() +''' + +__all__ = [ + "DEFAULT_REJECTION_PHRASES", + "RESPONSE_REJECTION_GUARDRAIL_CODE", +] diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py index 069603afdf..4a42e493d3 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py @@ -12,10 +12,11 @@ All /policy management endpoints import copy import json import os -from typing import TYPE_CHECKING, AsyncIterator, List, Literal, Optional, cast +from typing import (TYPE_CHECKING, Any, AsyncIterator, List, Literal, Optional, + cast) from fastapi import APIRouter, Depends, HTTPException, Request -from fastapi.responses import StreamingResponse +from fastapi.responses import Response, StreamingResponse from pydantic import BaseModel, Field from typing_extensions import TypedDict @@ -24,8 +25,12 @@ from litellm.constants import (COMPETITOR_LLM_TEMPERATURE, DEFAULT_COMPETITOR_DISCOVERY_MODEL, MAX_COMPETITOR_NAMES) from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.openai.chat.guardrail_translation.handler import \ + OpenAIChatCompletionsHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.guardrails.guardrail_hooks.custom_code import ( + RESPONSE_REJECTION_GUARDRAIL_CODE, CustomCodeGuardrail) from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.policy_engine.policy_registry import get_policy_registry @@ -39,7 +44,7 @@ from litellm.types.proxy.policy_engine import (PolicyGuardrailsResponse, PolicyTestResponse, PolicyValidateRequest, PolicyValidationResponse) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import \ @@ -69,20 +74,32 @@ class GuardrailErrorEntry(TypedDict): message: str -class ApplyPoliciesResult(TypedDict): - """Result of apply_policies: inputs plus any guardrail failures.""" +class _ApplyPoliciesResultBase(TypedDict): + """Base result of apply_policies: inputs plus any guardrail failures.""" inputs: GenericGuardrailAPIInputs guardrail_errors: List[GuardrailErrorEntry] -class ApplyPoliciesPerItemResult(TypedDict): - """Result for one input when using inputs_list.""" +class ApplyPoliciesResult(_ApplyPoliciesResultBase, total=False): + """Result of apply_policies. agent_response set when agent_id provided.""" + + agent_response: Any + + +class _ApplyPoliciesPerItemResultBase(TypedDict): + """Base result for one input when using inputs_list.""" inputs: GenericGuardrailAPIInputs guardrail_errors: List[GuardrailErrorEntry] +class ApplyPoliciesPerItemResult(_ApplyPoliciesPerItemResultBase, total=False): + """Result for one input when using inputs_list. agent_response set when agent_id provided.""" + + agent_response: Any + + class ApplyPoliciesListResult(TypedDict): """Result when using inputs_list: one result per input.""" @@ -185,21 +202,78 @@ async def apply_policies( return {"inputs": current_inputs, "guardrail_errors": guardrail_errors} +def _chat_body_from_inputs( + inputs: GenericGuardrailAPIInputs, agent_id: str, request_data: dict +) -> dict: + """Build a chat completion request body from guardrail inputs and agent_id.""" + messages: List[dict] + structured = inputs.get("structured_messages") + texts = inputs.get("texts") + if structured: + messages = list(structured) # type: ignore[arg-type] + elif texts: + if len(texts) == 1: + messages = [{"role": "user", "content": texts[0]}] + else: + messages = [{"role": "user", "content": "\n".join(texts)}] + else: + messages = [{"role": "user", "content": "Hello"}] + body: dict = {"model": agent_id, "messages": messages, "stream": False} + if request_data: + body.setdefault("metadata", {}).update(request_data) + return body + + +def _request_with_json_body(body: dict) -> Request: + """Create a Starlette Request that will return the given dict as parsed JSON body.""" + body_bytes = json.dumps(body).encode() + received: List[bool] = [False] + + async def receive() -> dict: + if received[0]: + return {"type": "http.disconnect"} + received[0] = True + return {"type": "http.request", "body": body_bytes, "more_body": False} + + scope: dict = { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "query_string": b"", + "headers": [(b"content-type", b"application/json")], + "scheme": "http", + "server": ("localhost", 8000), + "client": ("127.0.0.1", 0), + "root_path": "", + "app": None, + "asgi": {"version": "3.0", "spec_version": "2.0"}, + } + return Request(scope, receive=receive) + + class TestPoliciesAndGuardrailsRequest(BaseModel): """Request body for POST /utils/test_policies_and_guardrails.""" - policy_names: Optional[List[str]] = Field(default=None, description="Policy names to resolve guardrails from") - guardrail_names: Optional[List[str]] = Field(default=None, description="Guardrail names to apply directly") - inputs: Optional[dict] = Field( - default=None, - description="GenericGuardrailAPIInputs, e.g. { \"texts\": [\"...\"] }. Use inputs_list for per-input processing.", + policy_names: Optional[List[str]] = Field( + default=None, description="Policy names to resolve guardrails from" ) - inputs_list: Optional[List[dict]] = Field( - default=None, + guardrail_names: Optional[List[str]] = Field( + default=None, description="Guardrail names to apply directly" + ) + inputs_list: List[GenericGuardrailAPIInputs] = Field( + default=[], description="List of GenericGuardrailAPIInputs; each item processed separately (for batch compliance testing).", ) - request_data: dict = Field(default_factory=dict, description="Request context (model, user_id, etc.)") - input_type: Literal["request", "response"] = Field(default="request", description="Whether inputs are request or response") + request_data: dict = Field( + default_factory=dict, description="Request context (model, user_id, etc.)" + ) + input_type: Literal["request", "response"] = Field( + default="request", description="Whether inputs are request or response" + ) + agent_id: Optional[str] = Field( + default=None, + description="When set, call chat completion with this model/agent for each input and include the response in the result.", + ) @router.post( @@ -223,40 +297,86 @@ async def test_policies_and_guardrails( """ from litellm.litellm_core_utils.litellm_logging import \ Logging as LiteLLMLoggingObj - from litellm.proxy.proxy_server import proxy_logging_obj + from litellm.proxy.proxy_server import chat_completion, proxy_logging_obj from litellm.proxy.utils import handle_exception_on_proxy + def _serialize_chat_response(response: Any) -> Any: + if hasattr(response, "model_dump"): + return response.model_dump(exclude_unset=True) + if isinstance(response, dict): + return response + return response + + async def _get_agent_response( + inputs: GenericGuardrailAPIInputs, + agent_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> Any: + body = _chat_body_from_inputs(inputs, agent_id, data.request_data) + req = _request_with_json_body(body) + resp = Response() + result = await chat_completion( + request=req, + fastapi_response=resp, + model=agent_id, + user_api_key_dict=user_api_key_dict, + ) + return _serialize_chat_response(result) + try: logging_obj = cast(LiteLLMLoggingObj, proxy_logging_obj) - if data.inputs_list is not None: - results: List[ApplyPoliciesPerItemResult] = [] - for inp in data.inputs_list: - inputs_typed = cast(GenericGuardrailAPIInputs, inp) - item_result = await apply_policies( - policy_names=data.policy_names, - inputs=inputs_typed, - request_data=data.request_data, - input_type=data.input_type, - proxy_logging_obj=logging_obj, - guardrail_names=data.guardrail_names, - ) - results.append( - ApplyPoliciesPerItemResult( - inputs=item_result["inputs"], - guardrail_errors=item_result["guardrail_errors"], - ) - ) - return ApplyPoliciesListResult(results=results) - if data.inputs is not None: - inputs_typed = cast(GenericGuardrailAPIInputs, data.inputs) - return await apply_policies( + + results: List[ApplyPoliciesPerItemResult] = [] + for inp in data.inputs_list: + item_result = await apply_policies( policy_names=data.policy_names, - inputs=inputs_typed, + inputs=inp, request_data=data.request_data, input_type=data.input_type, proxy_logging_obj=logging_obj, guardrail_names=data.guardrail_names, ) + item: ApplyPoliciesPerItemResult = { + "inputs": item_result["inputs"], + "guardrail_errors": item_result["guardrail_errors"], + } + if data.agent_id is not None: + item["agent_response"] = await _get_agent_response( + item_result["inputs"], + data.agent_id, + user_api_key_dict, + ) + # run response through response_rejection_guardrail (reuses handler extraction + apply) + response_rejection_guardrail = CustomCodeGuardrail( + custom_code=RESPONSE_REJECTION_GUARDRAIL_CODE, + guardrail_name="response_rejection", + ) + try: + model_response = ModelResponse.model_validate( + item["agent_response"] + ) + handler = OpenAIChatCompletionsHandler() + await handler.process_output_response( + response=model_response, + guardrail_to_apply=response_rejection_guardrail, + litellm_logging_obj=logging_obj, + user_api_key_dict=user_api_key_dict, + ) + except Exception as guardrail_err: + item["guardrail_errors"] = list(item["guardrail_errors"]) + detail = getattr(guardrail_err, "detail", None) + if isinstance(detail, dict) and "error" in detail: + err_msg = detail["error"] + else: + err_msg = str(detail if detail is not None else guardrail_err) + item["guardrail_errors"].append( + GuardrailErrorEntry( + guardrail_name="response_rejection", + message=err_msg, + ) + ) + results.append(item) + return ApplyPoliciesListResult(results=results) raise ValueError("Either inputs or inputs_list must be provided") except Exception as e: raise handle_exception_on_proxy(e) @@ -581,11 +701,15 @@ def _validate_enrichment_request(data: EnrichTemplateRequest) -> tuple[dict, dic templates = _load_policy_templates_from_local_backup() template = next((t for t in templates if t.get("id") == data.template_id), None) if template is None: - raise HTTPException(status_code=404, detail=f"Template '{data.template_id}' not found") + raise HTTPException( + status_code=404, detail=f"Template '{data.template_id}' not found" + ) llm_enrichment = template.get("llm_enrichment") if llm_enrichment is None: - raise HTTPException(status_code=400, detail="Template does not support LLM enrichment") + raise HTTPException( + status_code=400, detail="Template does not support LLM enrichment" + ) # Validate competitors list size if provided if data.competitors and len(data.competitors) > MAX_COMPETITOR_NAMES: @@ -695,7 +819,11 @@ async def _stream_llm_competitor_names( while "\n" in buffer: line, buffer = buffer.split("\n", 1) name = _clean_competitor_line(line) - if name and name.lower() not in existing_lower and count < MAX_COMPETITOR_NAMES: + if ( + name + and name.lower() not in existing_lower + and count < MAX_COMPETITOR_NAMES + ): existing_lower.add(name.lower()) count += 1 yield name, False @@ -744,9 +872,7 @@ async def _stream_competitor_events( "{{" + llm_enrichment["parameter"] + "}}", brand_name ) try: - async for name, _ in _stream_llm_competitor_names( - prompt, model, [] - ): + async for name, _ in _stream_llm_competitor_names(prompt, model, []): if name: competitors.append(name) yield f"data: {json.dumps({'type': 'competitor', 'name': name})}\n\n" @@ -893,10 +1019,7 @@ def _build_all_names_per_competitor( competitors: list[str], variations_map: dict[str, list[str]] ) -> dict[str, list[str]]: """Build canonical + variation name lists for each competitor.""" - return { - comp: [comp] + variations_map.get(comp, []) - for comp in competitors - } + return {comp: [comp] + variations_map.get(comp, []) for comp in competitors} def _build_competitor_guardrail_definitions( @@ -912,7 +1035,9 @@ def _build_competitor_guardrail_definitions( output_blocked = _build_name_blocked_words(competitors, all_names) recommendation_blocked = _build_recommendation_blocked_words(competitors, all_names) - comparison_blocked = _build_comparison_blocked_words(competitors, all_names, brand_name) + comparison_blocked = _build_comparison_blocked_words( + competitors, all_names, brand_name + ) blocked_words_map = { "competitor-output-blocker": output_blocked, @@ -943,7 +1068,11 @@ def _build_name_blocked_words( result = [] for comp in competitors: for name in all_names[comp]: - desc = f"Competitor: {comp}" if name == comp else f"Competitor variation ({comp}): {name}" + desc = ( + f"Competitor: {comp}" + if name == comp + else f"Competitor variation ({comp}): {name}" + ) result.append({"keyword": name, "action": "BLOCK", "description": desc}) return result @@ -956,11 +1085,13 @@ def _build_recommendation_blocked_words( for comp in competitors: for name in all_names[comp]: for prefix in ["try", "use", "switch to", "consider"]: - result.append({ - "keyword": f"{prefix} {name}", - "action": "BLOCK", - "description": f"Recommendation to competitor ({comp})", - }) + result.append( + { + "keyword": f"{prefix} {name}", + "action": "BLOCK", + "description": f"Recommendation to competitor ({comp})", + } + ) return result @@ -971,23 +1102,29 @@ def _build_comparison_blocked_words( result = [] for comp in competitors: for name in all_names[comp]: - result.append({ - "keyword": f"{name} is better", - "action": "BLOCK", - "description": f"Unfavorable comparison ({comp})", - }) + result.append( + { + "keyword": f"{name} is better", + "action": "BLOCK", + "description": f"Unfavorable comparison ({comp})", + } + ) # Brand-level comparisons (only need one entry each, not per-competitor) - result.append({ - "keyword": f"better than {brand_name}", - "action": "BLOCK", - "description": "Unfavorable comparison", - }) - result.append({ - "keyword": f"{brand_name} is worse", - "action": "BLOCK", - "description": "Unfavorable comparison", - }) + result.append( + { + "keyword": f"better than {brand_name}", + "action": "BLOCK", + "description": "Unfavorable comparison", + } + ) + result.append( + { + "keyword": f"{brand_name} is worse", + "action": "BLOCK", + "description": "Unfavorable comparison", + } + ) return result @@ -1121,7 +1258,9 @@ async def _test_guardrail_definitions( request_data={}, input_type="request", ) - output_text = output.get("texts", [text])[0] if output.get("texts") else text + output_text = ( + output.get("texts", [text])[0] if output.get("texts") else text + ) if output_text != text: action = "masked" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ac795ba57c..7760a894c7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1,61 +1,47 @@ import json import time from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union +from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, + Union) from openai._models import BaseModel as OpenAIObject -from openai.types.audio.transcription_create_params import ( - FileTypes as FileTypes, # type: ignore -) +from openai.types.audio.transcription_create_params import \ + FileTypes as FileTypes # type: ignore from openai.types.chat.chat_completion import ChatCompletion as ChatCompletion -from openai.types.completion_usage import ( - CompletionTokensDetails, - CompletionUsage, - PromptTokensDetails, -) +from openai.types.completion_usage import (CompletionTokensDetails, + CompletionUsage, + PromptTokensDetails) from openai.types.moderation import Categories as Categories -from openai.types.moderation import ( - CategoryAppliedInputTypes as CategoryAppliedInputTypes, -) +from openai.types.moderation import \ + CategoryAppliedInputTypes as CategoryAppliedInputTypes from openai.types.moderation import CategoryScores as CategoryScores from openai.types.moderation_create_response import Moderation as Moderation -from openai.types.moderation_create_response import ( - ModerationCreateResponse as ModerationCreateResponse, -) +from openai.types.moderation_create_response import \ + ModerationCreateResponse as ModerationCreateResponse from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator from typing_extensions import Required, TypedDict from litellm._uuid import uuid -from litellm.types.llms.base import ( - BaseLiteLLMOpenAIResponseObject, - LiteLLMPydanticObjectBase, -) +from litellm.types.llms.base import (BaseLiteLLMOpenAIResponseObject, + LiteLLMPydanticObjectBase) from litellm.types.mcp import MCPServerCostInfo from ..litellm_core_utils.core_helpers import map_finish_reason from .agents import LiteLLMSendMessageResponse from .guardrails import GuardrailEventHooks -from .llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from .llms.anthropic_messages.anthropic_response import \ + AnthropicMessagesResponse from .llms.base import HiddenParams -from .llms.openai import ( - AllMessageValues, - Batch, - ChatCompletionAnnotation, - ChatCompletionRedactedThinkingBlock, - ChatCompletionThinkingBlock, - ChatCompletionToolCallChunk, - ChatCompletionToolParam, - ChatCompletionUsageBlock, - FileSearchTool, - FineTuningJob, - ImageURLListItem, - OpenAIChatCompletionChunk, - OpenAIChatCompletionFinishReason, - OpenAIFileObject, - OpenAIRealtimeStreamList, - ResponsesAPIResponse, - WebSearchOptions, -) +from .llms.openai import (AllMessageValues, Batch, ChatCompletionAnnotation, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, + ChatCompletionToolCallChunk, ChatCompletionToolParam, + ChatCompletionUsageBlock, FileSearchTool, + FineTuningJob, ImageURLListItem, + OpenAIChatCompletionChunk, + OpenAIChatCompletionFinishReason, OpenAIFileObject, + OpenAIRealtimeStreamList, ResponsesAPIResponse, + WebSearchOptions) from .rerank import RerankResponse as RerankResponse if TYPE_CHECKING: @@ -3173,6 +3159,7 @@ class LlmProviders(str, Enum): POE = "poe" CHUTES = "chutes" XIAOMI_MIMO = "xiaomi_mimo" + LITELLM_AGENT = "litellm_agent" # Create a set of all provider values for quick lookup @@ -3203,6 +3190,7 @@ class SearchProviders(str, Enum): LINKUP = "linkup" DUCKDUCKGO = "duckduckgo" + # Create a set of all search provider values for quick lookup SearchProvidersSet = {provider.value for provider in SearchProviders} diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py new file mode 100644 index 0000000000..149a0b5eae --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_response_rejection_guardrail_code.py @@ -0,0 +1,78 @@ +"""Tests for the response-rejection custom guardrail code (input_type response, block on refusal).""" + +import pytest +from fastapi import HTTPException + +from litellm.proxy.guardrails.guardrail_hooks.custom_code import ( + RESPONSE_REJECTION_GUARDRAIL_CODE, CustomCodeGuardrail) + + +@pytest.fixture +def response_rejection_guardrail(): + """Guardrail instance using the response-rejection custom code.""" + return CustomCodeGuardrail( + guardrail_name="response_rejection", + custom_code=RESPONSE_REJECTION_GUARDRAIL_CODE, + ) + + +@pytest.mark.asyncio +async def test_response_rejection_allows_request_input_type(response_rejection_guardrail): + """Should allow when input_type is 'request' (no response check).""" + result = await response_rejection_guardrail.apply_guardrail( + inputs={"texts": ["some user message"]}, + request_data={}, + input_type="request", + ) + assert result == {"texts": ["some user message"]} + + +@pytest.mark.asyncio +async def test_response_rejection_allows_helpful_response(response_rejection_guardrail): + """Should allow when response text does not contain rejection phrases.""" + result = await response_rejection_guardrail.apply_guardrail( + inputs={"texts": ["Here is how you can do that: step 1, step 2."]}, + request_data={}, + input_type="response", + ) + assert result["texts"] == ["Here is how you can do that: step 1, step 2."] + + +@pytest.mark.asyncio +async def test_response_rejection_blocks_refusal_phrase(response_rejection_guardrail): + """Should block when response contains a known rejection phrase.""" + with pytest.raises(HTTPException) as exc_info: + await response_rejection_guardrail.apply_guardrail( + inputs={"texts": ["That's not something I can help with."]}, + request_data={}, + input_type="response", + ) + assert exc_info.value.status_code == 400 + detail = exc_info.value.detail + assert isinstance(detail, dict) + assert "error" in detail + assert "rejected" in detail["error"].lower() or "reject" in detail["error"].lower() + assert detail.get("guardrail") == "response_rejection" + assert detail.get("detection_info", {}).get("matched_phrase") is not None + + +@pytest.mark.asyncio +async def test_response_rejection_blocks_case_insensitive(response_rejection_guardrail): + """Should block on refusal phrase regardless of case.""" + with pytest.raises(HTTPException): + await response_rejection_guardrail.apply_guardrail( + inputs={"texts": ["I'M SORRY, I CAN'T do that."]}, + request_data={}, + input_type="response", + ) + + +@pytest.mark.asyncio +async def test_response_rejection_empty_texts_allowed(response_rejection_guardrail): + """Should allow when texts is empty or missing.""" + result = await response_rejection_guardrail.apply_guardrail( + inputs={}, + request_data={}, + input_type="response", + ) + assert result == {} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 3ed71ecab1..a16afe9c8a 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -5453,6 +5453,8 @@ export interface TestPoliciesAndGuardrailsRequest { inputs_list?: GuardrailInputs[] | null; request_data?: Record; input_type?: "request" | "response"; + /** When set, backend runs chat completion with this model/agent per input and includes agent_response in each result. */ + agent_id?: string | null; } export interface GuardrailErrorEntry { @@ -5460,16 +5462,24 @@ export interface GuardrailErrorEntry { message: string; } +export interface TestPoliciesAndGuardrailsResultItem { + inputs: Record; + guardrail_errors: GuardrailErrorEntry[]; + /** Present when request included agent_id; serialized chat completion response. */ + agent_response?: Record; +} + export interface TestPoliciesAndGuardrailsResponse { inputs?: Record; guardrail_errors?: GuardrailErrorEntry[]; /** Present when inputs_list was used; one result per input. */ - results?: Array<{ inputs: Record; guardrail_errors: GuardrailErrorEntry[] }>; + results?: TestPoliciesAndGuardrailsResultItem[]; } export const testPoliciesAndGuardrails = async ( accessToken: string, - body: TestPoliciesAndGuardrailsRequest + body: TestPoliciesAndGuardrailsRequest, + signal?: AbortSignal ): Promise => { try { const url = proxyBaseUrl @@ -5477,6 +5487,7 @@ export const testPoliciesAndGuardrails = async ( : `/utils/test_policies_and_guardrails`; const response = await fetch(url, { method: "POST", + signal, headers: { [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", @@ -5488,6 +5499,7 @@ export const testPoliciesAndGuardrails = async ( inputs_list: body.inputs_list ?? null, request_data: body.request_data ?? {}, input_type: body.input_type ?? "request", + agent_id: body.agent_id ?? null, }), }); diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx index 10719b7ff4..c47c201d07 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/AgentBuilderView.tsx @@ -1,11 +1,14 @@ "use client"; -import { CommentOutlined, ExperimentOutlined, PlusOutlined, RobotOutlined, SaveOutlined } from "@ant-design/icons"; -import { Button, Input, Select, Spin, Tabs } from "antd"; +import { CommentOutlined, DeleteOutlined, ExperimentOutlined, LinkOutlined, PlusOutlined, RobotOutlined, SaveOutlined } from "@ant-design/icons"; +import { Button, Input, Modal, Select, Spin, Tabs } from "antd"; import React, { useCallback, useEffect, useState } from "react"; +import CodeBlock from "@/app/(dashboard)/api-reference/components/CodeBlock"; import NotificationsManager from "../../molecules/notifications_manager"; -import { modelCreateCall } from "../../networking"; -import { AgentModel, fetchAvailableAgentModels } from "../llm_calls/fetch_agents"; +import { keyCreateCall, modelCreateCall, modelDeleteCall, modelPatchUpdateCall, proxyBaseUrl } from "../../networking"; +import { fetchMCPServers } from "../../networking"; +import { MCPServer } from "../../mcp_tools/types"; +import { AgentModel, fetchAvailableAgentModels, MCPToolEntry } from "../llm_calls/fetch_agents"; import { fetchAvailableModels, ModelGroup } from "../llm_calls/fetch_models"; import ComplianceUI from "../complianceUI/ComplianceUI"; import ChatUI from "./ChatUI"; @@ -28,6 +31,133 @@ export interface AgentBuilderViewProps { const NEW_AGENT_ID = "__new__"; +function getConnectTabBaseUrl( + proxySettings: AgentBuilderViewProps["proxySettings"], + customProxyBaseUrl?: string, +): string { + const customDocBaseUrl = proxySettings?.LITELLM_UI_API_DOC_BASE_URL; + if (customDocBaseUrl && customDocBaseUrl.trim()) return customDocBaseUrl; + if (proxySettings?.PROXY_BASE_URL) return proxySettings.PROXY_BASE_URL; + if (customProxyBaseUrl?.trim()) return customProxyBaseUrl; + return ""; +} + +interface ConnectTabContentProps { + agentName: string; + proxySettings: AgentBuilderViewProps["proxySettings"]; + customProxyBaseUrl?: string; + accessToken: string | null; + userID: string | null; + disabledPersonalKeyCreation: boolean; + creatingKey: boolean; + createdKeyValue: string | null; + onCreateKey: () => void; +} + +function ConnectTabContent({ + agentName, + proxySettings, + customProxyBaseUrl, + disabledPersonalKeyCreation, + creatingKey, + createdKeyValue, + onCreateKey, +}: ConnectTabContentProps) { + const baseUrl = proxyBaseUrl ?? getConnectTabBaseUrl(proxySettings, customProxyBaseUrl); + const apiKeyForCurl = + createdKeyValue ? + createdKeyValue.startsWith("Bearer ") ? createdKeyValue : `Bearer ${createdKeyValue}` + : "Bearer sk-1234"; + const curlExample = `curl -L -X POST '${baseUrl}/v1/chat/completions' \\ +-H 'x-litellm-api-key: ${apiKeyForCurl}' \\ +-d '{ + "model": "${agentName}", + "stream": true, + "stream_options": { + "include_usage": true + }, + "messages": [ + { + "role": "user", + "content": "hey" + } + ] +}'`; + return ( +
+
+

Proxy base URL

+

+ {baseUrl} +

+
+
+

Call your agent (cURL)

+ +
+
+

Create a key for this agent

+

+ Create a virtual key that can only call this agent. The key will be scoped to you (user_id) and restricted to + the model {agentName}. +

+ + {disabledPersonalKeyCreation && ( +

Key creation is disabled for your account.

+ )} + {createdKeyValue && ( +

+ Key created. It is shown in the cURL example above — copy the snippet to use it. +

+ )} +
+
+ ); +} + +function getAgentModelId(agent: AgentModel): string | null { + const info = agent.model_info as { id?: string } | null | undefined; + return info?.id ?? null; +} + +function parseUnderlyingModel(litellmModel: string | undefined): string | undefined { + if (!litellmModel || !litellmModel.startsWith("litellm_agent/")) return undefined; + return litellmModel.slice("litellm_agent/".length) || undefined; +} + +const MCP_TOOLS_PREFIX = "litellm_proxy/mcp/"; + +function buildToolsFromServerIds(serverIds: string[], servers: MCPServer[]): MCPToolEntry[] { + return serverIds.map((serverId) => { + const server = servers.find((s) => s.server_id === serverId); + const serverName = server?.alias || server?.server_name || serverId; + return { + type: "mcp", + server_label: "litellm", + server_url: `${MCP_TOOLS_PREFIX}${serverName}`, + require_approval: "never", + }; + }); +} + +function getServerIdsFromTools(tools: MCPToolEntry[], servers: MCPServer[]): string[] { + return tools + .filter((t) => t.type === "mcp" && t.server_url?.startsWith(MCP_TOOLS_PREFIX)) + .map((t) => { + const suffix = t.server_url.slice(MCP_TOOLS_PREFIX.length); + const server = servers.find((s) => (s.alias || s.server_name || s.server_id) === suffix); + return server?.server_id; + }) + .filter((id): id is string => id != null); +} + export default function AgentBuilderView({ accessToken, token, @@ -42,7 +172,9 @@ export default function AgentBuilderView({ const [modelGroups, setModelGroups] = useState([]); const [loadingAgents, setLoadingAgents] = useState(true); const [selectedId, setSelectedId] = useState(null); - const [activeTab, setActiveTab] = useState<"configure" | "chat" | "test">("configure"); + const [activeTab, setActiveTab] = useState<"configure" | "chat" | "test" | "connect">("configure"); + const [creatingKey, setCreatingKey] = useState(false); + const [createdKeyValue, setCreatedKeyValue] = useState(null); // Draft for new agent const [draftName, setDraftName] = useState(""); @@ -50,12 +182,18 @@ export default function AgentBuilderView({ const [draftUnderlyingModel, setDraftUnderlyingModel] = useState(undefined); const [draftTemperature, setDraftTemperature] = useState(0.7); const [draftMaxTokens, setDraftMaxTokens] = useState(4096); + const [draftTools, setDraftTools] = useState([]); + + const [mcpServers, setMCPServers] = useState([]); + const [loadingMCPServers, setLoadingMCPServers] = useState(false); const [saving, setSaving] = useState(false); + const [deleting, setDeleting] = useState(false); const effectiveApiKey = apiKey || accessToken || ""; const selectedAgent = selectedId === NEW_AGENT_ID ? null : agentModels.find((a) => a.model_name === selectedId) ?? null; const isNewAgent = selectedId === NEW_AGENT_ID; + const selectedAgentModelId = selectedAgent ? getAgentModelId(selectedAgent) : null; const loadAgents = useCallback(async () => { if (!accessToken || !userID || !userRole) return; @@ -95,6 +233,52 @@ export default function AgentBuilderView({ loadModels(); }, [loadModels]); + const loadMCPServers = useCallback(async () => { + if (!effectiveApiKey) return; + setLoadingMCPServers(true); + try { + const servers = await fetchMCPServers(effectiveApiKey); + setMCPServers(Array.isArray(servers) ? servers : (servers as { data?: MCPServer[] })?.data ?? []); + } catch (e) { + console.error("Error fetching MCP servers:", e); + } finally { + setLoadingMCPServers(false); + } + }, [effectiveApiKey]); + + useEffect(() => { + loadMCPServers(); + }, [loadMCPServers]); + + // Clear created key when switching to another agent + useEffect(() => { + setCreatedKeyValue(null); + }, [selectedId]); + + // Sync draft fields when selecting an existing agent + useEffect(() => { + if (selectedAgent && !isNewAgent) { + setDraftName(selectedAgent.model_name); + setDraftSystemPrompt(selectedAgent.litellm_params?.litellm_system_prompt ?? ""); + const underlying = parseUnderlyingModel(selectedAgent.litellm_params?.model); + setDraftUnderlyingModel(underlying ?? modelGroups[0]?.model_group); + const p = selectedAgent.litellm_params as { temperature?: number; max_tokens?: number } | undefined; + setDraftTemperature(typeof p?.temperature === "number" ? p.temperature : 0.7); + setDraftMaxTokens(typeof p?.max_tokens === "number" ? p.max_tokens : 4096); + const rawTools = selectedAgent.litellm_params?.tools; + const tools: MCPToolEntry[] = Array.isArray(rawTools) + ? rawTools.filter((t): t is MCPToolEntry => t && typeof t === "object" && (t as MCPToolEntry).type === "mcp" && typeof (t as MCPToolEntry).server_url === "string") + : []; + setDraftTools(tools); + } + }, [selectedId, isNewAgent, selectedAgent?.model_name, selectedAgent?.litellm_params?.tools]); + + const selectedMCPServerIds = getServerIdsFromTools(draftTools, mcpServers); + + const handleMCPServerChange = (serverIds: string[]) => { + setDraftTools(buildToolsFromServerIds(serverIds, mcpServers)); + }; + const handleAddAgent = () => { setSelectedId(NEW_AGENT_ID); setDraftName(""); @@ -102,6 +286,7 @@ export default function AgentBuilderView({ setDraftUnderlyingModel(modelGroups[0]?.model_group); setDraftTemperature(0.7); setDraftMaxTokens(4096); + setDraftTools([]); setActiveTab("configure"); }; @@ -119,6 +304,7 @@ export default function AgentBuilderView({ litellm_system_prompt: draftSystemPrompt.trim() || undefined, temperature: draftTemperature, max_tokens: draftMaxTokens, + tools: draftTools, }, model_info: {}, }); @@ -133,6 +319,86 @@ export default function AgentBuilderView({ } }; + const handleUpdateAgent = async () => { + if (!accessToken || !selectedAgent || !selectedAgentModelId || !draftName?.trim() || !draftUnderlyingModel) { + NotificationsManager.fromBackend("Name and underlying model are required"); + return; + } + setSaving(true); + try { + await modelPatchUpdateCall( + accessToken, + { + model_name: draftName.trim(), + litellm_params: { + model: `litellm_agent/${draftUnderlyingModel}`, + litellm_system_prompt: draftSystemPrompt.trim() || undefined, + temperature: draftTemperature, + max_tokens: draftMaxTokens, + tools: draftTools, + }, + model_info: selectedAgent.model_info ?? {}, + }, + selectedAgentModelId, + ); + NotificationsManager.success("Agent updated successfully"); + await loadAgents(); + setSelectedId(draftName.trim()); + } catch (e) { + NotificationsManager.fromBackend("Failed to update agent"); + } finally { + setSaving(false); + } + }; + + const handleCreateKeyForAgent = async () => { + if (!accessToken || !userID || !selectedAgent) return; + setCreatingKey(true); + setCreatedKeyValue(null); + try { + const response = await keyCreateCall(accessToken, userID, { + models: [selectedAgent.model_name], + key_alias: `Agent: ${selectedAgent.model_name}`, + }); + const keyValue = response?.key ?? null; + if (keyValue) { + setCreatedKeyValue(keyValue); + NotificationsManager.success("Virtual key created. Use it in the curl example below."); + } else { + NotificationsManager.fromBackend("Key created but value not returned"); + } + } catch (e) { + NotificationsManager.fromBackend("Failed to create key for agent"); + } finally { + setCreatingKey(false); + } + }; + + const handleDeleteAgent = () => { + if (!selectedAgent || !selectedAgentModelId || !accessToken) return; + Modal.confirm({ + title: "Delete agent", + content: `Are you sure you want to delete "${selectedAgent.model_name}"? This cannot be undone.`, + okText: "Delete", + okType: "danger", + cancelText: "Cancel", + onOk: async () => { + setDeleting(true); + try { + await modelDeleteCall(accessToken, selectedAgentModelId); + NotificationsManager.success("Agent deleted"); + await loadAgents(); + const remaining = agentModels.filter((a) => a.model_name !== selectedAgent.model_name); + setSelectedId(remaining.length > 0 ? remaining[0].model_name : null); + } catch (e) { + NotificationsManager.fromBackend("Failed to delete agent"); + } finally { + setDeleting(false); + } + }, + }); + }; + if (!accessToken || !userID || !userRole) { return (
@@ -163,7 +429,11 @@ export default function AgentBuilderView({
- Agent Builder is experimental and may change or be removed without notice. + Agent Builder is experimental and may change or be removed without notice. We’d love your feedback—email us at{" "} + + product@berri.ai + + .
@@ -220,7 +490,7 @@ export default function AgentBuilderView({ <> setActiveTab(k as "configure" | "chat" | "test")} + onChange={(k) => setActiveTab(k as "configure" | "chat" | "test" | "connect")} className="flex-1 overflow-hidden [&_.ant-tabs-content]:h-full [&_.ant-tabs-tabpane]:h-full [&_.ant-tabs-nav]:pl-4" items={[ { @@ -232,8 +502,13 @@ export default function AgentBuilderView({ ), children: (
- {isNewAgent ? ( + {(isNewAgent || selectedAgent) ? (
+ {!selectedAgentModelId && selectedAgent && ( +
+ This agent cannot be updated or deleted here (missing model id). Manage it from Models & Endpoints. +
+ )}
-
- ) : selectedAgent ? ( -
- -
- {selectedAgent.model_name} -
+ +