diff --git a/deepteam/__init__.py b/deepteam/__init__.py index d8fd9bd9..c9f61025 100644 --- a/deepteam/__init__.py +++ b/deepteam/__init__.py @@ -9,11 +9,13 @@ from deepteam.guardrails import ( Guardrails, ) +from deepteam.utils import set_progress_callback from deepteam._version import __version__ __all__ = [ "red_team", "Guardrails", + "set_progress_callback", "__version__", ] diff --git a/deepteam/red_team.py b/deepteam/red_team.py index 81e7383e..f2f915c5 100644 --- a/deepteam/red_team.py +++ b/deepteam/red_team.py @@ -1,4 +1,5 @@ from typing import List, Optional +from contextlib import nullcontext from deepeval.models import DeepEvalBaseLLM from deepteam.vulnerabilities import BaseVulnerability @@ -7,6 +8,7 @@ from deepteam.red_teamer import RedTeamer from deepteam.attacks.attack_engine import AttackEngine from deepteam.frameworks.frameworks import AISafetyFramework +from deepteam.utils import progress_callback_context, ProgressCallback def red_team( @@ -22,6 +24,7 @@ def red_team( max_concurrent: int = 10, target_purpose: Optional[str] = None, attack_engine: Optional[AttackEngine] = None, + on_progress: Optional[ProgressCallback] = None, ): red_teamer = RedTeamer( async_mode=async_mode, @@ -31,15 +34,21 @@ def red_team( evaluation_model=evaluation_model, attack_engine=attack_engine, ) - risk_assessment = red_teamer.red_team( - model_callback=model_callback, - vulnerabilities=vulnerabilities, - attacks=attacks, - simulator_model=simulator_model, - evaluation_model=evaluation_model, - framework=framework, - attacks_per_vulnerability_type=attacks_per_vulnerability_type, - ignore_errors=ignore_errors, - attack_engine=attack_engine, + ctx = ( + progress_callback_context(on_progress) + if on_progress is not None + else nullcontext() ) + with ctx: + risk_assessment = red_teamer.red_team( + model_callback=model_callback, + vulnerabilities=vulnerabilities, + attacks=attacks, + simulator_model=simulator_model, + evaluation_model=evaluation_model, + framework=framework, + attacks_per_vulnerability_type=attacks_per_vulnerability_type, + ignore_errors=ignore_errors, + attack_engine=attack_engine, + ) return risk_assessment diff --git a/deepteam/utils.py b/deepteam/utils.py index 31f11438..eea8abf9 100644 --- a/deepteam/utils.py +++ b/deepteam/utils.py @@ -2,7 +2,7 @@ import inspect import os from contextlib import contextmanager -from typing import Optional, List, Callable +from typing import Optional, List, Callable, Dict, Any from rich.progress import ( Progress, SpinnerColumn, @@ -22,6 +22,66 @@ SPANS_CONTEXT_LIMIT = int(os.getenv("SPANS_CONTEXT_LIMIT", "40000")) +# Progress callback (issue #173): surface the progress deepteam shows on the +# console to a UI such as Streamlit. A single module-level observer fires from +# the `add_pbar` / `update_pbar` choke points, so every progress bar reports +# through it, with no callback threaded through the attack modules. +ProgressCallback = Callable[[Dict[str, Any]], None] +_progress_callback: Optional[ProgressCallback] = None + + +def set_progress_callback(callback: Optional[ProgressCallback]) -> None: + """Register (or clear) a callback that receives assessment progress events. + + The callback is invoked with a dict, for example:: + + {"event": "start", "description": "Simulating attacks", + "completed": 0, "total": 12} + {"event": "advance", "description": "Simulating attacks", + "completed": 1, "total": 12} + + This mirrors the progress deepteam renders to the console so it can be shown + in a custom UI such as Streamlit. Pass ``None`` to clear the callback. + """ + global _progress_callback + _progress_callback = callback + + +@contextmanager +def progress_callback_context(callback: Optional[ProgressCallback]): + """Temporarily register a progress callback, restoring the previous one.""" + global _progress_callback + previous = _progress_callback + _progress_callback = callback + try: + yield + finally: + _progress_callback = previous + + +def _emit_progress( + event: str, + description: str, + completed: Optional[int], + total: Optional[int], +) -> None: + callback = _progress_callback + if callback is None: + return + try: + callback( + { + "event": event, + "description": description, + "completed": completed, + "total": total, + } + ) + except Exception: + # A user-supplied callback must never break an assessment run. + pass + + def validate_model_callback_signature( model_callback: Callable, async_mode: bool, @@ -113,6 +173,7 @@ def add_pbar( total: Optional[int] = None, enabled: Optional[bool] = None, ) -> Optional[int]: + _emit_progress("start", description, 0, total) if progress is None or not hasattr(progress, "add_task"): return None return progress.add_task(description, total=total) @@ -135,8 +196,10 @@ def update_pbar( advance = task.remaining progress.update(pbar_id, advance=advance, total=total) task = next((t for t in progress.tasks if t.id == pbar_id), None) - if task is not None and task.finished and remove: - progress.remove_task(pbar_id) + if task is not None: + _emit_progress("advance", task.description, task.completed, task.total) + if task.finished and remove: + progress.remove_task(pbar_id) def remove_pbars( diff --git a/docs/content/docs/red-teaming-introduction.mdx b/docs/content/docs/red-teaming-introduction.mdx index 5212b824..a1e3418f 100644 --- a/docs/content/docs/red-teaming-introduction.mdx +++ b/docs/content/docs/red-teaming-introduction.mdx @@ -435,7 +435,7 @@ prompt_injection = PromptInjection() risk_assessment = red_team(model_callback=model_callback, vulnerabilities=[bias], attacks=[prompt_injection]) ``` -There is **ONE** mandatory and **ELEVEN** optional arguments when calling the `red_team()` function: +There is **ONE** mandatory and **TWELVE** optional arguments when calling the `red_team()` function: - `model_callback`: a callback of type `Callable[[str], str]` that wraps around the target LLM system you wish to red team. - [Optional] `vulnerabilities`: a list of type `BaseVulnerability`s that determines the weaknesses to detect for. @@ -449,6 +449,7 @@ There is **ONE** mandatory and **ELEVEN** optional arguments when calling the `r - [Optional] `max_concurrent`: an integer that determines the maximum number of coroutines that can be ran in parallel. You can decrease this value if your models are running into rate limit errors. Defaulted to `10`. - [Optional] `target_purpose`: a string specifying your target LLM application's intended purpose. This affects the passing and failing of simulated attacks that are evaluated. Defaulted to `None`. - [Optional] `attack_engine`: an optional `AttackEngine` used to refine baseline simulated attacks before they are sent to your target. Forwarded to the internal `RedTeamer`. Defaulted to `None` (a default engine is created per run). See [Attack engine](#attack-engine). +- [Optional] `on_progress`: a callback of type `Callable[[dict], None]` that receives progress events (e.g. `{"event": "advance", "description": ..., "completed": 1, "total": 12}`) as red teaming runs, useful for surfacing progress in a custom UI such as Streamlit. Defaulted to `None`. :::caution WARNING You **MUST** pass either `vulnerabilities` and `attacks` or `framework` in the `red_team` function. If not, you will get an error since red teaming cannot be performed without any of those. (In case you provide both, the `attacks` and `vulnerabilities` inside framework overwrite your regular `attacks` or `vulnerabilities`) diff --git a/tests/test_core/test_progress_callback.py b/tests/test_core/test_progress_callback.py new file mode 100644 index 00000000..851c6ec3 --- /dev/null +++ b/tests/test_core/test_progress_callback.py @@ -0,0 +1,76 @@ +import pytest + +from deepteam.utils import ( + create_progress, + add_pbar, + update_pbar, + set_progress_callback, + progress_callback_context, +) + + +class TestProgressCallback: + """Unit tests for the progress callback (issue #173). No model/API key.""" + + def setup_method(self): + set_progress_callback(None) + + def teardown_method(self): + set_progress_callback(None) + + def test_no_callback_is_noop(self): + # With no callback registered, progress helpers must still work. + with create_progress(enabled=True) as progress: + pbar_id = add_pbar(progress, "Task", total=2) + update_pbar(progress, pbar_id) + + def test_callback_receives_start_and_advance(self): + events = [] + set_progress_callback(events.append) + + with create_progress(enabled=True) as progress: + pbar_id = add_pbar(progress, "Simulating attacks", total=2) + update_pbar(progress, pbar_id) + update_pbar(progress, pbar_id) + + assert events[0]["event"] == "start" + assert events[0]["description"] == "Simulating attacks" + assert events[0]["completed"] == 0 + assert events[0]["total"] == 2 + + advances = [e for e in events if e["event"] == "advance"] + assert len(advances) == 2 + assert advances[-1]["completed"] == 2 + assert advances[-1]["total"] == 2 + + def test_context_manager_scopes_and_restores(self): + events = [] + with progress_callback_context(events.append): + add_pbar(None, "Within", total=1) + # Callback is restored (to None) after the context. + add_pbar(None, "Outside", total=1) + + descriptions = [e["description"] for e in events] + assert "Within" in descriptions + assert "Outside" not in descriptions + + def test_start_emitted_even_when_progress_disabled(self): + # A UI may want events without console rendering (progress disabled). + events = [] + set_progress_callback(events.append) + + pbar_id = add_pbar(None, "Headless", total=3) + + assert pbar_id is None + assert events and events[0]["description"] == "Headless" + + def test_callback_exception_does_not_break_run(self): + def boom(event): + raise RuntimeError("user callback error") + + set_progress_callback(boom) + + # A raising user callback must not propagate into the assessment. + with create_progress(enabled=True) as progress: + pbar_id = add_pbar(progress, "Task", total=1) + update_pbar(progress, pbar_id)