[Feat]: Adding vllm support.

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)

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

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,
             }
         }
 

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"],

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(

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

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"

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

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:

docs/providers/anthropic.md

@@ -0,0 +1,5 @@
+# Setup Anthropic
+
+```bash
+export ANTHROPIC_API_KEY="your-api-key"
+```

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/<your_deployment_name>"
+--strategy  "jailbreak"
+```
+
+Run from the UI:
+```bash
+compliant-llm dashboard
+```
+
+1. Provider: `Azure`
+2. Model: `<your_deployment_name>`
+3. Azure API Key: `<your_api_key>`
+4. Azure API Base: `<your_api_base>`
+5. Azure API Version: `<your_api_version>`

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.

docs/providers/openai.md

@@ -0,0 +1,5 @@
+# Setup OpenAI
+
+```bash
+export OPENAI_API_KEY="your-api-key"
+```

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: `<your_model>`
+3. Hosted LLM API Base: `<your_api_base>`
+
+### Run red-teaming using the CLI:
+```
+compliant-llm test --prompt "test" --provider "hosted_vllm/<your_model_name>"
+--strategy  "jailbreak"
+```

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"}
 ]