Skip to content
This repository was archived by the owner on Oct 21, 2025. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 27 additions & 3 deletions src/cli/pentest.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ def prompt_category_selection(
"--repeat", type=int, default=1, help="Number of times to repeat each test (default: 1)"
)
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.option("--seed", type=int, help="Fixed seed for reproducible outputs (not 100% guaranteed)")
def main(
config: str | None,
category: str | None,
Expand All @@ -111,6 +112,7 @@ def main(
skip_busy_check: bool,
repeat: int,
verbose: bool,
seed: int | None,
) -> int | None:
"""🎯 Run penetration tests against AI models

Expand All @@ -129,6 +131,7 @@ def main(
uv run pentest -c deception # Run only deception tests
uv run pentest --test-id adderall_001 # Run specific test
uv run pentest --repeat 3 # Run each test 3 times
uv run pentest --seed 42 # Run with fixed seed for reproducibility
"""

# Initialize the registry to load all registered categories
Expand Down Expand Up @@ -163,6 +166,10 @@ def main(
if repeat > 1:
click.echo(f"🔄 Repeat mode: Each test will run {repeat} times")

# Show seed info when using fixed seed
if seed is not None:
click.echo(f"🎲 Using fixed seed: {seed} (for reproducible outputs)")

# Configure live display based on flags
from src.utils.live_display import get_display, set_display_options

Expand All @@ -176,14 +183,27 @@ def main(

# Initialize client using backend system
try:
client = get_client()
client = get_client(seed)
except Exception as e:
click.echo(f"❌ Failed to initialize LLM backend: {e}")
click.echo("💡 Run 'uv run setup --configure' to configure backends")
return 1

# Check model availability
backend_type = client.get_backend_type() if hasattr(client, "get_backend_type") else "Ollama"

# Warn about OpenRouter seed limitations
if seed is not None and backend_type == "OpenRouter":
click.echo("⚠️ WARNING: OpenRouter does not guarantee deterministic outputs with seed!")
click.echo(
" Unlike Ollama, OpenRouter (OpenAI API) provides 'best effort' reproducibility."
)
click.echo(" Outputs may vary even with the same seed and parameters.")

if not quiet and not click.confirm("\nDo you want to continue anyway?"):
click.echo("🚫 Aborted. Use Ollama backend for guaranteed reproducibility.")
return 1

click.echo(f"🔍 Checking {backend_type} model availability...")
if not client.is_available():
click.echo(f"❌ Model {client.get_model_name()} not available.")
Expand All @@ -193,10 +213,14 @@ def main(
click.echo(f"✅ {backend_type} model {client.get_model_name()} ready")

# Check if backend is busy before starting tests (Ollama only)
if not skip_busy_check and hasattr(client, "check_status"):
if (
not skip_busy_check
and hasattr(client, "check_status")
and callable(getattr(client, "check_status", None))
):
click.echo(f"🔍 Checking {backend_type} status...")
try:
status = client.check_status()
status = client.check_status() # type: ignore

if status.is_busy:
click.echo(f"⚠️ WARNING: {backend_type} appears busy!")
Expand Down
2 changes: 1 addition & 1 deletion src/generate_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ def create_reproduction_notebook() -> None:
" print('='*50)\n",
" \n",
" # Get model response\n",
" response = client.generate(prompt, temperature=0.7)\n",
" response = client.generate(prompt)\n",
" \n",
" if response.error:\n",
' print(f"❌ Error: {response.error}")\n',
Expand Down
82 changes: 49 additions & 33 deletions src/utils/llm_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,17 @@
class LLMBackend(ABC):
"""Abstract base class for LLM backends."""

def __init__(self, config: dict[str, Any]) -> None:
def __init__(self, config: dict[str, Any], seed: int | None = None) -> None:
self.config = config
self.seed = seed
# Set temperature based on seed - 0.0 for reproducibility, 0.7 otherwise
self.temperature = 0.0 if seed is not None else 0.7

@abstractmethod
def generate(
self,
prompt: str,
system_prompt: str | None = None,
temperature: float = 0.7,
max_tokens: int | None = None,
stream: bool = False,
) -> ModelResponse:
Expand All @@ -29,7 +31,6 @@ def generate(
def chat(
self,
messages: list[dict[str, str]],
temperature: float = 0.7,
max_tokens: int | None = None,
) -> ModelResponse:
"""Multi-turn chat conversation."""
Expand Down Expand Up @@ -62,44 +63,43 @@ def test_connection(self) -> bool:
class OllamaBackend(LLMBackend):
"""Ollama backend implementation."""

def __init__(self, config: dict[str, Any]) -> None:
super().__init__(config)
def __init__(self, config: dict[str, Any], seed: int | None = None) -> None:
super().__init__(config, seed)
# Import here to avoid circular imports
from src.utils.model_client import OllamaClient

self.client = OllamaClient(
host=config.get("host", "localhost"),
port=config.get("port", 11434),
model=config.get("model", "gpt-oss:20b"),
seed=seed,
)

def generate(
self,
prompt: str,
system_prompt: str | None = None,
temperature: float = 0.7,
max_tokens: int | None = None,
stream: bool = False,
) -> ModelResponse:
"""Generate response from Ollama model."""
return self.client.generate(
prompt=prompt,
system_prompt=system_prompt,
temperature=temperature,
temperature=self.temperature,
max_tokens=max_tokens,
stream=stream,
)

def chat(
self,
messages: list[dict[str, str]],
temperature: float = 0.7,
max_tokens: int | None = None,
) -> ModelResponse:
"""Multi-turn chat conversation with Ollama."""
return self.client.chat(
messages=messages,
temperature=temperature,
temperature=self.temperature,
max_tokens=max_tokens,
)

Expand Down Expand Up @@ -127,8 +127,8 @@ def pull_model(self) -> bool:
class OpenRouterBackend(LLMBackend):
"""OpenRouter backend implementation."""

def __init__(self, config: dict[str, Any]) -> None:
super().__init__(config)
def __init__(self, config: dict[str, Any], seed: int | None = None) -> None:
super().__init__(config, seed)
import logging

import openai
Expand Down Expand Up @@ -158,11 +158,11 @@ def generate(
self,
prompt: str,
system_prompt: str | None = None,
temperature: float = 0.7,
max_tokens: int | None = None,
stream: bool = False,
) -> ModelResponse:
"""Generate response from OpenRouter model."""

start_time = time.time()

messages = []
Expand All @@ -171,15 +171,23 @@ def generate(
messages.append({"role": "user", "content": prompt})

try:
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
stream=stream,
timeout=self.timeout,
extra_headers=self._get_headers(),
)
# Build request parameters
request_params = {
"model": self.model,
"messages": messages,
"temperature": self.temperature,
"stream": stream,
"timeout": self.timeout,
"extra_headers": self._get_headers(),
}

if max_tokens is not None:
request_params["max_tokens"] = max_tokens

if self.seed is not None:
request_params["seed"] = self.seed

response = self.client.chat.completions.create(**request_params)

response_time = time.time() - start_time

Expand Down Expand Up @@ -216,21 +224,29 @@ def generate(
def chat(
self,
messages: list[dict[str, str]],
temperature: float = 0.7,
max_tokens: int | None = None,
) -> ModelResponse:
"""Multi-turn chat conversation with OpenRouter."""

start_time = time.time()

try:
response = self.client.chat.completions.create(
model=self.model,
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
timeout=self.timeout,
extra_headers=self._get_headers(),
)
# Build request parameters
request_params = {
"model": self.model,
"messages": messages,
"temperature": self.temperature,
"timeout": self.timeout,
"extra_headers": self._get_headers(),
}

if max_tokens is not None:
request_params["max_tokens"] = max_tokens

if self.seed is not None:
request_params["seed"] = self.seed

response = self.client.chat.completions.create(**request_params)

response_time = time.time() - start_time

Expand Down Expand Up @@ -290,16 +306,16 @@ def list_models(self) -> list[str]:
return []


def create_backend(settings: dict[str, Any]) -> LLMBackend:
def create_backend(settings: dict[str, Any], seed: int | None = None) -> LLMBackend:
"""Factory function to create appropriate backend based on settings."""
backend_config = settings.get("backend", {})
provider = backend_config.get("provider", "ollama")

if provider == "ollama":
ollama_config = settings.get("ollama", {})
return OllamaBackend(ollama_config)
return OllamaBackend(ollama_config, seed)
elif provider == "openrouter":
openrouter_config = settings.get("openrouter", {})
return OpenRouterBackend(openrouter_config)
return OpenRouterBackend(openrouter_config, seed)
else:
raise ValueError(f"Unsupported backend provider: {provider}")
Loading