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

1"""Enum Output Parser.""" 

2 

3from __future__ import annotations 

4 

5from enum import Enum 

6from typing import TYPE_CHECKING 

7 

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

9from lexigram.ai.llm.structured.parser import extract_json_block 

10from lexigram.logging import ( 

11 get_logger, 

12) 

13 

14if TYPE_CHECKING: 

15 from typing import TypeVar 

16 

17 T = TypeVar("T", bound=Enum) 

18else: 

19 T = None 

20 

21logger = get_logger(__name__) 

22 

23 

24class EnumOutputParser: 

25 """Parse LLM responses into Enum members. 

26 

27 Extracts JSON from the response and maps it to an Enum member. 

28 Supports both string values and integer values. 

29 

30 Example: 

31 >>> from enum import Enum 

32 >>> 

33 >>> class Status(Enum): 

34 ... ACTIVE = "active" 

35 ... INACTIVE = "inactive" 

36 >>> 

37 >>> parser = EnumOutputParser(Status) 

38 >>> result = parser.parse('"active"') 

39 >>> assert result == Status.ACTIVE 

40 """ 

41 

42 def __init__(self, enum: type[Enum]) -> None: 

43 """Initialize with an Enum class. 

44 

45 Args: 

46 enum: Enum subclass to parse into. 

47 """ 

48 self._enum = enum 

49 

50 def parse(self, text: str) -> Enum: 

51 """Parse text into an Enum member. 

52 

53 Args: 

54 text: Raw LLM response text that may contain JSON with enum value. 

55 

56 Returns: 

57 Corresponding Enum member. 

58 

59 Raises: 

60 ParseError: When JSON cannot be extracted or enum value is invalid. 

61 """ 

62 try: 

63 parsed = extract_json_block(text) 

64 except ValueError as exc: 

65 raise ParseError(str(exc)) from exc 

66 

67 if isinstance(parsed, str): 

68 value = parsed 

69 elif isinstance(parsed, int): 

70 try: 

71 return self._enum(parsed) 

72 except ValueError: 

73 raise ParseError(f"Invalid enum value: {parsed}") 

74 else: 

75 raise ParseError( 

76 f"Expected string or int for enum, got {type(parsed).__name__}" 

77 ) 

78 

79 try: 

80 return self._enum(value) 

81 except ValueError: 

82 valid_values = [e.value for e in self._enum] 

83 raise ParseError( 

84 f"Invalid enum value {value!r}. Valid values: {valid_values}" 

85 ) 

86 

87 def get_format_instructions(self) -> str: 

88 """Return format instructions for the LLM. 

89 

90 Returns: 

91 Format instruction string telling the model to output a valid 

92 enum value. 

93 """ 

94 valid_values = [e.value for e in self._enum] 

95 values_str = ", ".join(repr(v) for v in valid_values) 

96 return ( 

97 f"Your response should be one of: {values_str}. " 

98 "Return just the value, not a JSON object. " 

99 "Do not include any text before or after the value." 

100 ) 

101 

102 

103__all__ = ["EnumOutputParser"]