| import hashlib |
| import json |
| import os |
| import traceback |
| from augmentoolkit.generation_functions.pipeline_step_class import PipelineStep |
|
|
| |
|
|
|
|
| class SingleGenerationStep(PipelineStep): |
| def __init__(self, **kwargs): |
| super().__init__(**kwargs) |
|
|
| def generate_hash(self, input_data): |
| """Generate a unique hash from the input data to use as filename""" |
| |
| input_str = json.dumps(input_data, sort_keys=True) |
| |
| return hashlib.md5(input_str.encode()).hexdigest() |
|
|
| def make_output_path(self, output_dir, input_hash): |
| """Returns full path to the output text file based on input hash""" |
| return os.path.join(output_dir, f"{self.output_file}_{input_hash}.txt") |
|
|
| async def read_previous_output(self, input_data, output_dir): |
| """Check if output for this input already exists""" |
| input_hash = self.generate_hash(input_data) |
| output_path = self.make_output_path(output_dir, input_hash) |
|
|
| if os.path.exists(output_path): |
| with open(output_path, "r", encoding="utf-8") as f: |
| return json.load(f) |
| return None |
|
|
| def save( |
| self, |
| result=None, |
| input_data=None, |
| output_dir=None, |
| full_output=None, |
| full_response=None, |
| full_input=None, |
| include_details=False, |
| completion_mode=False, |
| ): |
| """Save result to a text file named by input hash""" |
| input_hash = self.generate_hash(input_data) |
| output_path = self.make_output_path(output_dir, input_hash) |
|
|
| |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) |
|
|
| if self.log_full_outputs and full_output: |
| print("\n" + "=" * 50) |
| print(f"FULL OUTPUT FOR HASH {input_hash}:") |
| print("=" * 50) |
|
|
| |
| if isinstance(full_output, dict): |
| for key, value in full_output.items(): |
| print(f"\n{key.upper()}:") |
| print("-" * 50) |
| print(value) |
| else: |
| |
| print(full_output) |
|
|
| print("=" * 50 + "\n") |
|
|
| if include_details: |
| obj = { |
| self.result_key: result, |
| "full_output": full_output, |
| "full_response": full_response, |
| "full_input": full_input, |
| "completion_mode": completion_mode, |
| } |
| else: |
| obj = {self.result_key: result} |
|
|
| |
| |
| temp_path = output_path + ".tmp" |
| with open(temp_path, "w", encoding="utf-8") as f: |
| json.dump(obj, f) |
| os.replace(temp_path, output_path) |
|
|
| return obj |
|
|
| async def run( |
| self, |
| input_data=None, |
| engine_wrapper=None, |
| default_prompt_folder=None, |
| prompt_folder=None, |
| output_dir=None, |
| completion_mode=False, |
| use_stop=True, |
| include_details=False, |
| **kwargs, |
| ): |
| full_prompt_path = ( |
| self.prompt_path + ".yaml" |
| if not completion_mode |
| else self.prompt_path + ".txt" |
| ) |
|
|
| |
| previous_result = await self.read_previous_output(input_data, output_dir) |
| if previous_result is not None: |
| return previous_result |
|
|
| processed_data, additional_kwargs = self.process_input_data(input_data) |
|
|
| error_message = "" |
| complete = False |
| max_retries = self.max_retries |
| while not complete and max_retries > 0: |
| try: |
| result, full_output, full_response, full_input = ( |
| await self.generate_data( |
| processed_data, |
| engine_wrapper, |
| full_prompt_path, |
| prompt_folder, |
| default_prompt_folder, |
| completion_mode, |
| use_stop, |
| error_message=error_message, |
| **kwargs, |
| **additional_kwargs, |
| ) |
| ) |
|
|
| validation_result = self.validation_function(result, input_data) |
| if validation_result["result"]: |
| complete = True |
| else: |
| error_message = validation_result["message"] |
| except Exception as e: |
| print(e) |
| error_message = str(e) |
| traceback.print_exc() |
| max_retries -= 1 |
|
|
| if not complete: |
| return None |
|
|
| return self.save( |
| result=result, |
| input_data=input_data, |
| output_dir=output_dir, |
| full_output=full_output, |
| full_response=full_response, |
| full_input=full_input, |
| include_details=include_details, |
| completion_mode=completion_mode, |
| ) |
|
|
