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

1"""Parser Registry for managing output parsers.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any 

6 

7from lexigram.logging import ( 

8 get_logger, 

9) 

10 

11logger = get_logger(__name__) 

12 

13 

14class ParserRegistry: 

15 """Registry for managing output parsers by name. 

16 

17 Provides a central registry for looking up parsers by name, 

18 similar to LangChain's parser registry. 

19 

20 Example: 

21 >>> registry = ParserRegistry() 

22 >>> registry.register("json", JSONOutputParser()) 

23 >>> parser = registry.get("json") 

24 >>> assert parser is not None 

25 """ 

26 

27 def __init__(self) -> None: 

28 """Initialize an empty registry.""" 

29 self._parsers: dict[str, Any] = {} 

30 

31 def register(self, name: str, parser: Any) -> None: 

32 """Register a parser with a name. 

33 

34 Args: 

35 name: Unique name for the parser. 

36 parser: Parser instance to register. 

37 """ 

38 self._parsers[name] = parser 

39 logger.debug("parser_registered", name=name) 

40 

41 def get(self, name: str) -> Any: 

42 """Get a parser by name. 

43 

44 Args: 

45 name: Name of the parser to retrieve. 

46 

47 Returns: 

48 The registered parser. 

49 

50 Raises: 

51 KeyError: If no parser is registered with that name. 

52 """ 

53 if name not in self._parsers: 

54 available = list(self._parsers.keys()) 

55 raise KeyError( 

56 f"No parser registered with name {name!r}. Available: {available}" 

57 ) 

58 return self._parsers[name] 

59 

60 def get_or_none(self, name: str) -> Any | None: 

61 """Get a parser by name, returning None if not found. 

62 

63 Args: 

64 name: Name of the parser to retrieve. 

65 

66 Returns: 

67 The registered parser, or None if not found. 

68 """ 

69 return self._parsers.get(name) 

70 

71 def list_parsers(self) -> list[str]: 

72 """List all registered parser names. 

73 

74 Returns: 

75 List of registered parser names. 

76 """ 

77 return list(self._parsers.keys()) 

78 

79 def unregister(self, name: str) -> None: 

80 """Unregister a parser by name. 

81 

82 Args: 

83 name: Name of the parser to unregister. 

84 

85 Raises: 

86 KeyError: If no parser is registered with that name. 

87 """ 

88 if name not in self._parsers: 

89 raise KeyError(f"No parser registered with name {name!r}") 

90 del self._parsers[name] 

91 logger.debug("parser_unregistered", name=name) 

92 

93 @classmethod 

94 def with_defaults(cls) -> ParserRegistry: 

95 """Create a registry with default parsers pre-registered. 

96 

97 Returns: 

98 A new ParserRegistry with default parsers. 

99 """ 

100 from lexigram.ai.llm.parsers.csv import CSVOutputParser 

101 from lexigram.ai.llm.parsers.enum import EnumOutputParser 

102 from lexigram.ai.llm.parsers.json import JSONOutputParser 

103 from lexigram.ai.llm.parsers.pydantic import PydanticOutputParser 

104 

105 registry = cls() 

106 registry.register("json", JSONOutputParser()) 

107 registry.register("pydantic", PydanticOutputParser) 

108 registry.register("enum", EnumOutputParser) 

109 registry.register("csv", CSVOutputParser()) 

110 return registry 

111 

112 

113__all__ = ["ParserRegistry"]