File size: 1,369 Bytes
005d862
80b34d1
005d862
 
 
 
 
 
 
 
 
80b34d1
005d862
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from typing import List, Union, Literal, Optional
from pydantic import BaseModel, ConfigDict, Field
import json
import re

class SemanticFilter(BaseModel):
    filter_type: Literal["substring", "regex_pattern", "length_limit", "entropy_threshold", "keyword_match"]
    value: Union[str, int, float]

class LogicNode(BaseModel):
    operator: Literal["AND", "OR", "NOT"]
    children: List[Union["LogicNode", SemanticFilter]] = Field(..., min_length=1)

class GuardrailGraph(BaseModel):
    graph_id: str
    description: str
    root: LogicNode

LogicNode.model_rebuild()

class Observation(BaseModel):
    adversarial_samples: List[str]
    benign_samples: List[str]

class Action(BaseModel):
    ast_json: str  # The model outputs a JSON string representing the GuardrailGraph
    baseline_ast_json: Optional[str] = None

class StepResult(BaseModel):
    observation: Observation
    reward: float
    done: bool
    info: dict

def extract_and_clean_json(text: str) -> str:
    # Extract JSON blocks from markdown explicitly and strip trailing commas
    text = text.strip()
    match = re.search(r"```(?:json)?\s*(.*?)\s*```", text, re.DOTALL)
    if match:
        text = match.group(1)
        
    # Replace trailing commas before closing braces/brackets
    text = re.sub(r',\s*}', '}', text)
    text = re.sub(r',\s*]', ']', text)
    
    return text.strip()