Coverage for baibu/submissions/cleaners.py: 100%

1"""Cleaners: the pluggable step that checks and cleans submitted text. 

2 

3Every submission first goes through the rules in :mod:`.text`. The cleaner 

4named by ``SUBMISSION_CLEANER`` then runs on the result: 

5 

6- ``RuleCleaner`` (default): the rules only. No model, no network. 

7- ``LiteLLMCleaner``: asks a language model, through LiteLLM, to redact 

8 remaining personal data and rate toxicity and quality. The deployment 

9 chooses the model and endpoint (``SUBMISSION_CLEANER_MODEL`` and friends). 

10- ``MockCleaner``: deterministic stand-in for a model, for tests and demos. 

11 

12A cleaner returns a :class:`CleaningResult` or raises :class:`CleanerError`. 

13""" 

14 

15import json 

16import re 

17from dataclasses import dataclass 

18from dataclasses import field 

19from functools import cache 

20 

21from django.conf import settings 

22from django.template.loader import render_to_string 

23from django.utils.module_loading import import_string 

24 

25from . import text as text_rules 

26 

27 

28class CleanerError(Exception): 

29 """The cleaner could not produce a result (for example, the model failed).""" 

30 

31 

32@dataclass 

33class CleaningResult: 

34 text: str 

35 cleaner: str 

36 redaction_count: int = 0 

37 toxicity_detected: bool | None = None 

38 quality_score: int | None = None 

39 detected_language: str = "" 

40 flags: list[str] = field(default_factory=list) 

41 

42 

43class Cleaner: 

44 name = "cleaner" 

45 

46 def clean(self, text: str, *, language_code: str) -> CleaningResult: 

47 raise NotImplementedError 

48 

49 

50def apply_rules(raw_text: str) -> text_rules.Redacted: 

51 return text_rules.redact(text_rules.normalise(raw_text)) 

52 

53 

54class RuleCleaner(Cleaner): 

55 """Normalisation and pattern redaction only.""" 

56 

57 name = "rules" 

58 

59 def clean(self, text: str, *, language_code: str) -> CleaningResult: 

60 ruled = apply_rules(text) 

61 return CleaningResult(text=ruled.text, cleaner=self.name, redaction_count=ruled.total) 

62 

63 

64class MockCleaner(RuleCleaner): 

65 """Behaves like a model cleaner without calling one. 

66 

67 Text containing a word from ``toxic_words`` is flagged as toxic; the 

68 quality score grows with length; the language is echoed back. 

69 """ 

70 

71 name = "mock" 

72 toxic_words = frozenset({"mocktoxic"}) 

73 

74 def clean(self, text: str, *, language_code: str) -> CleaningResult: 

75 result = super().clean(text, language_code=language_code) 

76 words = {w.strip(".,!?;:").lower() for w in result.text.split()} 

77 result.toxicity_detected = bool(words & self.toxic_words) 

78 result.quality_score = min(100, 10 * text_rules.word_count(result.text)) 

79 result.detected_language = language_code 

80 result.cleaner = self.name 

81 return result 

82 

83 

84_JSON_BLOCK = re.compile(r"\{.*\}", re.DOTALL) 

85 

86 

87class LiteLLMCleaner(RuleCleaner): 

88 """Rules first, then a language model reached through LiteLLM. 

89 

90 The instructions come from the template ``submissions/cleaner_prompt.txt``, 

91 which a deployment can override. The model must answer with one JSON 

92 object: ``text``, ``redactions``, ``toxic``, ``quality``, ``language``. 

93 """ 

94 

95 name = "litellm" 

96 

97 def clean(self, text: str, *, language_code: str) -> CleaningResult: 

98 import litellm 

99 

100 ruled = super().clean(text, language_code=language_code) 

101 model = settings.SUBMISSION_CLEANER_MODEL 

102 if not model: 

103 msg = "SUBMISSION_CLEANER_MODEL is not set." 

104 raise CleanerError(msg) 

105 prompt = render_to_string("submissions/cleaner_prompt.txt", {"language_code": language_code}) 

106 try: 

107 response = litellm.completion( 

108 model=model, 

109 messages=[{"role": "system", "content": prompt}, {"role": "user", "content": ruled.text}], 

110 api_base=settings.SUBMISSION_CLEANER_API_BASE or None, 

111 api_key=settings.SUBMISSION_CLEANER_API_KEY or None, 

112 timeout=settings.SUBMISSION_CLEANER_TIMEOUT, 

113 ) 

114 content = response.choices[0].message.content or "" 

115 except Exception as exc: 

116 msg = f"Model call failed: {type(exc).__name__}" 

117 raise CleanerError(msg) from exc 

118 answer = _parse_answer(content) 

119 return CleaningResult( 

120 text=text_rules.normalise(answer["text"]), 

121 cleaner=f"{self.name}:{getattr(response, 'model', None) or model}"[:128], 

122 redaction_count=ruled.redaction_count + answer["redactions"], 

123 toxicity_detected=answer["toxic"], 

124 quality_score=answer["quality"], 

125 detected_language=answer["language"], 

126 ) 

127 

128 

129def _parse_answer(content: str) -> dict: 

130 match = _JSON_BLOCK.search(content) 

131 try: 

132 data = json.loads(match.group(0) if match else content) 

133 except json.JSONDecodeError as exc: 

134 msg = "Model answer is not JSON." 

135 raise CleanerError(msg) from exc 

136 if not isinstance(data, dict) or not isinstance(data.get("text"), str): 

137 msg = "Model answer has no text." 

138 raise CleanerError(msg) 

139 quality = data.get("quality") 

140 redactions = data.get("redactions") 

141 return { 

142 "text": data["text"], 

143 "redactions": redactions if isinstance(redactions, int) and redactions >= 0 else 0, 

144 "toxic": data["toxic"] if isinstance(data.get("toxic"), bool) else None, 

145 "quality": max(0, min(100, quality)) if isinstance(quality, int) and not isinstance(quality, bool) else None, 

146 "language": str(data.get("language") or "")[:24], 

147 } 

148 

149 

150@cache 

151def _cleaner_class(path: str) -> type[Cleaner]: 

152 return import_string(path) 

153 

154 

155def get_cleaner() -> Cleaner: 

156 return _cleaner_class(settings.SUBMISSION_CLEANER)()