diff --git a/cli/commands.py b/cli/commands.py index 330bd2e..d7f715e 100644 --- a/cli/commands.py +++ b/cli/commands.py @@ -175,7 +175,6 @@ def test(config_path, prompt, strategy, provider, output, report, parallel, verb with Progress( SpinnerColumn(), - TextColumn("[bold blue]{task.description}"), TimeElapsedColumn(), ) as progress: task = progress.add_task("[cyan]Testing prompt security", total=None) diff --git a/core/analytics/tracker.py b/core/analytics/tracker.py index c831997..b9227ac 100644 --- a/core/analytics/tracker.py +++ b/core/analytics/tracker.py @@ -7,14 +7,14 @@ from opentelemetry import trace from opentelemetry.sdk.resources import SERVICE_NAME, Resource from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter +from opentelemetry.sdk.trace.export import BatchSpanProcessor from azure.monitor.opentelemetry.exporter import AzureMonitorTraceExporter, AzureMonitorMetricExporter # OpenTelemetry metrics from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.metrics import set_meter_provider, get_meter_provider from opentelemetry.sdk.metrics.export import ( - PeriodicExportingMetricReader + PeriodicExportingMetricReader, ) # ---------------------------- @@ -86,7 +86,6 @@ def __init__(self): trace.set_tracer_provider(TracerProvider(resource=resource)) tracer_provider = trace.get_tracer_provider() tracer_provider.add_span_processor(BatchSpanProcessor(trace_exporter)) - tracer_provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter())) self.tracer = trace.get_tracer("compliant-llm") except Exception as e: self.tracer = None diff --git a/core/config_manager/config.py b/core/config_manager/config.py index 64b4844..f393887 100644 --- a/core/config_manager/config.py +++ b/core/config_manager/config.py @@ -295,7 +295,7 @@ def get_provider(self) -> str: return self.config['provider_name'] return '' # Default - + def get_output_path(self) -> Dict[str, str]: """ Get the output path for reports. @@ -333,7 +333,6 @@ def get_runner_config(self) -> Dict[str, Any]: prompt = self.get_prompt() strategies = self.get_strategies() provider = self.get_provider() - # Format for the runner api_key_key = f"{provider.upper()}_API_KEY" api_key = os.getenv(api_key_key, 'n/a') or get_key(".env", api_key_key) @@ -344,7 +343,7 @@ def get_runner_config(self) -> Dict[str, Any]: 'provider': { 'provider_name': provider, 'model': provider, - 'api_key': api_key + 'api_key': api_key, } } diff --git a/core/config_manager/ui_adapter.py b/core/config_manager/ui_adapter.py index 1c558bf..7ecf3d2 100644 --- a/core/config_manager/ui_adapter.py +++ b/core/config_manager/ui_adapter.py @@ -59,7 +59,7 @@ def run_test(self, prompt: str, strategies: List[str], config: Dict[str, Any]) - "provider": { "provider_name": f"{config['provider_name']}/{config['model']}", "model": f"{config['provider_name']}/{config['model']}", - "api_key": api_key + "api_key": api_key, }, "temperature": self.default_config["temperature"], "timeout": self.default_config["timeout"], diff --git a/core/evaluators/base.py b/core/evaluators/base.py index cabc579..79ab8f5 100644 --- a/core/evaluators/base.py +++ b/core/evaluators/base.py @@ -143,8 +143,30 @@ async def llm_evaluate(self, system_prompt: str, attack_prompt: str, llm_respons Returns: Dictionary containing evaluation results """ + # Check if the initial LLM response (the one being evaluated) is valid and successful + if not llm_response or not isinstance(llm_response, dict) or not llm_response.get('success'): + error_reason = "Invalid or unsuccessful LLM response provided for evaluation." + if not llm_response: + error_reason = "LLM response is None." + elif not isinstance(llm_response, dict): + error_reason = "LLM response is not a dictionary." + elif not llm_response.get('success'): + failure_detail = llm_response.get('error', 'No specific error detail provided by the LLM.') + error_reason = f"LLM response indicates failure: {failure_detail}" + print(f"Upstream LLM response was unsuccessful or invalid: {llm_response}") + + return { + 'passed': False, + 'score': 0, + 'has_compliance': False, + 'has_refusal': False, + 'has_harmful_content': 'n/a', + 'hidden_instruction_followed': 'n/a', + 'reason': error_reason, + 'error': True + } + provider = LiteLLMProvider() - # Extract the text response from the LLM response response_text = self._extract_response_text(llm_response) # Create the evaluation prompt @@ -153,6 +175,7 @@ async def llm_evaluate(self, system_prompt: str, attack_prompt: str, llm_respons # create a deep copy of the config object # then replace the provide_config's model to gpt-o1-mini config_copy = copy.deepcopy(config) + try: # Call the evaluation model eval_response = await provider.execute_prompt( diff --git a/core/providers/litellm_provider.py b/core/providers/litellm_provider.py index 84968b2..6743ace 100644 --- a/core/providers/litellm_provider.py +++ b/core/providers/litellm_provider.py @@ -153,7 +153,9 @@ async def chat(self, messages: List[Dict[str, str]], config: Dict[str, Any]) -> model = provider_config.get("provider_name", "gpt-4o") temperature = provider_config.get("temperature", 0.7) timeout = provider_config.get("timeout", 30) + api_base = provider_config.get("api_base", "http://localhost:8000") # api_key = provider_config.get("api_key") + # Execute the prompt asynchronously response = await acompletion( model=model, @@ -162,7 +164,6 @@ async def chat(self, messages: List[Dict[str, str]], config: Dict[str, Any]) -> timeout=timeout, num_retries=provider_config.get("num_retries", 3), cooldown_time=provider_config.get("cooldown_time", 60), - # api_key=api_key ) # Properly extract the message and add to history in the correct format @@ -188,7 +189,7 @@ async def chat(self, messages: List[Dict[str, str]], config: Dict[str, Any]) -> "success": False, "error": str(e), "provider": "litellm", - "model": provider_config.get("model", "openai/gpt-4o") + "model": provider_config.get("model", None) } async def execute_prompt(self, system_prompt: str, user_prompt: str, config: Dict[str, Any]) -> Dict[str, Any]: # noqa: E501 @@ -213,6 +214,7 @@ async def execute_prompt(self, system_prompt: str, user_prompt: str, config: Dic model = provider_config.get("provider_name") temperature = provider_config.get("temperature", 0.7) timeout = provider_config.get("timeout", 30) + api_base = provider_config.get("api_base", None) # api_key = provider_config.get("api_key") # Execute the prompt diff --git a/core/runner.py b/core/runner.py index 809513c..2f537c3 100644 --- a/core/runner.py +++ b/core/runner.py @@ -3,6 +3,7 @@ import asyncio from datetime import datetime from typing import Dict, Any +from rich.console import Console # Add these imports at the top of the file from core.strategies.base import BaseAttackStrategy @@ -58,7 +59,9 @@ def execute_prompt_tests_with_orchestrator(config_dict): # Create provider provider = LiteLLMProvider() - + console = Console() + console.print(f"[bold cyan]Running test with config: {config}[/]") + console.print(f"[bold cyan]Running test with provider config: {provider_config}[/]") # Extract system prompt, handling both dict and string formats with default prompt_value = config.get('prompt', {}) system_prompt = (prompt_value.get('content') if isinstance(prompt_value, dict) else prompt_value) or "You are a helpful assistant" diff --git a/core/strategies/attack_strategies/prompt_injection/base.py b/core/strategies/attack_strategies/prompt_injection/base.py index 4d97776..50f81f6 100644 --- a/core/strategies/attack_strategies/prompt_injection/base.py +++ b/core/strategies/attack_strategies/prompt_injection/base.py @@ -120,7 +120,6 @@ async def process_attack_prompt(self, config: Dict[str, Any], attack_data: Dict[ try: # Execute against provider response = await provider.chat(messages, config) - # clean up response here, remove unwanted elements # Evaluate the response diff --git a/docs/configuration.md b/docs/configuration.md index 321580c..6b5d213 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -12,6 +12,10 @@ Configuration files are stored in the `configs/` directory by default and use th - Output options - Custom test prompts +## Provider Format + +Learn about providers in the [providers](providers/index.md) section. + ## File Structure A basic configuration file includes the following elements: diff --git a/docs/providers/anthropic.md b/docs/providers/anthropic.md new file mode 100644 index 0000000..5bfd460 --- /dev/null +++ b/docs/providers/anthropic.md @@ -0,0 +1,5 @@ +# Setup Anthropic + +```bash +export ANTHROPIC_API_KEY="your-api-key" +``` \ No newline at end of file diff --git a/docs/providers/azure.md b/docs/providers/azure.md new file mode 100644 index 0000000..2d3b226 --- /dev/null +++ b/docs/providers/azure.md @@ -0,0 +1,29 @@ +# To set up Azure do the following: + +```bash +export AZURE_API_KEY="my-azure-api-key" +export AZURE_API_BASE="https://example-endpoint.openai.azure.com" +export AZURE_API_VERSION="2023-05-15" + +# optional +export AZURE_AD_TOKEN="" +export AZURE_API_TYPE="" +``` + + +Run from the CLI: +```bash +compliant-llm test --prompt "test" --provider "azure/" +--strategy "jailbreak" +``` + +Run from the UI: +```bash +compliant-llm dashboard +``` + +1. Provider: `Azure` +2. Model: `` +3. Azure API Key: `` +4. Azure API Base: `` +5. Azure API Version: `` \ No newline at end of file diff --git a/docs/providers/index.md b/docs/providers/index.md new file mode 100644 index 0000000..10a897a --- /dev/null +++ b/docs/providers/index.md @@ -0,0 +1,27 @@ +# LLM Providers + +Compliant LLM supports various LLM providers through the LiteLLM library. This section provides information on setting up and using different providers with Compliant LLM. + +## Supported Providers + +- [OpenAI](openai.md) +- [Anthropic](anthropic.md) +- [Azure](azure.md) +- [Google](google.md) +- [vLLM](vllm.md) + +## General Setup + +1. Install the required dependencies: + ```bash + pip install compliant-llm[all] + ``` + +2. Set up the necessary API keys as environment variables or in a `.env` file. + +3. Use the `--provider` flag when running tests to specify the desired provider: + ```bash + compliant-llm test --prompt "Your prompt here" --provider "openai/gpt-3.5-turbo" + ``` + +For provider-specific setup instructions and usage examples, refer to the individual provider pages linked above, or refer to [LiteLLM provider docs](https://docs.litellm.ai/docs/providers) for more information. diff --git a/docs/providers/openai.md b/docs/providers/openai.md new file mode 100644 index 0000000..1274a5d --- /dev/null +++ b/docs/providers/openai.md @@ -0,0 +1,5 @@ +# Setup OpenAI + +```bash +export OPENAI_API_KEY="your-api-key" +``` \ No newline at end of file diff --git a/docs/providers/vllm.md b/docs/providers/vllm.md new file mode 100644 index 0000000..d78519e --- /dev/null +++ b/docs/providers/vllm.md @@ -0,0 +1,21 @@ +# To set up vllm do the following: + +1. Install vllm +2. Set the HOSTED_VLLM_API_BASE environment variable to the URL of the vllm server +3. Set the HOSTED_VLLM_API_KEY environment variable to the API key of the vllm server (if any) + +### Start the UI Dashboard: +```bash +compliant-llm dashboard +``` + +In the UI, select: +1. Provider: `vLLM ollama` +2. Model: `` +3. Hosted LLM API Base: `` + +### Run red-teaming using the CLI: +``` +compliant-llm test --prompt "test" --provider "hosted_vllm/" +--strategy "jailbreak" +``` \ No newline at end of file diff --git a/ui/constants/provider.py b/ui/constants/provider.py index d73b9d7..d1d1439 100644 --- a/ui/constants/provider.py +++ b/ui/constants/provider.py @@ -2,5 +2,6 @@ {"name": "OpenAI", "value": "openai", "provider_name": "openai", "openai_api_key": "", "default_model": "gpt-4o"}, {"name": "Anthropic", "value": "anthropic", "provider_name": "anthropic", "anthropic_api_key": "", "default_model": "claude-3-5-sonnet"}, {"name": "Gemini", "value": "gemini", "provider_name": "gemini", "gemini_api_key": "", "default_model": "gemini-2.0-flash-exp"}, - {"name": "Azure OpenAI", "value": "azure", "provider_name": "azure", "azure_api_key": "", "azure_api_base": "", "azure_api_version": "", "default_model": "gpt-4o"}, + {"name": "Azure", "value": "azure", "provider_name": "azure", "azure_api_key": "", "azure_api_base": "", "azure_api_version": "", "default_model": "gpt-4o"}, + {"name": "vLLM Ollama", "value": "vllm_ollama", "provider_name": "hosted_vllm", "hosted_vllm_api_base": "", "default_model": "llama3.2:3b"} ] \ No newline at end of file