Coverage for llm_dataset_engine/stages/prompt_formatter_stage.py: 19%

1"""Prompt formatting stage for template-based prompt generation."""

2

3 from decimal import Decimal

4 from typing import Any, Dict, List

5

6 import pandas as pd

7 from jinja2 import Template as Jinja2Template

8

9 from llm_dataset_engine.core.models import (

10 CostEstimate,

11 PromptBatch,

12 RowMetadata,

13 ValidationResult,

14)

15 from llm_dataset_engine.core.specifications import PromptSpec

16 from llm_dataset_engine.stages.pipeline_stage import PipelineStage

17

18

19 class PromptFormatterStage(

20 PipelineStage[tuple[pd.DataFrame, PromptSpec], List[PromptBatch]]

21):

22 """

23 Format prompts using template and row data.

24

25 Responsibilities:

26 - Extract input columns from rows

27 - Format prompts using template

28 - Batch prompts for efficient processing

29 - Attach metadata for tracking

30 """

31

32 def __init__(

33 self, batch_size: int = 100, use_jinja2: bool = False

34 ):

35 """

36 Initialize prompt formatter stage.

37

38 Args:

39 batch_size: Number of prompts per batch

40 use_jinja2: Use Jinja2 for template rendering

41 """

42 super().__init__("PromptFormatter")

43 self.batch_size = batch_size

44 self.use_jinja2 = use_jinja2

45

46 def process(

47 self, input_data: tuple[pd.DataFrame, PromptSpec], context: Any

48 ) -> List[PromptBatch]:

49 """Format prompts from DataFrame rows."""

50 df, prompt_spec = input_data

51

52 prompts: List[str] = []

53 metadata_list: List[RowMetadata] = []

54

55 # Extract template variables

56 template_str = prompt_spec.template

57

58 # Create template renderer

59 if self.use_jinja2:

60 template = Jinja2Template(template_str)

61

62 # Format prompt for each row

63 for idx, row in df.iterrows():

64 try:

65 # Extract input columns

66 row_data = {

67 col: row[col]

68 for col in df.columns

69 if col in template_str

70 }

71

72 # Format prompt (Jinja2 or f-string)

73 if self.use_jinja2:

74 prompt = template.render(**row_data)

75 else:

76 prompt = template_str.format(**row_data)

77

78 # Add few-shot examples if specified

79 if prompt_spec.few_shot_examples:

80 examples_text = self._format_few_shot_examples(

81 prompt_spec.few_shot_examples

82 )

83 prompt = f"{examples_text}\n\n{prompt}"

84

85 # Add system message if specified

86 if prompt_spec.system_message:

87 prompt = f"{prompt_spec.system_message}\n\n{prompt}"

88

89 prompts.append(prompt)

90

91 # Create metadata

92 metadata = RowMetadata(

93 row_index=idx,

94 row_id=row.get("id", None),

95 )

96 metadata_list.append(metadata)

97

98 except KeyError as e:

99 self.logger.warning(

100 f"Missing template variable at row {idx}: {e}"

101 )

102 continue

103 except Exception as e:

104 self.logger.error(

105 f"Error formatting prompt at row {idx}: {e}"

106 )

107 continue

108

109 # Create batches

110 batches: List[PromptBatch] = []

111 for i in range(0, len(prompts), self.batch_size):

112 batch_prompts = prompts[i : i + self.batch_size]

113 batch_metadata = metadata_list[i : i + self.batch_size]

114

115 batch = PromptBatch(

116 prompts=batch_prompts,

117 metadata=batch_metadata,

118 batch_id=i // self.batch_size,

119 )

120 batches.append(batch)

121

122 self.logger.info(

123 f"Formatted {len(prompts)} prompts into {len(batches)} batches"

124 )

125

126 return batches

127

128 def _format_few_shot_examples(

129 self, examples: List[Dict[str, str]]

130 ) -> str:

131 """

132 Format few-shot examples for prompt.

133

134 Args:

135 examples: List of example dicts with 'input' and 'output'

136

137 Returns:

138 Formatted examples text

139 """

140 formatted = ["Here are some examples:\n"]

141

142 for i, example in enumerate(examples, 1):

143 formatted.append(f"Example {i}:")

144 formatted.append(f"Input: {example.get('input', '')}")

145 formatted.append(f"Output: {example.get('output', '')}")

146 formatted.append("")

147

148 return "\n".join(formatted)

149

150 def validate_input(

151 self, input_data: tuple[pd.DataFrame, PromptSpec]

152 ) -> ValidationResult:

153 """Validate DataFrame and prompt specification."""

154 result = ValidationResult(is_valid=True)

155

156 df, prompt_spec = input_data

157

158 # Check DataFrame not empty

159 if df.empty:

160 result.add_error("DataFrame is empty")

161

162 # Check template variables exist in DataFrame

163 template = prompt_spec.template

164 import re

165

166 variables = re.findall(r"\{(\w+)\}", template)

167 missing_vars = set(variables) - set(df.columns)

168

169 if missing_vars:

170 result.add_error(

171 f"Template variables not in DataFrame: {missing_vars}"

172 )

173

174 return result

175

176 def estimate_cost(

177 self, input_data: tuple[pd.DataFrame, PromptSpec]

178 ) -> CostEstimate:

179 """Prompt formatting has no LLM cost."""

180 return CostEstimate(

181 total_cost=Decimal("0.0"),

182 total_tokens=0,

183 input_tokens=0,

184 output_tokens=0,

185 rows=len(input_data[0]),

186 )

187