From 5be279261230bc37ec925da5e7894701525e5edb Mon Sep 17 00:00:00 2001
From: user <user@localhost>
Date: Sat, 21 Dec 2024 15:40:19 -0500
Subject: [PATCH] Improve tool return values.

---
 ra_aid/tools/agent.py | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/ra_aid/tools/agent.py b/ra_aid/tools/agent.py
index 39d9e7b..4d14223 100644
--- a/ra_aid/tools/agent.py
+++ b/ra_aid/tools/agent.py
@@ -53,9 +53,10 @@ def request_research(query: str) -> Dict[str, Any]:
         
     return {
         "completion_message": completion_message,
-        "facts": get_memory_value("key_facts"),
-        "files": list(get_related_files()),
-        "notes": get_memory_value("research_notes"),
+        "key_facts": get_memory_value("key_facts"),
+        "related_files": list(get_related_files()),
+        "research_notes": get_memory_value("research_notes"),
+        "key_snippets": get_memory_value("key_snippets"),
         "success": success,
         "reason": reason
     }
@@ -103,9 +104,10 @@ def request_research_and_implementation(query: str) -> Dict[str, Any]:
         
     return {
         "completion_message": completion_message,
-        "facts": get_memory_value("key_facts"),
-        "files": list(get_related_files()),
-        "notes": get_memory_value("research_notes"),
+        "key_facts": get_memory_value("key_facts"),
+        "related_files": list(get_related_files()),
+        "research_notes": get_memory_value("research_notes"),
+        "key_snippets": get_memory_value("key_snippets"),
         "success": success,
         "reason": reason
     }
@@ -159,8 +161,9 @@ def request_task_implementation(task_spec: str) -> Dict[str, Any]:
     _global_memory['completion_state'] = False
         
     return {
-        "facts": get_memory_value("key_facts"),
-        "files": list(get_related_files()),
+        "key_facts": get_memory_value("key_facts"),
+        "related_files": list(get_related_files()),
+        "key_snippets": get_memory_value("key_snippets"),
         "completion_message": completion_message,
         "success": success,
         "reason": reason
@@ -208,8 +211,9 @@ def request_implementation(task_spec: str) -> Dict[str, Any]:
         
     return {
         "completion_message": completion_message,
-        "facts": get_memory_value("key_facts"),
-        "files": list(get_related_files()),
+        "key_facts": get_memory_value("key_facts"),
+        "related_files": list(get_related_files()),
+        "key_snippets": get_memory_value("key_snippets"),
         "success": success,
         "reason": reason
     }