Clean up (#231)

* Updates

* Lots of clean up

* Tell me how prompt

* Remove unnecessary ref

* Add custom prompt injection

* Lint, clean up

Author: Josh-XT

AgentLLM.py

@@ -14,6 +14,7 @@
 from json.decoder import JSONDecodeError
 import spacy
 from spacy.cli import download
+from CustomPrompt import CustomPrompt
 
 try:
     nlp = spacy.load("en_core_web_sm")
@@ -27,7 +28,7 @@ class AgentLLM:
     def __init__(self, agent_name: str = "AgentLLM", primary_objective=None):
         self.CFG = Agent(agent_name)
         self.primary_objective = primary_objective
-        self.initialize_task_list()
+        self.task_list = deque([])
         self.commands = Commands(agent_name)
         self.available_commands = self.commands.get_available_commands()
         self.web_requests = web_requests()
@@ -60,60 +61,56 @@ def __init__(self, agent_name: str = "AgentLLM", primary_objective=None):
         self.agent_config = self.CFG.load_agent_config(self.agent_name)
         self.output_list = []
         self.stop_running_event = None
-        self.instruct = self.CFG.instruct
 
     def get_output_list(self):
         return self.output_list
 
-    def trim_context(self, context: List[str], max_tokens: int) -> List[str]:
-        trimmed_context = []
-        total_tokens = 0
-        for item in context:
-            item_tokens = len(nlp(item))
-            if total_tokens + item_tokens <= max_tokens:
-                trimmed_context.append(item)
-                total_tokens += item_tokens
-            else:
-                break
-        return trimmed_context
+    def get_commands_string(self):
+        if len(self.available_commands) == 0:
+            return "No commands."
+
+        enabled_commands = filter(
+            lambda command: command.get("enabled", True), self.available_commands
+        )
+
+        friendly_names = map(
+            lambda command: f"{command['friendly_name']} - {command['name']}({command['args']})",
+            enabled_commands,
+        )
+        return "\n".join(friendly_names)
 
     def run(
         self,
         task: str,
         max_context_tokens: int = 500,
         long_term_access: bool = False,
         commands_enabled: bool = True,
-        instruction: bool = False,
+        prompt: str = "",
+        **kwargs,
     ):
-        if self.CFG.NO_MEMORY:
+        cp = CustomPrompt()
+        if prompt == "":
             prompt = task
+        elif prompt in ["execute", "task", "priority"]:
+            prompt = cp.get_model_prompt(prompt_name=prompt, model=self.CFG.AI_MODEL)
         else:
-            self.CFG.log_interaction("USER", task)
-            context = self.context_agent(
-                query=task, top_results_num=3, long_term_access=long_term_access
-            )
-            context = self.trim_context(context, max_context_tokens)
-            prompt = self.get_prompt_with_context(task=task, context=context)
-        if instruction:
-            # Command and prompt injection for instruction mode
-            instruction_prompt = self.CFG.INSTRUCT_PROMPT
-            prompt = instruction_prompt.replace("{task}", task)
-            prompt = prompt.replace("{AGENT_NAME}", self.agent_name)
-
-            enabled_commands = filter(
-                lambda command: command.get("enabled", True), self.available_commands
-            )
-
-            friendly_names = map(
-                lambda command: f"{command['friendly_name']} - {command['name']}({command['args']})",
-                enabled_commands,
-            )
+            prompt = CustomPrompt().get_prompt(prompt)
 
-            if len(self.available_commands) == 0:
-                prompt = prompt.replace("{COMMANDS}", "No commands.")
-            else:
-                prompt = prompt.replace("{COMMANDS}", "\n".join(friendly_names))
-        self.response = self.instruct(prompt)
+        context = self.context_agent(
+            query=task,
+            top_results_num=3,
+            long_term_access=long_term_access,
+            max_tokens=max_context_tokens,
+        )
+        formatted_prompt = prompt.format(
+            task=task,
+            agent_name=self.agent_name,
+            commands=self.get_commands_string(),
+            context=context,
+            **kwargs,
+        )
+        self.CFG.log_interaction("USER", task)
+        self.response = self.CFG.instruct(formatted_prompt)
         if not self.CFG.NO_MEMORY:
             self.store_result(task, self.response)
             self.CFG.log_interaction(self.agent_name, self.response)
@@ -190,7 +187,11 @@ def store_result(self, task_name: str, result: str):
                 )
 
     def context_agent(
-        self, query: str, top_results_num: int, long_term_access: bool = False
+        self,
+        query: str,
+        top_results_num: int,
+        long_term_access: bool = False,
+        max_tokens: int = 500,
     ) -> List[str]:
         if long_term_access:
             interactions = self.CFG.memory["interactions"]
@@ -209,12 +210,16 @@ def context_agent(
                 include=["metadatas"],
             )
             context = [item["result"] for item in results["metadatas"][0]]
