feat: add progress callback to surface assessment progress

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

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

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(

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`)

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)