Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/thinking/normalizer.py: 7%

1"""Extraction and normalization of thinking blocks from LLM responses.""" 

2 

3from __future__ import annotations 

4 

5from lexigram.ai.llm.thinking.patterns import THINKING_PATTERNS 

6 

7 

8def _extract_gemma_channel(text: str) -> tuple[str, str | None] | None: 

9 """Extract thinking from Gemma-4 channel format. 

10 

11 Google's canonical format (from official Gemma-4 docs): 

12 ``<|channel>thought\\n...thinking...\\n<channel|>\\n...response...`` 

13 

14 The end token ``<channel|>`` (note: opposite direction from the start token) 

15 marks the boundary between thinking and response content. 

16 

17 Gemma-4 26B/31B with thinking *disabled* still emits a ghost empty block: 

18 ``<|channel>thought\\n<channel|>\\n{json}`` — this is handled correctly since 

19 the thinking text between the markers is empty (returned as ``None``). 

20 

21 Fallback paths for LM Studio GGUF variants that may use different separators: 

22 1. ``\\n<|channel>response\\n`` — seen in some GGUF chat templates 

23 2. ``\\n<|channel>output\\n`` — alternative GGUF chat template variant 

24 3. Bare format: no end marker, content starts directly at the first ``{``/``[`` 

25 

26 Args: 

27 text: Raw model output potentially containing a Gemma channel block. 

28 

29 Returns: 

30 ``(clean_content, thinking_text)`` tuple, or ``None`` if no Gemma 

31 channel start marker (``<|channel>thought``) is found. 

32 """ 

33 start_marker = "<|channel>thought" 

34 

35 if start_marker not in text: 

36 return None 

37 

38 start_pos = text.find(start_marker) 

39 thinking_start = start_pos + len(start_marker) 

40 after_start = text[thinking_start:] 

41 

42 # --- Primary path: canonical <channel|> end token (Google official format) --- 

43 end_token = "<channel|>" # noqa: S105 # marker string, not a credential 

44 end_pos = after_start.find(end_token) 

45 if end_pos != -1: 

46 thinking_text = after_start[:end_pos].strip() 

47 clean_content = after_start[end_pos + len(end_token) :].strip() 

48 return clean_content, thinking_text or None 

49 

50 # --- Secondary paths: GGUF chat-template variants --- 

51 for sep in ("\n<|channel>response\n", "\n<|channel>output\n"): 

52 sep_pos = after_start.find(sep) 

53 if sep_pos != -1: 

54 thinking_text = after_start[:sep_pos].strip() 

55 clean_content = after_start[sep_pos + len(sep) :].strip() 

56 return clean_content, thinking_text or None 

57 

58 # --- Tertiary fallback: no separator — look for first { or [ --- 

59 # Some LM Studio GGUF variants output thinking directly followed by JSON 

60 # with no channel separator between them. 

61 for json_char in ("{", "["): 

62 json_pos = after_start.find(json_char) 

63 if json_pos != -1: 

64 thinking_text = after_start[:json_pos].strip() 

65 clean_content = after_start[json_pos:].strip() 

66 if thinking_text: # Only strip if there actually is thinking text 

67 return clean_content, thinking_text 

68 

69 return None 

70 

71 

72def normalize_thinking_text(text: str) -> tuple[str, str | None]: 

73 """Extract thinking text from raw LLM output. 

74 

75 Tries each pattern in THINKING_PATTERNS order. Returns (clean_content, thinking_text_or_None). 

76 clean_content has thinking block removed and is stripped. 

77 thinking_text is the raw thinking content (stripped), or None if not found. 

78 

79 Pattern matching is by substring presence of start_marker (and end_marker after it), 

80 NOT by model name. The bare-closing-tag pattern (end_marker="</think>", no start) 

81 matches only when start_marker is NOT found but end_marker IS found — this covers 

82 models that output ...thinking...</think>\\nresponse. 

83 

84 Falls back: after removing a thinking block, if clean_content is empty but thinking 

85 text was found, tries to extract from the first `{` or `[` in the original text to 

86 recover any JSON that may have been embedded. 

87 

88 Args: 

89 text: Raw LLM response text, possibly containing inline thinking tags. 

90 

91 Returns: 

92 A tuple of (clean_content, thinking_text_or_None). 

93 - clean_content: The response text with thinking stripped out, stripped of whitespace. 

94 - thinking_text_or_None: The thinking/reasoning text, or None if no thinking found. 

95 """ 

96 for pattern in THINKING_PATTERNS: 

97 # Special handling for Gemma channel format 

98 if pattern.name == "gemma_channel": 

99 result = _extract_gemma_channel(text) 

100 if result is not None: 

101 clean_content, thinking_text = result 

102 # Fallback: if clean_content is empty but thinking exists, try JSON recovery 

103 if not clean_content and thinking_text: 

104 # Try to find JSON starting with { or [ 

105 for json_char in ("{", "["): 

106 json_pos = text.find(json_char) 

107 if json_pos != -1: 

108 clean_content = text[json_pos:].strip() 

109 break 

110 return clean_content, thinking_text 

111 

112 # For bare closing tag pattern, only match if start_marker not found 

113 elif pattern.start_marker == "": 

114 if pattern.end_marker in text: 

115 # Find the position of the end marker 

116 end_pos = text.find(pattern.end_marker) 

117 thinking_text = text[:end_pos].strip() 

118 clean_content = text[end_pos + len(pattern.end_marker) :].strip() 

119 

120 # Fallback: if clean_content is empty but thinking exists, try JSON recovery 

121 if not clean_content and thinking_text: 

122 # Try to find JSON starting with { or [ 

123 for json_char in ("{", "["): 

124 json_pos = text.find(json_char) 

125 if json_pos != -1: 

126 clean_content = text[json_pos:].strip() 

127 break 

128 

129 return clean_content, thinking_text or None 

130 

131 # For patterns with start markers, check both start and end are present 

132 elif pattern.start_marker in text: 

133 start_pos = text.find(pattern.start_marker) 

134 # Look for end marker after the start marker 

135 end_search_start = start_pos + len(pattern.start_marker) 

136 end_pos = text.find(pattern.end_marker, end_search_start) 

137 

138 if end_pos != -1: 

139 # Extract thinking content (between markers) 

140 thinking_start = start_pos + len(pattern.start_marker) 

141 thinking_text = text[thinking_start:end_pos].strip() 

142 

143 # Extract clean content (after end marker) 

144 if pattern.strip_end_marker: 

145 clean_start = end_pos + len(pattern.end_marker) 

146 else: 

147 clean_start = end_pos 

148 

149 clean_content = text[clean_start:].strip() 

150 

151 # Fallback: if clean_content is empty but thinking exists, try JSON recovery 

152 if not clean_content and thinking_text: 

153 # Try to find JSON starting with { or [ 

154 after_thinking = text[clean_start:] 

155 for json_char in ("{", "["): 

156 json_pos = after_thinking.find(json_char) 

157 if json_pos != -1: 

158 clean_content = after_thinking[json_pos:].strip() 

159 break 

160 

161 return clean_content, thinking_text or None 

162 

163 # No pattern matched 

164 return text, None 

165 

166 

167__all__ = ["normalize_thinking_text"]