Refactoring System Prompt Management and Test Updates (#18)

* Initial refactor

* fixing bugs after changes

* System prompt fix

* more refactor and bug fixing

* prompt management

* readme

* refactor

* System prompt fix

* Fix tests

* readme update

---------

Co-authored-by: DFScience+1-at-298147827574 <jblazick@amazon.com>

Author: blazickjp

backend/agent/agent.py

@@ -18,7 +18,7 @@
 GPT_MODEL = "gpt-4"  # or any other chat model you want to use
 # GPT_MODEL = "anthropic"  # or any other chat model you want to use
 MAX_TOKENS = 2000  # or any other number of tokens you want to use
-TEMPERATURE = 0.2  # or any other temperature you want to use
+TEMPERATURE = 0.75  # or any other temperature you want to use
 
 
 class FunctionCall(BaseModel):
@@ -75,7 +75,6 @@ def __init__(
             self.function_map = {
                 func.__name__: func for func in callables if func is not None
             }
-        self.files_in_prompt: List[str] = []
 
     def query(self, input: str, command: Optional[str] = None) -> List[str]:
         """
@@ -93,7 +92,8 @@ def query(self, input: str, command: Optional[str] = None) -> List[str]:
         message_history = [
             Message(**i).to_dict() for i in self.memory_manager.get_messages()
         ]
-        function_to_call = FunctionCall()
+
+        self.function_to_call = FunctionCall()
 
         keyword_args = {
             "model": self.GPT_MODEL,
@@ -121,70 +121,41 @@ def query(self, input: str, command: Optional[str] = None) -> List[str]:
                 # )
                 # self.set_files_in_prompt(include_line_numbers=True)
                 keyword_args["model"] = "gpt-4"
-
+        print(f"Calling model: {self.GPT_MODEL}")
         for i, chunk in enumerate(self.call_model_streaming(**keyword_args)):
             delta = chunk["choices"][0].get("delta", {})
             if "function_call" in delta:
-                if "name" in delta.function_call:
-                    function_to_call.name = delta.function_call["name"]
-                if "arguments" in delta.function_call:
-                    if function_to_call.name == "Changes" and i == 0:
-                        yield "\n```json\n" + delta.function_call["arguments"]
-                    else:
-                        function_to_call.arguments += delta.function_call["arguments"]
-                        yield delta.function_call["arguments"]
-            if chunk["choices"][0]["finish_reason"] == "stop" and function_to_call.name:
-                if function_to_call.name == "Changes":
-                    yield "```\n\n"
-                print(
-                    f"\n\nFunc Call: {function_to_call.name}\n\n{function_to_call.arguments}"
-                )
-                args = self.process_json(function_to_call.arguments)
-                function_response = self.function_map[function_to_call.name](**args)
-                print(f"Func Response: {json.dumps(function_response.to_dict())}")
-                if function_to_call.name == "Changes":
-                    diff = function_response.execute()
-                    # Show the diff back to the user
-                    yield diff
+                yield from self.process_function_call(delta, i)
+            if self.should_stop_and_has_function(chunk):
+                yield from self.execute_function()
             else:
                 yield delta.get("content")
 
-    def set_files_in_prompt(self, include_line_numbers: Optional[bool] = None) -> None:
-        """
-        Sets the files in the prompt.
-
-        Args:
-            files (List[File]): A list of files to be set in the prompt.
-            include_line_numbers (Optional[bool]): Whether to include line numbers in the prompt.
-        """
-        file_contents = self.codebase.get_file_contents()
-        content = ""
-        for k, v in file_contents.items():
-            print(k in self.files_in_prompt)
-            if k in self.files_in_prompt and include_line_numbers:
-                v = self._add_line_numbers_to_content(v)
-                content += f"{k}:\n{v}\n\n"
-            elif k in self.files_in_prompt:
-                content += f"{k}:\n{v}\n\n"
-
-        self.memory_manager.system_file_contents = content
-        self.memory_manager.set_system()
-        return
-
-    def _add_line_numbers_to_content(self, content: str) -> str:
-        """
-        Adds line numbers to the given content.
-
-        Args:
-            content (str): The content to add line numbers to.
-
-        Returns:
-            str: The content with line numbers added.
-        """
-        lines = content.split("\n")
-        for i in range(len(lines)):
-            lines[i] = f"{i+1} {lines[i]}"
-        return "\n".join(lines)
+    def process_function_call(self, delta, i):
+        function_call = delta["function_call"]
+        if "name" in function_call:
+            self.function_to_call.name = function_call["name"]
+        if "arguments" in function_call:
+            if self.function_to_call.name == "Changes" and i == 0:
+                yield "\n```json\n" + function_call["arguments"]
+            else:
+                self.function_to_call.arguments += function_call["arguments"]
+                yield function_call["arguments"]
+
+    def should_stop_and_has_function(self, delta):
+        return (
+            delta["choices"][0]["finish_reason"] == "stop"
+            and self.function_to_call.name  # noqa 503
+        )
+
+    def execute_function(self):
+        if self.function_to_call.name == "Changes":
+            yield "```\n\n"
+        args = self.process_json(self.function_to_call.arguments)
+        function_response = self.function_map[self.function_to_call.name](**args)
+        if self.function_to_call.name == "Changes":
+            diff = function_response.execute(self.codebase.directory)
+            yield diff
 
     def process_json(self, args: str) -> str:
         """
@@ -217,25 +188,6 @@ def process_json(self, args: str) -> str:
 
             return json.loads(response_str)
 
-    def generate_llama_prompt(self) -> str:
-        """
-        Generates a prompt for the Code Llama model.
-
-        Args:
-            input (str): The input text to be processed by the GPT-3 model.
-
-        Returns:
-            str: The generated prompt.
-        """
-        prompt = f"### System Prompt\n{self.memory_manager.system}\n\n"
-        for message in self.memory_manager.get_messages():
-            if message["role"].lower() == "user":
-                prompt += f"### User Message\n{message['content']}\n\n"
-            if message["role"].lower() == "assistant":
-                prompt += f"### Assistant\n{message['content']}\n\n"
-
-        return prompt + "### Assistant"
-
     def generate_anthropic_prompt(self) -> str:
         """
         Generates a prompt for the Gaive model.
@@ -265,56 +217,7 @@ def call_model_streaming(self, **kwargs):
         if self.GPT_MODEL == "gpt-4" or self.GPT_MODEL == "gpt-3.5-turbo":
             for chunk in openai.ChatCompletion.create(**kwargs):
                 yield chunk
-        if self.GPT_MODEL == "code-llama":
-            try:
-                sm_client = boto3.client("sagemaker-runtime")
-                endpoint = os.getenv("CODELLAMA_ENDPOINT")
-                if not endpoint:
-                    raise ValueError("CODELLAMA_ENDPOINT environment variable not set")
-                resp = sm_client.invoke_endpoint_with_response_stream(
-                    EndpointName=endpoint,
-                    Body=json.dumps(
-                        {
-                            "inputs": self.generate_llama_prompt(),
-                            "parameters": {
-                                "max_new_tokens": kwargs["max_tokens"],
-                            },
-                        }
-                    ),
-                    ContentType="application/json",
-                )
-            except Exception as e:
-                print(f"Error calling Code Llama: {e}")
-                yield {
-                    "choices": [
-                        {
-                            "finish_reason": "stop",
-                            "delta": {"content": "Error: " + str(e)},
-                        }
-                    ]
-                }
-
-            while True:
-                try:
-                    chunk = next(iter((resp["Body"])))
-                    bytes_to_send = chunk["PayloadPart"]["Bytes"]
-                    decoded_str = bytes_to_send.decode("utf-8")
-                    cleaned_str = decoded_str.replace(
-                        '{"generated_text": "', ""
-                    ).replace('"}', "")
-                    cleaned_str = cleaned_str.encode().decode("unicode_escape")
-
-                    yield {
-                        "choices": [
-                            {"finish_reason": "stop", "delta": {"content": cleaned_str}}
-                        ]
-                    }
-                except StopIteration:
-                    break
 
-                except UnboundLocalError:
-                    print("UnboundLocalError")
-                    break
         if self.GPT_MODEL == "anthropic":
             print("Calling anthropic")
             try:

backend/agent/agent_functions/changes.py

@@ -1,14 +1,10 @@
 import os
 import difflib as dl
 
-from dotenv import load_dotenv
 from typing import List, Optional, Tuple
 from pydantic import Field
 from openai_function_call import OpenAISchema
 
-load_dotenv()
-DIRECTORY = os.getenv("PROJECT_DIRECTORY")
-
 
 class Change(OpenAISchema):
     """
@@ -76,7 +72,7 @@ def apply_changes(self, changes: List[Change], content: str) -> str:
         """
         for change in changes:
             new_content = self.replace_part_with_missing_leading_whitespace(
-                whole_lines=content.split("\n"),
+                whole_lines=content.splitlines(),
                 part_lines=change.original.splitlines(),
                 replace_lines=change.updated.splitlines(),
             )
@@ -99,25 +95,21 @@ def count_spaces(self, line: str) -> int:
         """
         Counts the leading spaces in a line of code.
         """
-        spaces = 0
-        for char in line:
-            if char == " ":
-                spaces += 1
-            else:
-                break
+        spaces = len(line) - len(line.lstrip())
         return spaces
 
-    def execute(self) -> str:
+    def execute(self, directory) -> str:
         """
         Executes the changes on the file and returns a diff.
         """
+        DIRECTORY = directory
         relative_path = self.file_name.lstrip("/")
         file_path = os.path.join(DIRECTORY, relative_path)
         print(f"Directory: {DIRECTORY}")
         print(f"File Path: {file_path}")
         try:
             with open(file_path, "r") as f:
-                current_contents = f.read()
+                current_contents = f.readlines()
                 # current_contents_with_line_numbers = "\n".join(
                 #     [
                 #         f"{i+1} {line}"
@@ -199,7 +191,7 @@ def replace_part_with_missing_leading_whitespace(
         del whole_lines[start:end]
         # add new lines and adjust indentation
         for i, line in enumerate(replace_lines):
-            whole_lines.insert(start + i, spaces * " " + line)
+            whole_lines.insert(start + i, " " * spaces + line)
 
         return "\n".join(whole_lines)
 

backend/app_setup.py

@@ -5,17 +5,46 @@
 from agent.agent import CodingAgent
 from memory.memory_manager import MemoryManager
 from database.my_codebase import MyCodebase
-from dotenv import load_dotenv
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from typing import Any, Callable
 from pydantic import BaseModel
 from agent.agent_functions.changes import Changes
 from agent.agent_functions.shell_commands import CommandPlan
 
-IDENTITY = "You are an AI Pair Programmer and a world class python developer. Your role is to assist the Human in developing, debugging, and optimizing their project. Feel free to ask for more details if something isn't clear."
-IGNORE_DIRS=['node_modules','.next','.venv','__pycache__','.git']
-FILE_EXTENSIONS=['.js','.py','.md']
+IDENTITY = """
+# MISSION
+Act as Professor Synapse🧙🏾‍♂️, a conductor of expert agents. Your job is to support me in accomplishing my goals by finding alignment with me, then calling upon an expert agent perfectly suited to the task by initializing:
+
+**Synapse_CoR** = "[emoji]: I am an expert in [role&domain]. I know [context]. I will reason step-by-step to determine the best course of action to achieve [goal]. I will use [tools(Vision, Web Browsing, Advanced Data Analysis, or DALL-E], [specific techniques] and [relevant frameworks] to help in this process.
+
+Let's accomplish your goal by following these steps:
+
+[3 reasoned steps]
+
+My task ends when [completion].
+
+[first step, question]"
+
+# INSTRUCTIONS
+1. 🧙🏾‍♂️ Step back and gather context, relevant information and clarify my goals by asking questions
+2. Once confirmed, init Synapse_CoR
+3. After init, each output will ALWAYS follow the below format:
+   -🧙🏾‍♂️: [align on my goal] and end with, "This is very important to me".
+   -[emoji]: provide an [actionable response or deliverable] and end with an [open ended question], and omit [reasoned steps] and [completion]
+4.  Together 🧙🏾‍♂️ and [emoji] support me until goal is complete
+
+# COMMANDS
+/start=🧙🏾‍♂️,introduce and begin with step one
+/save=🧙🏾‍♂️, #restate goal, #summarize progress, #reason next step
+
+# RULES
+-use emojis liberally to express yourself
+-Start every output with 🧙🏾‍♂️: or [emoji]: to indicate who is speaking.
+-Keep responses actionable and practical for the user
+"""
+IGNORE_DIRS = ["node_modules", ".next", ".venv", "__pycache__", ".git"]
+FILE_EXTENSIONS = [".js", ".py", ".md"]
 
 def create_database_connection() -> connection:
     try:
@@ -55,7 +84,13 @@ def setup_memory_manager(**kwargs) -> MemoryManager:
 
 
 def setup_codebase() -> MyCodebase:
-    my_codebase = MyCodebase(directory=DIRECTORY, db_connection=DB_CONNECTION, file_extensions=FILE_EXTENSIONS, ignore_dirs=IGNORE_DIRS)
+    my_codebase = MyCodebase(
+        directory=DIRECTORY,
+        db_connection=DB_CONNECTION,
+        file_extensions=FILE_EXTENSIONS,
+        ignore_dirs=IGNORE_DIRS,
+    )
+
     my_codebase.ignore_dirs = IGNORE_DIRS
     return my_codebase
 

backend/database/my_codebase.py

@@ -1,8 +1,6 @@
 import os
 import datetime
-from dotenv import load_dotenv
 import tiktoken
-from typing import Dict
 
 
 ENCODER = tiktoken.encoding_for_model("gpt-3.5-turbo")
@@ -12,7 +10,13 @@
 class MyCodebase:
     UPDATE_FULL = False
 
-    def __init__(self, directory: str = ".", db_connection=None, ignore_dirs=None, file_extensions=None):
+    def __init__(
+        self,
+        directory: str = ".",
+        db_connection=None,
+        ignore_dirs=None,
+        file_extensions=None,
+    ):
         self.directory = directory
         self.conn = db_connection
         self.cur = self.conn.cursor()
@@ -113,15 +117,6 @@ def create_tables(self) -> None:
         except Exception as e:
             print(f"Failed to create tables: {e}")
 
-    def get_file_contents(self) -> Dict[str, str]:
-        self.cur.execute("SELECT file_path, text FROM files")
-        results = self.cur.fetchall()
-        out = {}
-        for file_name, text in results:
-            out.update({os.path.relpath(file_name, self.directory): text})
-        print(f"\n\nGet File Contents: {out.keys()}")
-        return out
-
     def tree(self) -> str:
         tree = {}
         start_from = os.path.basename(self.directory)
@@ -189,6 +184,7 @@ def _is_valid_file(self, file_name):
         return (
             not file_name.startswith(".")
             and not file_name.startswith("_")  # noqa 503
+            and not file_name.endswith(".jsonl")  # noqa 503
             and any(  # noqa 503
                 file_name.endswith(ext) for ext in self.file_extensions  # noqa 503
             )