-        return context
-
-    def get_prompt_with_context(self, task: str, context: List[str]) -> str:
-        context_str = "\n\n".join(context)
-        prompt = f"Task: {task}\n\nContext: {context_str}\n\nResponse:"
-        return prompt
+        trimmed_context = []
+        total_tokens = 0
+        for item in context:
+            item_tokens = len(nlp(item))
+            if total_tokens + item_tokens <= max_tokens:
+                trimmed_context.append(item)
+                total_tokens += item_tokens
+            else:
+                break
+        return "\n".join(trimmed_context)
 
     def chunk_content(self, content: str, max_length: int = 500) -> List[str]:
         content_chunks = []
@@ -233,40 +238,30 @@ def chunk_content(self, content: str, max_length: int = 500) -> List[str]:
             content_chunks.append(" ".join(chunk))
         return content_chunks
 
-    def set_agent_name(self, agent_name):
-        self.agent_name = agent_name
-
     def get_status(self):
         try:
             return not self.stop_running_event.is_set()
         except:
             return False
 
-    def initialize_task_list(self):
-        self.task_list = deque([])
-
     def update_output_list(self, output):
         print(
             self.CFG.save_task_output(self.agent_name, output, self.primary_objective)
         )
-        # self.output_list.append(output)
-
-    def set_objective(self, new_objective):
-        self.primary_objective = new_objective
 
     def task_creation_agent(
         self, result: Dict, task_description: str, task_list: List[str]
     ) -> List[Dict]:
