Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/structured/utils.py: 23%

1"""Utility functions for structured LLM output handling.""" 

2 

3from __future__ import annotations 

4 

5from typing import TYPE_CHECKING, Any, TypeVar, cast 

6 

7if TYPE_CHECKING: 

8 from lexigram.contracts.ai import LLMClientProtocol 

9 from lexigram.contracts.core import JSON 

10 from lexigram.domain import DomainModel 

11 

12from lexigram.ai.llm.structured.parser import StructuredOutputParser 

13 

14T = TypeVar("T", bound="DomainModel") 

15 

16 

17def create_json_mode_messages( 

18 prompt: str, 

19 schema: type[DomainModel] | None = None, 

20 system_prompt: str | None = None, 

21) -> list[dict[str, str]]: 

22 """Create messages for JSON mode with optional schema. 

23 

24 Args: 

25 prompt: User prompt 

26 schema: Optional Pydantic model for schema 

27 system_prompt: Optional system prompt (default: JSON instruction) 

28 

29 Returns: 

30 Messages list for LLM 

31 

32 Example: 

33 >>> messages = create_json_mode_messages( 

34 ... "Extract person info", 

35 ... schema=Person 

36 ... ) 

37 """ 

38 # Default JSON system prompt 

39 if system_prompt is None: 

40 system_prompt = ( 

41 "You are a helpful assistant that responds in valid JSON format. " 

42 "Always return properly formatted JSON without any additional text or explanation." 

43 ) 

44 

45 messages = [{"role": "system", "content": system_prompt}] 

46 

47 # Add schema if provided 

48 if schema: 

49 parser = StructuredOutputParser(schema) 

50 schema_text = parser.get_schema_prompt() 

51 prompt = f"{schema_text}\n\n{prompt}" 

52 

53 messages.append({"role": "user", "content": prompt}) 

54 

55 return messages 

56 

57 

58async def complete_with_schema( 

59 client: LLMClientProtocol, 

60 prompt: str, 

61 schema: type[T], 

62 system_prompt: str | None = None, 

63 **kwargs: Any, 

64) -> T: 

65 """Complete with automatic schema parsing and validation. 

66 

67 Args: 

68 client: LLM client 

69 prompt: User prompt 

70 schema: Pydantic model for validation 

71 system_prompt: Optional system prompt 

72 **kwargs: Additional completion arguments 

73 

74 Returns: 

75 Validated schema instance 

76 

77 Example: 

78 >>> from lexigram.ai.llm import OpenAIClient 

79 >>> 

80 >>> client = OpenAIClient(api_key="sk-...") 

81 >>> person = await complete_with_schema( 

82 ... client, 

83 ... "Extract person from: John Doe, age 30", 

84 ... schema=Person 

85 ... ) 

86 """ 

87 messages = create_json_mode_messages(prompt, schema, system_prompt) 

88 

89 result = await client.complete(messages=messages, **kwargs) # type: ignore[arg-type] 

90 if result.is_err(): 

91 raise result.unwrap_err() 

92 completion = result.unwrap() 

93 

94 parser = StructuredOutputParser(schema) 

95 return cast("T", parser.parse(completion)) 

96 

97 

98async def complete_with_json( 

99 client: LLMClientProtocol, 

100 prompt: str, 

101 system_prompt: str | None = None, 

102 **kwargs: Any, 

103) -> JSON: 

104 """Complete and parse response as JSON. 

105 

106 Args: 

107 client: LLM client 

108 prompt: User prompt 

109 system_prompt: Optional system prompt 

110 **kwargs: Additional completion arguments 

111 

112 Returns: 

113 Parsed JSON 

114 

115 Example: 

116 >>> data = await complete_with_json( 

117 ... client, 

118 ... "Generate a config with 3 fields" 

119 ... ) 

120 """ 

121 from lexigram.ai.llm.structured.formatter import ResponseFormatter 

122 

123 messages = create_json_mode_messages(prompt, system_prompt=system_prompt) 

124 

125 result = await client.complete(messages=messages, **kwargs) # type: ignore[arg-type] 

126 if result.is_err(): 

127 raise result.unwrap_err() 

128 completion = result.unwrap() 

129 

130 return ResponseFormatter.to_json(completion) # type: ignore[arg-type]