Repository navigation
Expand file tree
/
Copy pathllm_string_synthesizer.py
More file actions
168 lines (139 loc) · 7.17 KB
/
Copy pathllm_string_synthesizer.py
File metadata and controls
168 lines (139 loc) · 7.17 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
"""
LLM-based Program Synthesis
This module uses Amazon Bedrock to generate string operation programs
from input-output examples through carefully crafted prompts.
"""
import json
import os
from typing import List, Tuple, Optional
import boto3
from strings import *
class LLMPromptAndResponseLogger:
"""
Logger for LLM prompt and response
"""
def __init__(self, file_path: str):
self.file_path = file_path
self.prompt = ""
self.response = ""
self.program = None
self.examples = []
self.error = None
def log_prompt(self, prompt: str, examples: List[Tuple[str, str]]):
self.prompt = str(prompt)
self.examples = str(examples)
def log_response(self, response: dict):
self.response = str(response)
def log_program(self, program: StringExpression):
self.program = str(program)
def log_error(self, error: Exception):
self.error = str(error)
def save(self):
# dump into jsonl file as a single line
with open(self.file_path, 'a') as f:
f.write(json.dumps({'prompt': self.prompt, 'response': self.response, 'examples': self.examples, 'program': self.program, 'error': self.error}) + '\n')
f.flush()
class LLMStringSynthesizer:
"""LLM-based synthesizer using Amazon Bedrock"""
def __init__(self, api_key: Optional[str] = None, logger: Optional[LLMPromptAndResponseLogger] = None):
"""
Initialize the LLM synthesizer
Args:
api_key: Bedrock API key. If None, will try to get from environment
"""
if api_key is None:
api_key = os.getenv('AWS_BEARER_TOKEN_BEDROCK')
if not api_key:
raise ValueError("Bedrock API key required. Set AWS_BEARER_TOKEN_BEDROCK environment variable or pass api_key parameter.")
os.environ['AWS_BEARER_TOKEN_BEDROCK'] = api_key
self.client = boto3.client('bedrock-runtime', region_name='us-east-1')
self.model_id = 'us.anthropic.claude-sonnet-4-6'
self.logger = logger
def synthesize(self, examples: List[Tuple[str, str]], max_iterations: int = 5) -> StringExpression:
"""
Synthesize a string expression from input-output examples using LLM
Args:
examples: List of (input, output) string pairs
max_iterations: Maximum number of growth iterations (unused in LLM approach)
Returns:
StringExpression that satisfies the examples
"""
if not examples:
raise ValueError("No examples provided")
# Create the prompt template with DSL description and examples
prompt = self.generate_prompt(examples)
if self.logger:
self.logger.log_prompt(prompt, examples)
try:
# Generate response from Amazon Bedrock
response = self.client.converse(
modelId=self.model_id,
messages=[
{'role': 'user', 'content': [{'text': prompt}]}
],
)
if self.logger:
self.logger.log_response(response)
program_text = ''.join(
block['text']
for block in response['output']['message']['content']
if 'text' in block
).strip()
# Extract and evaluate the program
program = self.extract_program(program_text)
if self.logger:
self.logger.log_program(program)
# Validate the program against examples
if self.validate_program(program, examples):
if self.logger:
self.logger.save()
return program
else:
raise ValueError(f"Generated program does not satisfy all examples: {program}")
except Exception as e:
if self.logger:
self.logger.log_error(e)
self.logger.save()
raise ValueError(f"Failed to synthesize program: {str(e)}")
def generate_prompt(self, examples: List[Tuple[str, str]]) -> str:
"""
Create a comprehensive prompt template for the LLM including DSL description and examples
"""
#####################################################################################################
# #
# Part 3 (a): Creating a Prompt Template (`generate_prompt`) #
# #
# TODO: Add your implementation here. #
# #
# NOTE: Below is a placeholder implementation that does not generate a prompt template. You need to #
# implement the actual prompt template generation logic here. #
# #
#####################################################################################################
return "Please write a program"
def extract_program(self, response_text: str) -> StringExpression:
"""
Extract the program from LLM response and evaluate it to get StringExpression object
"""
#####################################################################################################
# #
# Part 3 (b): Extract the StringExpression from the LLM response (`extract_program`) #
# #
# TODO: Add your implementation here. #
# #
# NOTE: Below is a placeholder implementation that does not extract the StringExpression from the #
# LLM response. You need to implement the actual extraction logic here. #
# #
#####################################################################################################
return StringLiteral("Dummy program")
def validate_program(self, program: StringExpression, examples: List[Tuple[str, str]]) -> bool:
"""
Validate that the generated program works correctly on all examples
"""
try:
for input_str, expected_output in examples:
actual_output = program.interpret(input_str)
if actual_output != expected_output:
return False
return True
except Exception:
return False