diff --git a/docs/my-website/docs/proxy/guardrails/qualifire.md b/docs/my-website/docs/proxy/guardrails/qualifire.md index 66961c92d9d..850af37e47f 100644 --- a/docs/my-website/docs/proxy/guardrails/qualifire.md +++ b/docs/my-website/docs/proxy/guardrails/qualifire.md @@ -8,13 +8,7 @@ Use [Qualifire](https://qualifire.ai) to evaluate LLM outputs for quality, safet ## Quick Start -### 1. Install the Qualifire SDK - -```bash -pip install qualifire -``` - -### 2. Define Guardrails on your LiteLLM config.yaml +### 1. Define Guardrails on your LiteLLM config.yaml Define your guardrails under the `guardrails` section: @@ -61,13 +55,13 @@ guardrails: - `post_call` Run **after** LLM call, on **input & output** - `during_call` Run **during** LLM call, on **input**. Same as `pre_call` but runs in parallel as LLM call. Response not returned until guardrail check completes -### 3. Start LiteLLM Gateway +### 2. Start LiteLLM Gateway ```shell litellm --config config.yaml --detailed_debug ``` -### 4. Test request +### 3. Test request **[Langchain, OpenAI SDK Usage Examples](../proxy/user_keys#request-format)** @@ -142,7 +136,7 @@ guardrails: evaluation_id: eval_abc123 # Your evaluation ID from Qualifire dashboard ``` -When `evaluation_id` is provided, LiteLLM will use `invoke_evaluation()` instead of `evaluate()`, running the pre-configured evaluation from your dashboard. +When `evaluation_id` is provided, LiteLLM will use the invoke evaluation API endpoint instead of the evaluate endpoint, running the pre-configured evaluation from your dashboard. ## Available Checks @@ -213,19 +207,19 @@ guardrails: ### Parameter Reference -| Parameter | Type | Default | Description | -| ------------------------------ | ----------- | --------------------------- | -------------------------------------------------------- | -| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key | -| `api_base` | `str` | `None` | Custom API base URL (optional) | -| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard | -| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection | -| `hallucinations_check` | `bool` | `None` | Enable hallucination detection | -| `grounding_check` | `bool` | `None` | Enable grounding verification | -| `pii_check` | `bool` | `None` | Enable PII detection | -| `content_moderation_check` | `bool` | `None` | Enable content moderation | -| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check | -| `assertions` | `List[str]` | `None` | Custom assertions to validate | -| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` | +| Parameter | Type | Default | Description | +| ------------------------------ | ----------- | ---------------------------- | -------------------------------------------------------- | +| `api_key` | `str` | `QUALIFIRE_API_KEY` env var | Your Qualifire API key | +| `api_base` | `str` | `https://proxy.qualifire.ai` | Custom API base URL (optional) | +| `evaluation_id` | `str` | `None` | Pre-configured evaluation ID from Qualifire dashboard | +| `prompt_injections` | `bool` | `true` (if no other checks) | Enable prompt injection detection | +| `hallucinations_check` | `bool` | `None` | Enable hallucination detection | +| `grounding_check` | `bool` | `None` | Enable grounding verification | +| `pii_check` | `bool` | `None` | Enable PII detection | +| `content_moderation_check` | `bool` | `None` | Enable content moderation | +| `tool_selection_quality_check` | `bool` | `None` | Enable tool selection quality check | +| `assertions` | `List[str]` | `None` | Custom assertions to validate | +| `on_flagged` | `str` | `"block"` | Action when content is flagged: `"block"` or `"monitor"` | ### Default Behavior @@ -261,4 +255,3 @@ This evaluates whether the LLM selected the appropriate tools and provided corre - [Qualifire Documentation](https://docs.qualifire.ai) - [Qualifire Dashboard](https://app.qualifire.ai) -- [Qualifire Python SDK](https://github.com/qualifire-dev/qualifire-python-sdk) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 5d2f096156e..004132ca05b 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -55,6 +55,7 @@ const sidebars = { "proxy/guardrails/test_playground", "proxy/guardrails/litellm_content_filter", ...[ + "proxy/guardrails/qualifire", "proxy/guardrails/aim_security", "proxy/guardrails/onyx_security", "proxy/guardrails/aporia_api", diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index a6971b49f3b..87da11efad0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -5,6 +5,7 @@ # +-------------------------------------------------------------+ # Qualifire - Evaluate LLM outputs for quality, safety, and reliability +import json import os from typing import Any, Dict, List, Literal, Optional, Type @@ -15,12 +16,17 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import ( Logging as LiteLLMLoggingObj, ) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs GUARDRAIL_NAME = "qualifire" +DEFAULT_QUALIFIRE_API_BASE = "https://proxy.qualifire.ai" class QualifireGuardrail(CustomGuardrail): @@ -44,7 +50,7 @@ class QualifireGuardrail(CustomGuardrail): Args: api_key: API key for Qualifire (or use QUALIFIRE_API_KEY env var) - api_base: Optional custom API base URL + api_base: Optional custom API base URL (defaults to https://api.qualifire.ai) evaluation_id: Pre-configured evaluation ID from Qualifire dashboard prompt_injections: Enable prompt injection detection (default if no other checks) hallucinations_check: Enable hallucination detection @@ -64,6 +70,7 @@ class QualifireGuardrail(CustomGuardrail): api_base or get_secret_str("QUALIFIRE_BASE_URL") or os.environ.get("QUALIFIRE_BASE_URL") + or DEFAULT_QUALIFIRE_API_BASE ) self.evaluation_id = evaluation_id self.prompt_injections = prompt_injections @@ -79,7 +86,11 @@ class QualifireGuardrail(CustomGuardrail): if not self._has_any_check_enabled() and not self.evaluation_id: self.prompt_injections = True - self._client = None + # Initialize async HTTP client for direct API calls + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + super().__init__(**kwargs) def _has_any_check_enabled(self) -> bool: @@ -96,43 +107,22 @@ class QualifireGuardrail(CustomGuardrail): ] ) - def _get_client(self): - """Lazy initialization of Qualifire client.""" - if self._client is None: - try: - from qualifire.client import Client - except ImportError: - raise ImportError( - "qualifire package is required for QualifireGuardrail. " - "Install it with: pip install qualifire" - ) - - client_kwargs: Dict[str, Any] = {} - if self.qualifire_api_key: - client_kwargs["api_key"] = self.qualifire_api_key - if self.qualifire_api_base: - client_kwargs["base_url"] = self.qualifire_api_base - - self._client = Client(**client_kwargs) - - return self._client - - def _convert_messages_to_qualifire_format( + def _convert_messages_to_api_format( self, messages: List[AllMessageValues] - ) -> List[Any]: + ) -> List[Dict[str, Any]]: """ - Convert LiteLLM messages to Qualifire's LLMMessage format. + Convert LiteLLM messages to Qualifire API format. Supports tool calls for tool_selection_quality_check. - """ - try: - from qualifire.types import LLMMessage, LLMToolCall - except ImportError: - raise ImportError( - "qualifire package is required for QualifireGuardrail. " - "Install it with: pip install qualifire" - ) - qualifire_messages = [] + Returns a list of dicts matching the API's ModelInvocationCanonicalMessage schema: + { + "role": "user" | "assistant" | "system" | "tool", + "content": "...", + "tool_call_id": "...", # optional + "tool_calls": [{"id": "...", "name": "...", "arguments": {...}}] # optional + } + """ + api_messages = [] for msg in messages: role = msg.get("role", "user") content = msg.get("content", "") @@ -147,42 +137,86 @@ class QualifireGuardrail(CustomGuardrail): text_parts.append(part) content = "\n".join(text_parts) - llm_message_kwargs: Dict[str, Any] = { + api_message: Dict[str, Any] = { "role": role, "content": content if isinstance(content, str) else str(content), } + # Handle tool_call_id for tool response messages + tool_call_id = msg.get("tool_call_id") + if tool_call_id: + api_message["tool_call_id"] = tool_call_id + # Handle tool calls if present tool_calls = msg.get("tool_calls") if tool_calls and isinstance(tool_calls, list): - qualifire_tool_calls = [] + api_tool_calls = [] for tc in tool_calls: if isinstance(tc, dict): function_info = tc.get("function", {}) # Arguments can be a string (JSON) or dict args = function_info.get("arguments", {}) if isinstance(args, str): - import json - try: args = json.loads(args) except json.JSONDecodeError: args = {} - qualifire_tool_calls.append( - LLMToolCall( - id=tc.get("id") or "", - name=function_info.get("name") or "", - arguments=args if isinstance(args, dict) else {}, - ) + api_tool_calls.append( + { + "id": tc.get("id") or "", + "name": function_info.get("name") or "", + "arguments": args if isinstance(args, dict) else {}, + } ) - if qualifire_tool_calls: - llm_message_kwargs["tool_calls"] = qualifire_tool_calls + if api_tool_calls: + api_message["tool_calls"] = api_tool_calls - qualifire_messages.append(LLMMessage(**llm_message_kwargs)) + api_messages.append(api_message) - return qualifire_messages + return api_messages - def _check_if_flagged(self, result: Any) -> bool: + def _convert_tools_to_api_format( + self, tools: Optional[List[Any]] + ) -> Optional[List[Dict[str, Any]]]: + """ + Convert OpenAI-format tools to Qualifire API format. + + Returns a list of dicts matching the API's ModelInvocationToolDefinition schema: + { + "name": "...", + "description": "...", + "parameters": {...} + } + """ + if not tools: + return None + + api_tools = [] + for tool in tools: + if isinstance(tool, dict): + # Handle OpenAI function tool format + if tool.get("type") == "function": + function_def = tool.get("function", {}) + api_tools.append( + { + "name": function_def.get("name", ""), + "description": function_def.get("description", ""), + "parameters": function_def.get("parameters", {}), + } + ) + # Handle direct tool format + elif "name" in tool: + api_tools.append( + { + "name": tool.get("name", ""), + "description": tool.get("description", ""), + "parameters": tool.get("parameters", {}), + } + ) + + return api_tools if api_tools else None + + def _check_if_flagged(self, result: Dict[str, Any]) -> bool: """ Check if the Qualifire evaluation result indicates flagged content. @@ -190,65 +224,53 @@ class QualifireGuardrail(CustomGuardrail): A high score (close to 100) indicates GOOD content, low score indicates problems. """ # Check evaluation results for any flagged items - evaluation_results = getattr(result, "evaluationResults", None) or [] - if isinstance(result, dict): - evaluation_results = result.get("evaluationResults", []) or [] + evaluation_results = result.get("evaluationResults", []) or [] for eval_result in evaluation_results: - results: List[Any] = [] - if isinstance(eval_result, dict): - results = eval_result.get("results", []) or [] - else: - results = getattr(eval_result, "results", []) or [] - + results = eval_result.get("results", []) or [] for r in results: - flagged = ( - r.get("flagged") - if isinstance(r, dict) - else getattr(r, "flagged", False) - ) - if flagged: + if r.get("flagged"): return True return False - def _build_evaluate_kwargs( + def _build_evaluate_payload( self, - qualifire_messages: List[Any], + api_messages: List[Dict[str, Any]], output: Optional[str], assertions: Optional[List[str]], - available_tools: Optional[List[Any]], + available_tools: Optional[List[Dict[str, Any]]], ) -> Dict[str, Any]: - """Build kwargs dictionary for the evaluate call.""" - kwargs: Dict[str, Any] = {"messages": qualifire_messages} + """Build payload dictionary for the /api/evaluation/evaluate endpoint.""" + payload: Dict[str, Any] = {"messages": api_messages} if output is not None: - kwargs["output"] = output + payload["output"] = output # Add enabled checks if self.prompt_injections: - kwargs["prompt_injections"] = True + payload["prompt_injections"] = True if self.hallucinations_check: - kwargs["hallucinations_check"] = True + payload["hallucinations_check"] = True if self.grounding_check: - kwargs["grounding_check"] = True + payload["grounding_check"] = True if self.pii_check: - kwargs["pii_check"] = True + payload["pii_check"] = True if self.content_moderation_check: - kwargs["content_moderation_check"] = True + payload["content_moderation_check"] = True if self.tool_selection_quality_check: # Only enable tool_selection_quality_check if available_tools is provided if available_tools: - kwargs["tool_selection_quality_check"] = True - kwargs["available_tools"] = available_tools + payload["tool_selection_quality_check"] = True + payload["available_tools"] = available_tools else: verbose_proxy_logger.debug( "Qualifire Guardrail: tool_selection_quality_check enabled but no available_tools provided, skipping this check" ) if assertions: - kwargs["assertions"] = assertions + payload["assertions"] = assertions - return kwargs + return payload async def _run_qualifire_check( self, @@ -274,11 +296,17 @@ class QualifireGuardrail(CustomGuardrail): assertions = dynamic_params.get("assertions") or self.assertions on_flagged = dynamic_params.get("on_flagged") or self.on_flagged - try: - client = self._get_client() - qualifire_messages = self._convert_messages_to_qualifire_format(messages) + # Prepare headers + headers = { + "X-Qualifire-API-Key": self.qualifire_api_key or "", + "Content-Type": "application/json", + } - # Use invoke_evaluation if evaluation_id is provided + try: + # Convert messages to API format + api_messages = self._convert_messages_to_api_format(messages) + + # Use invoke endpoint if evaluation_id is provided if evaluation_id: # For invoke_evaluation, we need to extract input/output input_text = "" @@ -291,25 +319,47 @@ class QualifireGuardrail(CustomGuardrail): input_text = content break - result = client.invoke_evaluation( - evaluation_id=evaluation_id, - input=input_text, - output=output or "", - ) + payload = { + "evaluation_id": evaluation_id, + "input": input_text, + "output": output or "", + "messages": api_messages, + } + + # Convert tools if provided + api_tools = self._convert_tools_to_api_format(available_tools) + if api_tools: + payload["available_tools"] = api_tools + + url = f"{self.qualifire_api_base}/api/evaluation/invoke" else: - # Use evaluate with individual checks - kwargs = self._build_evaluate_kwargs( - qualifire_messages=qualifire_messages, + # Use evaluate endpoint with individual checks + api_tools = self._convert_tools_to_api_format(available_tools) + payload = self._build_evaluate_payload( + api_messages=api_messages, output=output, assertions=assertions, - available_tools=available_tools, + available_tools=api_tools, ) - result = client.evaluate(**kwargs) + url = f"{self.qualifire_api_base}/api/evaluation/evaluate" - # Convert result to dict for logging + verbose_proxy_logger.debug( + f"Qualifire Guardrail: Making request to {url}" + ) + + # Make the API request + response = await self.async_handler.post( + url=url, + headers=headers, + json=payload, + ) + response.raise_for_status() + result = response.json() + + # Extract response info for logging qualifire_response = { - "score": getattr(result, "score", None), - "status": getattr(result, "status", None), + "score": result.get("score"), + "status": result.get("status"), } verbose_proxy_logger.debug( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py index 35ed49a84ed..fd72185d1e7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_qualifire.py @@ -2,7 +2,6 @@ Unit tests for Qualifire guardrail integration. """ -import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -75,9 +74,37 @@ class TestQualifireGuardrailInit: assert guardrail.on_flagged == "monitor" + def test_init_with_default_api_base(self): + """Test that default API base is set when not provided.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + DEFAULT_QUALIFIRE_API_BASE, + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + assert guardrail.qualifire_api_base == DEFAULT_QUALIFIRE_API_BASE + + def test_init_with_custom_api_base(self): + """Test initialization with custom API base URL.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + api_base="https://custom.qualifire.ai", + guardrail_name="test_guardrail", + ) + + assert guardrail.qualifire_api_base == "https://custom.qualifire.ai" + class TestQualifireGuardrailMessageConversion: - """Tests for message conversion to Qualifire format.""" + """Tests for message conversion to API format.""" def test_convert_simple_messages(self): """Test conversion of simple text messages.""" @@ -95,15 +122,13 @@ class TestQualifireGuardrailMessageConversion: {"role": "assistant", "content": "Hi there!"}, ] - # Create mock LLMMessage class - mock_llm_message = MagicMock() + result = guardrail._convert_messages_to_api_format(messages) - with patch( - "litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire.QualifireGuardrail._convert_messages_to_qualifire_format" - ) as mock_convert: - mock_convert.return_value = [mock_llm_message, mock_llm_message] - result = guardrail._convert_messages_to_qualifire_format(messages) - assert len(result) == 2 + assert len(result) == 2 + assert result[0]["role"] == "user" + assert result[0]["content"] == "Hello, world!" + assert result[1]["role"] == "assistant" + assert result[1]["content"] == "Hi there!" def test_convert_multimodal_messages(self): """Test conversion of multimodal messages with text parts.""" @@ -126,112 +151,258 @@ class TestQualifireGuardrailMessageConversion: }, ] - with patch( - "litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire.QualifireGuardrail._convert_messages_to_qualifire_format" - ) as mock_convert: - mock_convert.return_value = [MagicMock()] - result = guardrail._convert_messages_to_qualifire_format(messages) - assert len(result) == 1 + result = guardrail._convert_messages_to_api_format(messages) + + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"] == "First part\nSecond part" + + def test_convert_messages_with_tool_calls(self): + """Test conversion of messages with tool calls.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + messages = [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location": "NYC"}', + }, + } + ], + }, + ] + + result = guardrail._convert_messages_to_api_format(messages) + + assert len(result) == 1 + assert result[0]["role"] == "assistant" + assert "tool_calls" in result[0] + assert len(result[0]["tool_calls"]) == 1 + assert result[0]["tool_calls"][0]["id"] == "call_123" + assert result[0]["tool_calls"][0]["name"] == "get_weather" + assert result[0]["tool_calls"][0]["arguments"] == {"location": "NYC"} -class TestQualifireGuardrailEvaluateKwargs: - """Tests for evaluate kwargs passed to Qualifire client.""" +class TestQualifireGuardrailToolConversion: + """Tests for tool definition conversion.""" + + def test_convert_openai_function_tools(self): + """Test conversion of OpenAI function tool format.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get weather for a location", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + result = guardrail._convert_tools_to_api_format(tools) + + assert result is not None + assert len(result) == 1 + assert result[0]["name"] == "get_weather" + assert result[0]["description"] == "Get weather for a location" + + def test_convert_empty_tools(self): + """Test that empty tools returns None.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + guardrail_name="test_guardrail", + ) + + result = guardrail._convert_tools_to_api_format(None) + assert result is None + + result = guardrail._convert_tools_to_api_format([]) + assert result is None + + +class TestQualifireGuardrailAPICall: + """Tests for API call with httpx client.""" @pytest.mark.asyncio async def test_evaluate_called_with_prompt_injections(self): - """Test that evaluate is called with prompt_injections enabled.""" - # Mock the qualifire module and its types - mock_qualifire_types = MagicMock() - mock_llm_message = MagicMock() - mock_llm_tool_call = MagicMock() - mock_message_instance = MagicMock() - mock_llm_message.return_value = mock_message_instance - - mock_qualifire_types.LLMMessage = mock_llm_message - mock_qualifire_types.LLMToolCall = mock_llm_tool_call - - with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + """Test that evaluate endpoint is called with prompt_injections enabled.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output=None, dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output=None, dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert "messages" in call_kwargs + # Verify the API was called + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + assert "json" in call_kwargs + payload = call_kwargs["json"] + assert payload["prompt_injections"] is True + assert "messages" in payload + assert call_kwargs["url"].endswith("/api/evaluation/evaluate") @pytest.mark.asyncio async def test_evaluate_called_with_multiple_checks(self): """Test that evaluate is called with multiple checks enabled.""" - # Mock the qualifire module and its types - mock_qualifire_types = MagicMock() - mock_llm_message = MagicMock() - mock_llm_tool_call = MagicMock() - mock_message_instance = MagicMock() - mock_llm_message.return_value = mock_message_instance - - mock_qualifire_types.LLMMessage = mock_llm_message - mock_qualifire_types.LLMToolCall = mock_llm_tool_call - - with patch.dict('sys.modules', {'qualifire': MagicMock(), 'qualifire.types': mock_qualifire_types}): - from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( - QualifireGuardrail, - ) + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) - guardrail = QualifireGuardrail( - api_key="test_key", - prompt_injections=True, - pii_check=True, - hallucinations_check=True, - assertions=["Output must be valid JSON"], - guardrail_name="test_guardrail", - ) + guardrail = QualifireGuardrail( + api_key="test_key", + prompt_injections=True, + pii_check=True, + hallucinations_check=True, + assertions=["Output must be valid JSON"], + guardrail_name="test_guardrail", + ) - # Mock the client - mock_client = MagicMock() - mock_result = MagicMock() - mock_result.score = 100 - mock_result.status = "completed" - mock_result.evaluationResults = [] - mock_client.evaluate.return_value = mock_result - guardrail._client = mock_client + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) - messages = [{"role": "user", "content": "Hello, world!"}] + messages = [{"role": "user", "content": "Hello, world!"}] - await guardrail._run_qualifire_check( - messages=messages, output="Test output", dynamic_params={} - ) + await guardrail._run_qualifire_check( + messages=messages, output="Test output", dynamic_params={} + ) - # Verify evaluate was called with correct kwargs - mock_client.evaluate.assert_called_once() - call_kwargs = mock_client.evaluate.call_args[1] - assert call_kwargs["prompt_injections"] is True - assert call_kwargs["pii_check"] is True - assert call_kwargs["hallucinations_check"] is True - assert call_kwargs["assertions"] == ["Output must be valid JSON"] - assert call_kwargs["output"] == "Test output" + # Verify the API was called with correct payload + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + payload = call_kwargs["json"] + assert payload["prompt_injections"] is True + assert payload["pii_check"] is True + assert payload["hallucinations_check"] is True + assert payload["assertions"] == ["Output must be valid JSON"] + assert payload["output"] == "Test output" + + @pytest.mark.asyncio + async def test_invoke_endpoint_used_with_evaluation_id(self): + """Test that invoke endpoint is used when evaluation_id is provided.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="test_key", + evaluation_id="eval_123", + guardrail_name="test_guardrail", + ) + + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + messages = [{"role": "user", "content": "Hello, world!"}] + + await guardrail._run_qualifire_check( + messages=messages, output="Test output", dynamic_params={} + ) + + # Verify the invoke endpoint was called + guardrail.async_handler.post.assert_called_once() + call_kwargs = guardrail.async_handler.post.call_args[1] + + assert call_kwargs["url"].endswith("/api/evaluation/invoke") + payload = call_kwargs["json"] + assert payload["evaluation_id"] == "eval_123" + assert payload["input"] == "Hello, world!" + assert payload["output"] == "Test output" + + @pytest.mark.asyncio + async def test_correct_headers_sent(self): + """Test that correct headers are sent with the API request.""" + from litellm.proxy.guardrails.guardrail_hooks.qualifire.qualifire import ( + QualifireGuardrail, + ) + + guardrail = QualifireGuardrail( + api_key="my_api_key", + guardrail_name="test_guardrail", + ) + + # Mock the async HTTP handler + mock_response = MagicMock() + mock_response.json.return_value = { + "score": 100, + "status": "completed", + "evaluationResults": [], + } + mock_response.raise_for_status = MagicMock() + guardrail.async_handler.post = AsyncMock(return_value=mock_response) + + messages = [{"role": "user", "content": "Hello!"}] + + await guardrail._run_qualifire_check( + messages=messages, output=None, dynamic_params={} + ) + + call_kwargs = guardrail.async_handler.post.call_args[1] + headers = call_kwargs["headers"] + + assert headers["X-Qualifire-API-Key"] == "my_api_key" + assert headers["Content-Type"] == "application/json" class TestQualifireGuardrailCheckIfFlagged: @@ -248,12 +419,14 @@ class TestQualifireGuardrailCheckIfFlagged: guardrail_name="test_guardrail", ) - # Mock result with completed status and no flagged items - mock_result = MagicMock() - mock_result.status = "completed" - mock_result.evaluationResults = [] + # Result with completed status and no flagged items (dict format) + result = { + "status": "completed", + "score": 100, + "evaluationResults": [], + } - assert guardrail._check_if_flagged(mock_result) is False + assert guardrail._check_if_flagged(result) is False def test_check_if_flagged_returns_true_for_flagged_content(self): """Test that _check_if_flagged returns True when content is flagged.""" @@ -266,18 +439,25 @@ class TestQualifireGuardrailCheckIfFlagged: guardrail_name="test_guardrail", ) - # Mock result with flagged item - mock_inner_result = MagicMock() - mock_inner_result.flagged = True + # Result with flagged item (dict format matching API response) + result = { + "status": "completed", + "score": 15, + "evaluationResults": [ + { + "type": "prompt_injection", + "results": [ + { + "flagged": True, + "score": 0.15, + "reason": "Prompt injection detected", + } + ], + } + ], + } - mock_eval_result = MagicMock() - mock_eval_result.results = [mock_inner_result] - - mock_result = MagicMock() - mock_result.status = "completed" - mock_result.evaluationResults = [mock_eval_result] - - assert guardrail._check_if_flagged(mock_result) is True + assert guardrail._check_if_flagged(result) is True def test_check_if_flagged_returns_false_when_no_flagged_items(self): """Test that _check_if_flagged returns False when no items are flagged.""" @@ -291,17 +471,24 @@ class TestQualifireGuardrailCheckIfFlagged: ) # Result with evaluation results but nothing flagged - mock_inner_result = MagicMock() - mock_inner_result.flagged = False + result = { + "status": "completed", + "score": 95, + "evaluationResults": [ + { + "type": "prompt_injection", + "results": [ + { + "flagged": False, + "score": 0.95, + "reason": "No issues detected", + } + ], + } + ], + } - mock_eval_result = MagicMock() - mock_eval_result.results = [mock_inner_result] - - mock_result = MagicMock() - mock_result.status = "success" - mock_result.evaluationResults = [mock_eval_result] - - assert guardrail._check_if_flagged(mock_result) is False + assert guardrail._check_if_flagged(result) is False class TestQualifireGuardrailShouldRun: