subhdotsol commited on
Commit
c651e61
·
1 Parent(s): f15e9b2

refactor(models): clean up imports and add module docstring

Browse files
Files changed (1) hide show
  1. models.py +13 -1
models.py CHANGED
@@ -1,5 +1,11 @@
1
- from enum import Enum
 
 
 
2
  from pydantic import BaseModel, Field, field_validator
 
 
 
3
 
4
  class StrategyType(str, Enum):
5
  ROLEPLAY = "roleplay"
@@ -10,6 +16,7 @@ class StrategyType(str, Enum):
10
  ENCODING = "encoding"
11
  MULTI_TURN = "multi_turn"
12
 
 
13
  class TargetCategory(str, Enum):
14
  PRIVACY = "privacy"
15
  MISINFORMATION = "misinformation"
@@ -17,6 +24,7 @@ class TargetCategory(str, Enum):
17
  MANIPULATION = "manipulation"
18
  ILLEGAL_ACTIVITY = "illegal_activity"
19
 
 
20
  class AttackAction(BaseModel):
21
  strategy_type: StrategyType
22
  target_category: TargetCategory
@@ -31,6 +39,7 @@ class AttackAction(BaseModel):
31
  v = v.replace(pattern, "[FILTERED]")
32
  return v.strip()
33
 
 
34
  class RedTeamObservation(BaseModel):
35
  defender_response: str
36
  defense_score: float = Field(ge=0.0, le=1.0)
@@ -41,6 +50,7 @@ class RedTeamObservation(BaseModel):
41
  feedback: str
42
  episode_id: str
43
 
 
44
  class EpisodeState(BaseModel):
45
  episode_id: str
46
  turn: int
@@ -48,10 +58,12 @@ class EpisodeState(BaseModel):
48
  attacks_so_far:int
49
  is_active: bool
50
 
 
51
  class StepResult(BaseModel):
52
  observation: RedTeamObservation
53
  reward: float
54
 
 
55
  class ResetResponse(BaseModel):
56
  observation: RedTeamObservation
57
  episode_id: str
 
1
+ """
2
+ Shared Pydantic models and Enums for the RedTeamOS environment.
3
+ These define the API contract between the server, environment, and clients.
4
+ """
5
  from pydantic import BaseModel, Field, field_validator
6
+ from typing import Optional
7
+ from enum import Enum
8
+
9
 
10
  class StrategyType(str, Enum):
11
  ROLEPLAY = "roleplay"
 
16
  ENCODING = "encoding"
17
  MULTI_TURN = "multi_turn"
18
 
19
+
20
  class TargetCategory(str, Enum):
21
  PRIVACY = "privacy"
22
  MISINFORMATION = "misinformation"
 
24
  MANIPULATION = "manipulation"
25
  ILLEGAL_ACTIVITY = "illegal_activity"
26
 
27
+
28
  class AttackAction(BaseModel):
29
  strategy_type: StrategyType
30
  target_category: TargetCategory
 
39
  v = v.replace(pattern, "[FILTERED]")
40
  return v.strip()
41
 
42
+
43
  class RedTeamObservation(BaseModel):
44
  defender_response: str
45
  defense_score: float = Field(ge=0.0, le=1.0)
 
50
  feedback: str
51
  episode_id: str
52
 
53
+
54
  class EpisodeState(BaseModel):
55
  episode_id: str
56
  turn: int
 
58
  attacks_so_far:int
59
  is_active: bool
60
 
61
+
62
  class StepResult(BaseModel):
63
  observation: RedTeamObservation
64
  reward: float
65
 
66
+
67
  class ResetResponse(BaseModel):
68
  observation: RedTeamObservation
69
  episode_id: str