merge

euclid/database/database.py

@@ -3,6 +3,7 @@ from peewee import *
 from termcolor import colored
 from functools import reduce
 import operator
+from const.common import PROMPT_DATA_TO_IGNORE
 
 from utils.utils import hash_data
 from database.models.components.base_models import database
@@ -19,6 +20,7 @@ from database.models.development import Development
 from database.models.file_snapshot import FileSnapshot
 from database.models.command_runs import CommandRuns
 from database.models.user_inputs import UserInputs
+from database.models.files import File
 
 
 def save_user(user_id, email, password):
@@ -152,20 +154,20 @@ def get_progress_steps(app_id, step=None):
         return steps
 
 
-def save_development_step(app_id, prompt_path, prompt_data, llm_req_num, messages, response):
+def save_development_step(app_id, prompt_path, prompt_data, llm_req_num, messages, response, previous_step=None):
     app = get_app(app_id)
     hash_id = hash_data({
         'prompt_path': prompt_path,
-        'prompt_data': {k: v for k, v in (prompt_data.items() if prompt_data is not None else {}) if k not in {"directory_tree"}},
+        'prompt_data': {k: v for k, v in (prompt_data.items() if prompt_data is not None else {}) if k not in PROMPT_DATA_TO_IGNORE},
         'llm_req_num': llm_req_num
     })
     try:
         inserted_id = (DevelopmentSteps
-                       .insert(app=app, hash_id=hash_id, messages=messages, llm_response=response)
-                       .on_conflict(conflict_target=[DevelopmentSteps.app, DevelopmentSteps.hash_id],
-                                    preserve=[DevelopmentSteps.messages, DevelopmentSteps.llm_response],
-                                    update={})
-                       .execute())
+            .insert(app=app, hash_id=hash_id, messages=messages, llm_response=response, previous_dev_step=previous_step)
+            .on_conflict(conflict_target=[DevelopmentSteps.app, DevelopmentSteps.hash_id],
+                        preserve=[DevelopmentSteps.messages, DevelopmentSteps.llm_response],
+                        update={})
+            .execute())
 
         dev_step = DevelopmentSteps.get_by_id(inserted_id)
         print(colored(f"Saved DEV step => {dev_step.id}", "yellow"))
@@ -199,8 +201,8 @@ def hash_and_save_step(Model, app_id, hash_data_args, data_fields, message):
         inserted_id = (Model
                        .insert(**data_to_insert)
                        .on_conflict(conflict_target=[Model.app, Model.hash_id],
-                                    preserve=[field for field in data_fields.keys()],
-                                    update={})
+                                    preserve=[],
+                                    update=data_fields)
                        .execute())
 
         record = Model.get_by_id(inserted_id)
@@ -219,8 +221,11 @@ def save_command_run(project, command, cli_response):
     data_fields = {
         'command': command,
         'cli_response': cli_response,
+        'previous_command_run': project.checkpoints['last_command_run'],
     }
-    return hash_and_save_step(CommandRuns, project.args['app_id'], hash_data_args, data_fields, "Saved Command Run")
+    command_run = hash_and_save_step(CommandRuns, project.args['app_id'], hash_data_args, data_fields, "Saved Command Run")
+    project.checkpoints['last_command_run'] = command_run
+    return command_run
 
 
 def get_command_run_from_hash_id(project, command):
@@ -228,7 +233,8 @@ def get_command_run_from_hash_id(project, command):
         'command': command,
         'command_runs_count': project.command_runs_count
     }
-    return get_db_model_from_hash_id(data_to_hash, CommandRuns, project.args['app_id'])
+    command_run = get_db_model_from_hash_id(data_to_hash, CommandRuns, project.args['app_id'])
+    return command_run
 
 def save_user_input(project, query, user_input):
     hash_data_args = {
@@ -238,22 +244,26 @@ def save_user_input(project, query, user_input):
     data_fields = {
         'query': query,
         'user_input': user_input,
+        'previous_user_input': project.checkpoints['last_user_input'],
     }
-    return hash_and_save_step(UserInputs, project.args['app_id'], hash_data_args, data_fields, "Saved User Input")
+    user_input = hash_and_save_step(UserInputs, project.args['app_id'], hash_data_args, data_fields, "Saved User Input")
+    project.checkpoints['last_user_input'] = user_input
+    return user_input
 
 def get_user_input_from_hash_id(project, query):
     data_to_hash = {
         'query': query,
         'user_inputs_count': project.user_inputs_count
     }
-    return get_db_model_from_hash_id(data_to_hash, UserInputs, project.args['app_id'])
+    user_input = get_db_model_from_hash_id(data_to_hash, UserInputs, project.args['app_id'])
+    return user_input
 
 
 def get_development_step_from_hash_id(app_id, prompt_path, prompt_data, llm_req_num):
     if prompt_data is None:
         prompt_data_dict = {}
     else:
-        prompt_data_dict = {k: v for k, v in prompt_data.items() if k not in {"directory_tree"}}
+        prompt_data_dict = {k: v for k, v in prompt_data.items() if k not in PROMPT_DATA_TO_IGNORE}
 
     hash_id = hash_data({
         'prompt_path': prompt_path,
@@ -285,6 +295,7 @@ def create_tables():
             FileSnapshot,
             CommandRuns,
             UserInputs,
+            File,
         ])
 
 
@@ -304,6 +315,7 @@ def drop_tables():
             FileSnapshot,
             CommandRuns,
             UserInputs,
+            File,
             ]:
             database.execute_sql(f'DROP TABLE IF EXISTS "{table._meta.table_name}" CASCADE')