From ed5952782fbdf388f335768d12a958088c611d68 Mon Sep 17 00:00:00 2001
From: slavakurilyak <hello@slavakurilyak.com>
Date: Tue, 4 Apr 2023 15:53:59 -0500
Subject: [PATCH] Improve key validation and handling in overwrite_memory and
 message_agent functions

---
 scripts/commands.py | 38 ++++++++++++++++++++++++++++++++------
 1 file changed, 32 insertions(+), 6 deletions(-)

diff --git a/scripts/commands.py b/scripts/commands.py
index f8f96fe7..c1d5db9d 100644
--- a/scripts/commands.py
+++ b/scripts/commands.py
@@ -16,6 +16,13 @@ from googleapiclient.errors import HttpError
 cfg = Config()
 
 
+def is_valid_int(value):
+    try:
+        int(value)
+        return True
+    except ValueError:
+        return False
+
 def get_command(response):
     try:
         response_json = fix_and_parse_json(response)
@@ -194,14 +201,28 @@ def delete_memory(key):
 
 
 def overwrite_memory(key, string):
-    if int(key) >= 0 and key < len(mem.permanent_memory):
-        _text = "Overwriting memory with key " + \
-            str(key) + " and string " + string
-        mem.permanent_memory[key] = string
+    # Check if the key is a valid integer
+    if is_valid_int(key):
+        key_int = int(key)
+        # Check if the integer key is within the range of the permanent_memory list
+        if 0 <= key_int < len(mem.permanent_memory):
+            _text = "Overwriting memory with key " + str(key) + " and string " + string
+            # Overwrite the memory slot with the given integer key and string
+            mem.permanent_memory[key_int] = string
+            print(_text)
+            return _text
+        else:
+            print(f"Invalid key '{key}', out of range.")
+            return None
+    # Check if the key is a valid string
+    elif isinstance(key, str):
+        _text = "Overwriting memory with key " + key + " and string " + string
+        # Overwrite the memory slot with the given string key and string
+        mem.string_key_memory[key] = string
         print(_text)
         return _text
     else:
-        print("Invalid key, cannot overwrite memory.")
+        print(f"Invalid key '{key}', must be an integer or a string.")
         return None
 
 
@@ -235,7 +256,12 @@ def start_agent(name, task, prompt, model=cfg.fast_llm_model):
 
 def message_agent(key, message):
     global cfg
-    agent_response = agents.message_agent(key, message)
+    
+    # Check if the key is a valid integer
+    if not is_valid_int(key):
+        return "Invalid key, cannot message agent."
+    
+    agent_response = agents.message_agent(int(key), message)
 
     # Speak response
     if cfg.speak_mode: