Add files via upload

Signed-off-by: Birdup <34012548+birdup000@users.noreply.github.com>

agixt/Chain.py

@@ -523,20 +523,29 @@ def move_step(self, chain_name, current_step_number, new_step_number):
     def get_step_response(self, chain_name, chain_run_id=None, step_number="all"):
         if chain_run_id is None:
             chain_run_id = self.get_last_chain_run_id(chain_name=chain_name)
+
         chain_data = self.get_chain(chain_name=chain_name)
-
-        if isinstance(chain_data, list) and not chain_data:  # Check if chain not found
+
+        # Handle empty chain data
+        if not chain_data:
             return None
-
+
+        # Get chain ID safely
+        chain_id = chain_data.get("id") if isinstance(chain_data, dict) else None
+
+        if not chain_id:
+            return None
+
         session = get_session()
+
         if step_number == "all":
             chain_steps = (
                 session.query(ChainStep)
-                .filter(ChainStep.chain_id == chain_data["id"])
+                .filter(ChainStep.chain_id == chain_id)
                 .order_by(ChainStep.step_number)
                 .all()
             )
-
+            
             responses = {}
             for step in chain_steps:
                 chain_step_responses = (
@@ -548,21 +557,24 @@ def get_step_response(self, chain_name, chain_run_id=None, step_number="all"):
                     .order_by(ChainStepResponse.timestamp)
                     .all()
                 )
+
                 step_responses = [response.content for response in chain_step_responses]
                 responses[str(step.step_number)] = step_responses
+
             session.close()
             return responses
+
         else:
             step_number = int(step_number)
             chain_step = (
                 session.query(ChainStep)
                 .filter(
-                    ChainStep.chain_id == chain_data["id"],
+                    ChainStep.chain_id == chain_id,
                     ChainStep.step_number == step_number,
                 )
                 .first()
             )
-
+            
             if chain_step:
                 chain_step_responses = (
                     session.query(ChainStepResponse)
@@ -573,9 +585,11 @@ def get_step_response(self, chain_name, chain_run_id=None, step_number="all"):
                     .order_by(ChainStepResponse.timestamp)
                     .all()
                 )
+
                 step_responses = [response.content for response in chain_step_responses]
                 session.close()
                 return step_responses
+
             else:
                 session.close()
                 return None