From c5bbefa60065589419315abb138edc5980cba21a Mon Sep 17 00:00:00 2001
From: Om Raheja <omr@rslp.org>
Date: Fri, 30 Aug 2024 15:52:20 -0400
Subject: [PATCH] added actual pyfzf implementation

---
 agent_repl.py       | 40 +++++++++++++++++++++++++++++-----------
 aios/utils/utils.py |  5 +++++
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/agent_repl.py b/agent_repl.py
index 2cedb0bf..fdce5dda 100644
--- a/agent_repl.py
+++ b/agent_repl.py
@@ -1,8 +1,10 @@
 import argparse
 from typing import List, Tuple
 
-from aios.utils.utils import delete_directories, parse_global_args
+from aios.utils.utils import delete_directories, humanify_agent, parse_global_args
 from aios.hooks.llm import useFactory, useKernel, useFIFOScheduler
+from pyopenagi.agents.interact import Interactor
+
 from dotenv import load_dotenv
 import warnings
 
@@ -15,6 +17,16 @@ def clean_cache(root_directory):
     }
     delete_directories(root_directory, targets)
 
+def get_all_agents() -> dict[str, str]:
+    interactor = Interactor()
+    agents = interactor.list_available_agents()
+    agent_names = {}
+    for a in agents:
+        agent_names[humanify_agent(a["agent"])] = a["agent"]
+
+    return agent_names
+
+
 def get_agent_list() -> List[Tuple[str, str]]:
     return [
         ("academic_agent", "Research academic topics"),
@@ -63,11 +75,6 @@ def main():
     parser = parse_global_args()
     args = parser.parse_args()
 
-    # try to load pyfzf
-    try:
-        import pyfzf
-    except ImportError:
-        print("pyfzf is not installed. Please install it using 'pip install pyfzf'")
 
     load_dotenv()
 
@@ -91,16 +98,27 @@ def main():
         max_workers=500
     )
 
-    agents = get_agent_list()
-    display_agents(agents)
-    chosen_agent, _ = get_user_choice(agents)
-    task = get_user_task()
 
+    # try to load pyfzf
+    chosen_agent = ""
+    try:
+        from pyfzf.pyfzf import FzfPrompt
+        fzf = FzfPrompt()
+        agents = get_all_agents()
+        chosen_agent = agents[fzf.prompt(list(agents.keys()))[0]]
+
+    except ImportError:
+        print("pyfzf is not installed. Falling back to default reader.")
+        agents = get_agent_list()
+        display_agents(agents)
+        chosen_agent, _ = "example/" + get_user_choice(agents)
+
+    task = get_user_task()
     startScheduler()
 
     try:
         agent_id = submitAgent(
-            agent_name=f"example/{chosen_agent}",
+            agent_name=chosen_agent,
             task_input=task
         )
 
diff --git a/aios/utils/utils.py b/aios/utils/utils.py
index feb4465f..d1a4d935 100755
--- a/aios/utils/utils.py
+++ b/aios/utils/utils.py
@@ -78,3 +78,8 @@ def delete_directories(root_dir, target_dirs):
                 full_path = os.path.join(dirpath, dirname)
                 # print(f"Deleting {full_path}...")
                 shutil.rmtree(full_path, ignore_errors=True)
+
+def humanify_agent(input_string: str):
+    """ turns 'author/example_agent' into 'Example Agent' """
+    last_part = input_string.split('/')[-1].replace('_', ' ')
+    return ' '.join(word.capitalize() for word in last_part.split())