From 43d8c5511d811598f93bcbe8b2160fc543e837af Mon Sep 17 00:00:00 2001
From: Zvonimir Sabljic <zvonimir@pythagora.io>
Date: Tue, 12 Sep 2023 20:57:11 +0200
Subject: [PATCH] Added current high level step to saved command runs, user
 inputs and dev steps because we skip high level steps in one way and then in
 the coding step, we skip them with these functions. REFACTOR eventually so
 that we skip everything with saved steps since the beginning of the project.

---
 pilot/database/database.py                 | 19 ++++++++++++++-----
 pilot/database/models/command_runs.py      |  1 +
 pilot/database/models/development_steps.py |  1 +
 pilot/database/models/user_inputs.py       |  1 +
 4 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/pilot/database/database.py b/pilot/database/database.py
index 52d9711..bb6cdcc 100644
--- a/pilot/database/database.py
+++ b/pilot/database/database.py
@@ -157,9 +157,9 @@ def get_progress_steps(app_id, step=None):
         return steps
 
 
-def get_db_model_from_hash_id(model, app_id, previous_step):
+def get_db_model_from_hash_id(model, app_id, previous_step, high_level_step):
     try:
-        db_row = model.get((model.app == app_id) & (model.previous_step == previous_step))
+        db_row = model.get((model.app == app_id) & (model.previous_step == previous_step) & (model.high_level_step == high_level_step))
     except DoesNotExist:
         return None
     return db_row
@@ -207,6 +207,7 @@ def save_development_step(project, prompt_path, prompt_data, messages, llm_respo
         'messages': messages,
         'llm_response': llm_response,
         'previous_step': project.checkpoints['last_development_step'],
+        'high_level_step': project.current_step,
     }
 
     development_step = hash_and_save_step(DevelopmentSteps, project.args['app_id'], hash_data_args, data_fields, "Saved Development Step")
@@ -227,11 +228,14 @@ def get_development_step_from_hash_id(project, prompt_path, prompt_data, llm_req
         'llm_req_num': llm_req_num
     }
     development_step = get_db_model_from_hash_id(DevelopmentSteps, project.args['app_id'],
-                                                 project.checkpoints['last_development_step'])
+                                                 project.checkpoints['last_development_step'], project.current_step)
     return development_step
 
 
 def save_command_run(project, command, cli_response):
+    if project.current_step != 'coding':
+        return
+
     hash_data_args = {
         'command': command,
         'command_runs_count': project.command_runs_count,
@@ -240,6 +244,7 @@ def save_command_run(project, command, cli_response):
         'command': command,
         'cli_response': cli_response,
         'previous_step': project.checkpoints['last_command_run'],
+        'high_level_step': project.current_step,
     }
     command_run = hash_and_save_step(CommandRuns, project.args['app_id'], hash_data_args, data_fields,
                                      "Saved Command Run")
@@ -253,11 +258,14 @@ def get_command_run_from_hash_id(project, command):
         'command_runs_count': project.command_runs_count
     }
     command_run = get_db_model_from_hash_id(CommandRuns, project.args['app_id'],
-                                            project.checkpoints['last_command_run'])
+                                            project.checkpoints['last_command_run'], project.current_step)
     return command_run
 
 
 def save_user_input(project, query, user_input):
+    if project.current_step != 'coding':
+        return
+
     hash_data_args = {
         'query': query,
         'user_inputs_count': project.user_inputs_count,
@@ -266,6 +274,7 @@ def save_user_input(project, query, user_input):
         'query': query,
         'user_input': user_input,
         'previous_step': project.checkpoints['last_user_input'],
+        'high_level_step': project.current_step,
     }
     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
@@ -277,7 +286,7 @@ def get_user_input_from_hash_id(project, query):
         'query': query,
         'user_inputs_count': project.user_inputs_count
     }
-    user_input = get_db_model_from_hash_id(UserInputs, project.args['app_id'], project.checkpoints['last_user_input'])
+    user_input = get_db_model_from_hash_id(UserInputs, project.args['app_id'], project.checkpoints['last_user_input'], project.current_step)
     return user_input
 
 
diff --git a/pilot/database/models/command_runs.py b/pilot/database/models/command_runs.py
index b6c34c4..bcce4c9 100644
--- a/pilot/database/models/command_runs.py
+++ b/pilot/database/models/command_runs.py
@@ -11,6 +11,7 @@ class CommandRuns(BaseModel):
     command = TextField(null=True)
     cli_response = TextField(null=True)
     previous_step = ForeignKeyField('self', null=True, column_name='previous_step')
+    high_level_step = CharField(null=True)
 
     class Meta:
         db_table = 'command_runs'
diff --git a/pilot/database/models/development_steps.py b/pilot/database/models/development_steps.py
index 6492a4d..83988a5 100644
--- a/pilot/database/models/development_steps.py
+++ b/pilot/database/models/development_steps.py
@@ -18,6 +18,7 @@ class DevelopmentSteps(BaseModel):
         llm_response = JSONField(null=False)  # Custom JSON field for SQLite
 
     previous_step = ForeignKeyField('self', null=True, column_name='previous_step')
+    high_level_step = CharField(null=True)
 
     class Meta:
         db_table = 'development_steps'
diff --git a/pilot/database/models/user_inputs.py b/pilot/database/models/user_inputs.py
index 7d2451c..b2539f2 100644
--- a/pilot/database/models/user_inputs.py
+++ b/pilot/database/models/user_inputs.py
@@ -11,6 +11,7 @@ class UserInputs(BaseModel):
     query = TextField(null=True)
     user_input = TextField(null=True)
     previous_step = ForeignKeyField('self', null=True, column_name='previous_step')
+    high_level_step = CharField(null=True)
 
     class Meta:
         db_table = 'user_inputs'