asofter commited on
Commit
7bcd3fa
1 Parent(s): 6c4c4e6

* update version

Browse files
Files changed (4) hide show
  1. app.py +6 -2
  2. output.py +15 -2
  3. prompt.py +9 -2
  4. requirements.txt +1 -1
app.py CHANGED
@@ -41,6 +41,10 @@ Scanning prompt and output using [LLM Guard](https://laiyer-ai.github.io/llm-gua
41
 
42
  scanner_type = st.sidebar.selectbox("Type", [PROMPT, OUTPUT], index=0)
43
 
 
 
 
 
44
  enabled_scanners = None
45
  settings = None
46
  if scanner_type == PROMPT:
@@ -105,11 +109,11 @@ try:
105
  start_time = time.monotonic()
106
  if scanner_type == PROMPT:
107
  st_result_text, results_valid, results_score = scan_prompt(
108
- vault, enabled_scanners, settings, st_prompt_text
109
  )
110
  elif scanner_type == OUTPUT:
111
  st_result_text, results_valid, results_score = scan_output(
112
- vault, enabled_scanners, settings, st_prompt_text, st_output_text
113
  )
114
  end_time = time.monotonic()
115
  st_time_delta = timedelta(seconds=end_time - start_time)
 
41
 
42
  scanner_type = st.sidebar.selectbox("Type", [PROMPT, OUTPUT], index=0)
43
 
44
+ st_fail_fast = st.sidebar.checkbox(
45
+ "Fail fast", value=False, help="Stop scanning after first failure"
46
+ )
47
+
48
  enabled_scanners = None
49
  settings = None
50
  if scanner_type == PROMPT:
 
109
  start_time = time.monotonic()
110
  if scanner_type == PROMPT:
111
  st_result_text, results_valid, results_score = scan_prompt(
112
+ vault, enabled_scanners, settings, st_prompt_text, st_fail_fast
113
  )
114
  elif scanner_type == OUTPUT:
115
  st_result_text, results_valid, results_score = scan_output(
116
+ vault, enabled_scanners, settings, st_prompt_text, st_output_text, st_fail_fast
117
  )
118
  end_time = time.monotonic()
119
  st_time_delta = timedelta(seconds=end_time - start_time)
output.py CHANGED
@@ -371,13 +371,22 @@ def get_scanner(scanner_name: str, vault: Vault, settings: Dict):
371
 
372
 
373
  def scan(
374
- vault: Vault, enabled_scanners: List[str], settings: Dict, prompt: str, text: str
 
 
 
 
 
375
  ) -> (str, Dict[str, bool], Dict[str, float]):
376
  sanitized_output = text
377
  results_valid = {}
378
  results_score = {}
379
 
380
- with st.status("Scanning output...", expanded=True) as status:
 
 
 
 
381
  for scanner_name in enabled_scanners:
382
  st.write(f"{scanner_name} scanner...")
383
  scanner = get_scanner(
@@ -386,6 +395,10 @@ def scan(
386
  sanitized_output, is_valid, risk_score = scanner.scan(prompt, sanitized_output)
387
  results_valid[scanner_name] = is_valid
388
  results_score[scanner_name] = risk_score
 
 
 
 
389
  status.update(label="Scanning complete", state="complete", expanded=False)
390
 
391
  return sanitized_output, results_valid, results_score
 
371
 
372
 
373
  def scan(
374
+ vault: Vault,
375
+ enabled_scanners: List[str],
376
+ settings: Dict,
377
+ prompt: str,
378
+ text: str,
379
+ fail_fast: bool = False,
380
  ) -> (str, Dict[str, bool], Dict[str, float]):
381
  sanitized_output = text
382
  results_valid = {}
383
  results_score = {}
384
 
385
+ status_text = "Scanning prompt..."
386
+ if fail_fast:
387
+ status_text = "Scanning prompt (fail fast mode)..."
388
+
389
+ with st.status(status_text, expanded=True) as status:
390
  for scanner_name in enabled_scanners:
391
  st.write(f"{scanner_name} scanner...")
392
  scanner = get_scanner(
 
395
  sanitized_output, is_valid, risk_score = scanner.scan(prompt, sanitized_output)
396
  results_valid[scanner_name] = is_valid
397
  results_score[scanner_name] = risk_score
398
+
399
+ if fail_fast and not is_valid:
400
+ break
401
+
402
  status.update(label="Scanning complete", state="complete", expanded=False)
403
 
404
  return sanitized_output, results_valid, results_score
prompt.py CHANGED
@@ -320,19 +320,26 @@ def get_scanner(scanner_name: str, vault: Vault, settings: Dict):
320
 
321
 
322
  def scan(
323
- vault: Vault, enabled_scanners: List[str], settings: Dict, text: str
324
  ) -> (str, Dict[str, bool], Dict[str, float]):
325
  sanitized_prompt = text
326
  results_valid = {}
327
  results_score = {}
328
 
329
- with st.status("Scanning prompt...", expanded=True) as status:
 
 
 
 
330
  for scanner_name in enabled_scanners:
331
  st.write(f"{scanner_name} scanner...")
332
  scanner = get_scanner(scanner_name, vault, settings[scanner_name])
333
  sanitized_prompt, is_valid, risk_score = scanner.scan(sanitized_prompt)
334
  results_valid[scanner_name] = is_valid
335
  results_score[scanner_name] = risk_score
 
 
 
336
  status.update(label="Scanning complete", state="complete", expanded=False)
337
 
338
  return sanitized_prompt, results_valid, results_score
 
320
 
321
 
322
  def scan(
323
+ vault: Vault, enabled_scanners: List[str], settings: Dict, text: str, fail_fast: bool = False
324
  ) -> (str, Dict[str, bool], Dict[str, float]):
325
  sanitized_prompt = text
326
  results_valid = {}
327
  results_score = {}
328
 
329
+ status_text = "Scanning prompt..."
330
+ if fail_fast:
331
+ status_text = "Scanning prompt (fail fast mode)..."
332
+
333
+ with st.status(status_text, expanded=True) as status:
334
  for scanner_name in enabled_scanners:
335
  st.write(f"{scanner_name} scanner...")
336
  scanner = get_scanner(scanner_name, vault, settings[scanner_name])
337
  sanitized_prompt, is_valid, risk_score = scanner.scan(sanitized_prompt)
338
  results_valid[scanner_name] = is_valid
339
  results_score[scanner_name] = risk_score
340
+
341
+ if fail_fast and not is_valid:
342
+ break
343
  status.update(label="Scanning complete", state="complete", expanded=False)
344
 
345
  return sanitized_prompt, results_valid, results_score
requirements.txt CHANGED
@@ -1,4 +1,4 @@
1
- llm-guard==0.2.2
2
  pandas==2.1.0
3
  streamlit==1.26.0
4
  streamlit-tags==1.2.8
 
1
+ llm-guard==0.2.3
2
  pandas==2.1.0
3
  streamlit==1.26.0
4
  streamlit-tags==1.2.8