-        prompt = self.CFG.TASK_PROMPT
-        # Prompt Engineering - Objective
-        prompt = prompt.replace("{objective}", self.primary_objective)
-        # Prompt Engineering - Result
-        prompt = prompt.replace("{result}", str(result))
-        # Prompt Engineering - Task Description
-        prompt = prompt.replace("{task_description}", task_description)
-        # Prompt Engineering - Task List
-        prompt = prompt.replace("{tasks}", ", ".join(task_list))
-        response = self.run(prompt, commands_enabled=False)
+        response = self.run(
+            task=self.primary_objective,
+            commands_enabled=False,
+            prompt="task",
+            objective=self.primary_objective,
+            result=result,
+            task_description=task_description,
+            tasks=", ".join(task_list),
+        )
+
         lines = response.split("\n") if "\n" in response else [response]
         new_tasks = [
             re.sub(r"^.*?(\d)", r"\1", line)
@@ -278,14 +273,16 @@ def task_creation_agent(
     def prioritization_agent(self):
         task_names = [t["task_name"] for t in self.task_list]
         next_task_id = len(self.task_list) + 1
-        prompt = self.CFG.PRIORITY_PROMPT
-        # Prompt Engineering - Objective
-        prompt = prompt.replace("{objective}", self.primary_objective)
-        # Prompt Engineering - Task ID
-        prompt = prompt.replace("{next_task_id}", str(next_task_id))
-        # Prompt Engineering - Task Names
-        prompt = prompt.replace("{task_names}", ", ".join(task_names))
-        response = self.run(prompt, commands_enabled=False)
+
+        response = self.run(
+            task=self.primary_objective,
+            commands_enabled=False,
+            prompt="priority",
+            objective=self.primary_objective,
+            tasks=", ".join(task_names),
+            next_task_id=next_task_id,
+        )
+
         lines = response.split("\n") if "\n" in response else [response]
         new_tasks = [
             re.sub(r"^.*?(\d)", r"\1", line)
@@ -300,35 +297,8 @@ def prioritization_agent(self):
                 task_name = task_parts[1].strip()
                 self.task_list.append({"task_id": task_id, "task_name": task_name})
 
-    def execution_agent(self, task: str, task_id: int) -> str:
-        context = self.context_agent(
-            query=f"{self.primary_objective} {task}", top_results_num=5
-        )
-        prompt = self.CFG.EXECUTION_PROMPT
-        # Prompt Engineering - Objective
-        prompt = prompt.replace("{objective}", self.primary_objective)
-        # Prompt Engineering - Task
-        prompt = prompt.replace("{task}", task)
-        # Prompt Engineering - Context
-        prompt = prompt.replace("{context}", "\n".join(context))
-        # Prompt Engineering - Commands
-
-        enabled_commands = filter(
-            lambda command: command.get("enabled", True), self.available_commands
-        )
-
-        friendly_names = map(
-            lambda command: f"{command['friendly_name']} - {command['name']}({command['args']})",
-            enabled_commands,
-        )
-
-        if task_id == 0 or len(self.available_commands) == 0:
-            prompt = prompt.replace("{COMMANDS}", "No commands.")
-        else:
-            prompt = prompt.replace("{COMMANDS}", "\n".join(friendly_names))
-        return self.run(prompt)
-
-    def run_task(self, stop_event):
+    def run_task(self, stop_event, objective):
+        self.primary_objective = objective
         self.update_output_list(
             f"Starting task with objective: {self.primary_objective}.\n\n"
         )
@@ -343,7 +313,7 @@ def run_task(self, stop_event):
             self.update_output_list(
                 f"\nExecuting task {task['task_id']}: {task['task_name']}\n"
             )
-            result = self.execution_agent(task["task_name"], task["task_id"])
+            result = self.run(task=task["task_name"], prompt="execution")
             self.update_output_list(f"\nTask Result:\n\n{result}\n")
             new_tasks = self.task_creation_agent(
                 {"data": result},

CustomPrompt.py

@@ -32,3 +32,7 @@ def delete_prompt(self, prompt_name):
     def update_prompt(self, prompt_name, prompt):
         with open(os.path.join("prompts", f"{prompt_name}.txt"), "w") as f:
             f.write(prompt)
+
+    def get_model_prompt(self, prompt_name, model="default"):
+        with open(f"model-prompts/{model}/{prompt_name}.txt", "r") as f:
+            return f.read()

app.py

@@ -263,11 +263,11 @@ async def start_task_agent(agent_name: str, objective: Objective) -> ResponseMes
     # If it doesn't exist, create it.
     if agent_name not in CFG.agent_instances:
         CFG.agent_instances[agent_name] = AgentLLM(agent_name)
-    CFG.agent_instances[agent_name].set_objective(objective.objective)
     stop_event = threading.Event()
     agent_stop_events[agent_name] = stop_event
     agent_thread = threading.Thread(
-        target=CFG.agent_instances[agent_name].run_task, args=(stop_event,)
+        target=CFG.agent_instances[agent_name].run_task,
+        args=(stop_event, objective.objective),
     )
     agent_threads[agent_name] = agent_thread
     agent_thread.start()

prompts/Tell Me How.txt

@@ -0,0 +1,6 @@
+Mark: Tony, help!
+Tony: What's wrong, Mark?
+Mark: I forgot how to {task}!
+Tony: That's easy! Do you want me to show you how?
+Mark: Sure!
+(continue as Tony)

provider/gpt4free.py

@@ -1,5 +1,6 @@
 import gpt4free
 
+
 class Gpt4freeProvider:
     def __init__(
         self,
@@ -8,11 +9,13 @@ def __init__(
         MAX_TOKENS: int = 4096,
         **kwargs
     ):
-        self.requirements = ['gpt4free']
+        self.requirements = ["gpt4free"]
         self.AI_MODEL = AI_MODEL
         self.AI_TEMPERATURE = AI_TEMPERATURE
         self.MAX_TOKENS = MAX_TOKENS
 
     def instruct(self, prompt):
-        response = gpt4free.Completion.create(gpt4free.Provider.UseLess, prompt=prompt, model=self.AI_MODEL)
+        response = gpt4free.Completion.create(
+            gpt4free.Provider.UseLess, prompt=prompt, model=self.AI_MODEL
+        )
         return response["text"]

provider/huggingface.py

@@ -5,16 +5,17 @@
     AutoModel,
     AutoModelForSeq2SeqLM,
     T5Tokenizer,
-    AutoConfig
+    AutoConfig,
 )
 
+
 class HuggingfaceProvider:
     def __init__(
         self,
         AI_MODEL: str = "gpt2",
         AI_TEMPERATURE: float = 0.7,
         MAX_TOKENS: int = 1024,
-        **kwargs
+        **kwargs,
     ):
         self.requirements = ["transformers", "torch", "accelerate"]
         self.AI_MODEL = AI_MODEL
@@ -25,7 +26,9 @@ def instruct(self, prompt):
         try:
             model_path = self.AI_MODEL
             if "chatglm" in model_path:
-                tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_path, trust_remote_code=True
+                )
                 model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
             elif "dolly" in model_path:
                 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
@@ -50,18 +53,18 @@ def instruct(self, prompt):
                     model_path, low_cpu_mem_usage=True
                 )
 
-            input_ids = tokenizer(prompt, return_tensors='pt').input_ids
+            input_ids = tokenizer(prompt, return_tensors="pt").input_ids
 
             output_ids = model.generate(
                 input_ids,
                 pad_token_id=tokenizer.eos_token_id,
                 temperature=self.AI_TEMPERATURE,
                 max_new_tokens=self.MAX_TOKENS,
-                no_repeat_ngram_size=2
+                no_repeat_ngram_size=2,
             )
             input_length = 1 if model.config.is_encoder_decoder else len(input_ids[0])
             output_ids = output_ids[0][input_length:]
-    
+
             return tokenizer.decode(output_ids, skip_special_tokens=True).strip()
         except Exception as e:
             return f"Transformer Error: {e}"