improved docs and tui

om-raheja · Oct 10, 2024 · 995a117 · 995a117
1 parent 5fe06a0
commit 995a117
Show file tree

Hide file tree

Showing 5 changed files with 69 additions and 32 deletions.
diff --git a/agent_repl.py b/agent_repl.py
@@ -1,13 +1,43 @@
-import argparse
 from typing import List, Tuple
 
 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 aios.utils.state import useGlobalState
+
+from pyopenagi.manager.manager import AgentManager
 
 from dotenv import load_dotenv
 import warnings
 
+parser = parse_global_args()
+parser.add_argument("--agent_manager", default="https://my.aios.foundation",
+                        type=str, help="Agent manager url")
+parser.add_argument("--no-color-output", help="Disable color output")
+
+WHITE = "\033[0m"
+BLUE = "\033[34m"
+GREEN = "\033[32m"
+RED = "\033[31m"
+RESET = "\033[0m"
+BOLD = "\033[1m"
+
+if parser.parse_args().no_color_output:
+    WHITE = ""
+    BLUE = ""
+    GREEN = ""
+    RED = ""
+    RESET = ""
+
+
+
+args = parser.parse_args()
+
+getLLMState, setLLMState, setLLMCallback = useGlobalState()
+getFactory, setFactory, setFactoryCallback = useGlobalState()
+getManager, setManager, setManagerCallback = useGlobalState()
+
+
+setManager(AgentManager(args.agent_manager))
 
 def clean_cache(root_directory):
     targets = {
@@ -19,15 +49,15 @@ def clean_cache(root_directory):
     delete_directories(root_directory, targets)
 
 
-def get_all_agents() -> dict[str, str]:
-    interactor = Interactor()
-    agents = interactor.list_available_agents()
+def get_all_agents(
+            manager: AgentManager = getManager()
+        ) -> dict[str, str]:
+    agents = (manager.list_available_agents())
     agent_names = {}
     for a in agents:
         agent_names[humanify_agent(a["agent"])] = a["agent"]
 
-    return agent_names
-
+    return [agent["agent"] for agent in agents]
 
 def display_agents(agents: List[str]):
     print("Available Agents:")
@@ -48,13 +78,12 @@ def get_user_choice(agents: List[Tuple[str, str]]) -> Tuple[str, str]:
 
 
 def get_user_task(chosen_agent) -> str:
-    return input(f"Enter the task for the {chosen_agent}: ")
+    return input(f"{BLUE}@{chosen_agent}{RESET}: ")
 
 
 def main():
     warnings.filterwarnings("ignore")
-    parser = parse_global_args()
-    args = parser.parse_args()
+
 
     load_dotenv()
 
@@ -63,16 +92,16 @@ def main():
         max_gpu_memory=args.max_gpu_memory,
         eval_device=args.eval_device,
         max_new_tokens=args.max_new_tokens,
-        log_mode=args.llm_kernel_log_mode,
+        log_mode=args.log_mode,
         use_backend=args.use_backend,
     )
 
     startScheduler, stopScheduler = useFIFOScheduler(
-        llm=llm, log_mode=args.scheduler_log_mode, get_queue_message=None
+        llm=llm, log_mode=args.log_mode, get_queue_message=None
     )
 
     submitAgent, awaitAgentExecution = useFactory(
-        log_mode=args.agent_log_mode, max_workers=500
+        log_mode=args.log_mode, max_workers=500
     )
 
     startScheduler()
@@ -106,7 +135,7 @@ def main():
             agents = get_all_agents()
 
             # shell prompt
-            print(f"[{args.llm_name}]> ", end="")
+            print(f"{BOLD}[{args.llm_name}]>{RESET} ", end="")
 
             # check if the user put in a tab
             if getch() != "\t":
@@ -117,13 +146,13 @@ def main():
 
                 fzf = FzfPrompt()
 
-                selected = fzf.prompt(list(agents.keys()))
+                selected = fzf.prompt(agents)
 
                 if len(selected) == 0:
                     print("No agent selected. Please try again.")
                     continue
 
-                chosen_agent = agents[selected[0]]
+                chosen_agent = selected[0]
 
             except ImportError:
                 print("pyfzf is not installed. Falling back to default reader.")
@@ -138,7 +167,7 @@ def main():
                 awaitAgentExecution(agent_id)
             except Exception as e:
                 print(f"An error occurred: {str(e)}")
-                
+
     except KeyboardInterrupt:
         stopScheduler()
         clean_cache(root_directory="./")

diff --git a/aios/utils/utils.py b/aios/utils/utils.py
@@ -23,10 +23,8 @@ def parse_global_args():
     parser.add_argument('--max_gpu_memory', type=json.loads, help="Max gpu memory allocated for the LLM")
     parser.add_argument('--eval_device', type=str, help="Evaluation device")
     parser.add_argument('--max_new_tokens', type=int, default=256, help="The maximum number of new tokens for generation")
-    parser.add_argument("--agent_log_mode", type=str,default="console",choices=["console", "file"])
-    parser.add_argument("--scheduler_log_mode", type=str,default="console",choices=["console", "file"])
-    parser.add_argument("--llm_kernel_log_mode", type=str, default="console", choices=["console", "file"])
-    parser.add_argument("--use_backend", type=str, default="ollama", choices=["ollama", "vllm"])
+    parser.add_argument("--log-mode", type=str, default="console", choices=["console", "file"])
+    parser.add_argument("--use_backend", type=str, default=None, choices=["ollama", "vllm"])
 
     return parser
 

diff --git a/docs/source/aios_developer/frontend_developer.rst b/docs/source/aios_developer/frontend_developer.rst
@@ -0,0 +1,10 @@
+.. _agent_index.rst:
+
+Frontend development standards
+==================================
+
+You may choose to implement your own frontend so functionality similar to "Hey Jarvis" may be activated on voice, a potential webui can be used, or maybe you just want this to work on Windows 2000.
+
+You can use `agent_repl.py` and the `launch.py` to use the TUI and webui, respectively.
+
+more information tbd
diff --git a/pyopenagi/manager/manager.py b/pyopenagi/manager/manager.py
@@ -1,4 +1,6 @@
-# agent_manager.py
+# more advanced agent manager
+# as opposed to the simple interact.py where you could
+# upload or download agents
 
 import importlib
 import os
@@ -206,5 +208,3 @@ def load_agent(self, author: str, name: str, version: str = "latest"):
 if __name__ == '__main__':
     manager = AgentManager('http://localhost:3000/')
     manager.upload_agent('pyopenagi/agents/example/academic_agent')
-
-
diff --git a/server.py b/server.py
@@ -38,19 +38,19 @@
 
 setManager(AgentManager('https://my.aios.foundation'))
 
-# parser = parse_global_args()
-# args = parser.parse_args()
+parser = parse_global_args()
+args = parser.parse_args()
 
 # check if the llm information was specified in args
 
 setLLMState(
     useKernel(
-        llm_name='gpt-4o-mini',
-        max_gpu_memory=None,
-        eval_device=None,
-        max_new_tokens=256,
-        log_mode='console',
-        use_backend=None
+        llm_name=args.llm_name,
+        max_gpu_memory=args.max_gpu_memory,
+        eval_device=args.eval_device,
+        max_new_tokens=args.max_new_tokens,
+        log_mode=args.log_mode,
+        use_backend=args.use_backend
     )
 )