Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/parsers/fixing.py: 28%

1"""Format Fixing Parser with retry budget.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from typing import Any, TypeVar 

7 

8from lexigram.ai.llm.structured.exceptions import ParseError 

9from lexigram.logging import ( 

10 get_logger, 

11) 

12 

13logger = get_logger(__name__) 

14 

15T = TypeVar("T") 

16 

17 

18class FormatFixingParser: 

19 """Parser that retries with LLM-assisted fixing on parse failure. 

20 

21 Wraps a base parser and, on parse failure, calls the LLM with a fixing 

22 prompt that includes the original output and the parse error. Retries 

23 are bounded by the retry_budget. 

24 

25 Example: 

26 >>> parser = FormatFixingParser( 

27 ... base_parser=JSONOutputParser(), 

28 ... llm_client=llm_client, 

29 ... retry_budget=3 

30 ... ) 

31 >>> result = parser.parse('not valid json') 

32 """ 

33 

34 def __init__( 

35 self, 

36 base_parser: Any, 

37 llm_client: Any, 

38 *, 

39 retry_budget: int = 3, 

40 guard_check: Callable[[str], bool] | None = None, 

41 ) -> None: 

42 """Initialize the format fixing parser. 

43 

44 Args: 

45 base_parser: The underlying parser to use for parsing. 

46 llm_client: LLM client to use for fixing attempts. 

47 retry_budget: Maximum number of fix attempts (default 3). 

48 guard_check: Optional guard function to validate malformed input 

49 before sending to LLM. Should return True if safe. 

50 """ 

51 self._base_parser = base_parser 

52 self._llm_client = llm_client 

53 self._retry_budget = retry_budget 

54 self._guard_check = guard_check 

55 

56 async def parse(self, text: str) -> Any: 

57 """Parse text, attempting fixes on failure. 

58 

59 Args: 

60 text: Raw LLM response text to parse. 

61 

62 Returns: 

63 Parsed output from the base parser. 

64 

65 Raises: 

66 ParseError: When all fix attempts fail or guard check fails. 

67 """ 

68 attempt = 0 

69 last_error: ParseError | None = None 

70 current_text = text 

71 

72 while attempt <= self._retry_budget: 

73 try: 

74 return self._base_parser.parse(current_text) 

75 except ParseError as exc: 

76 last_error = exc 

77 attempt += 1 

78 

79 if attempt > self._retry_budget: 

80 logger.warning( 

81 "format_fixing_exhausted", 

82 attempts=attempt, 

83 error=str(exc), 

84 ) 

85 break 

86 

87 logger.debug( 

88 "format_fixing_attempt", 

89 attempt=attempt, 

90 error=str(exc), 

91 ) 

92 

93 current_text = await self._fix_with_llm(current_text, str(exc)) 

94 

95 if last_error is not None: 

96 raise last_error 

97 raise ParseError("Format fixing failed with no error recorded") 

98 

99 async def _fix_with_llm(self, original_output: str, error_message: str) -> str: 

100 """Call LLM to fix the malformed output. 

101 

102 Args: 

103 original_output: The original malformed output. 

104 error_message: The parse error message. 

105 

106 Returns: 

107 Fixed output from the LLM. 

108 

109 Raises: 

110 ParseError: When the LLM call fails. 

111 """ 

112 if self._guard_check is not None: 

113 if not self._guard_check(original_output): 

114 raise ParseError( 

115 "Guard check failed: input flagged as potentially unsafe" 

116 ) 

117 

118 fixing_prompt = self._build_fixing_prompt(original_output, error_message) 

119 

120 try: 

121 response = await self._llm_client.complete( 

122 messages=[{"role": "user", "content": fixing_prompt}] 

123 ) 

124 fixed = response.content if hasattr(response, "content") else str(response) 

125 logger.debug("format_fixing_llm_response", length=len(fixed)) 

126 return fixed 

127 except Exception as exc: 

128 raise ParseError(f"LLM fix call failed: {exc}") from exc 

129 

130 def _build_fixing_prompt(self, original_output: str, error_message: str) -> str: 

131 """Build the prompt for the LLM to fix the output. 

132 

133 Args: 

134 original_output: The original malformed output. 

135 error_message: The parse error message. 

136 

137 Returns: 

138 The complete prompt to send to the LLM. 

139 """ 

140 return f"""The following output failed to parse: 

141 

142Original output: 

143``` 

144{original_output} 

145``` 

146 

147Parse error: 

148{error_message} 

149 

150Please fix the output so it can be parsed correctly. Return only the fixed output, no explanations.""" 

151 

152 def get_format_instructions(self) -> str: 

153 """Return format instructions from the base parser. 

154 

155 Returns: 

156 Format instructions from the wrapped parser. 

157 """ 

158 if hasattr(self._base_parser, "get_format_instructions"): 

159 return str(self._base_parser.get_format_instructions()) 

160 return "" 

161 

162 

163__all__ = ["FormatFixingParser"]