From 80118f1799304297378ba048e5301f40542efdea Mon Sep 17 00:00:00 2001
From: Nick Sullivan <nick@sullivanflock.com>
Date: Sat, 1 Jul 2023 14:47:08 -0600
Subject: [PATCH] Add AI sidekick command

This commit introduces a new feature where an AI sidekick can assist with tasks. The AI sidekick is initialized with a set of tools for file management and human input. The sidekick can be invoked via the new `sidekick` command in the CLI.

Additionally, the `review` command has been updated to clarify that it can review either unstaged changes or a specified commit. The help text for the `--commit` option has been updated to reflect this.

Minor changes have also been made to the `get_llm_model` function to allow for a default token size.
---
 aicodebot/agents.py | 34 ++++++++++++++++++++++++++++++++++
 aicodebot/cli.py    | 21 +++++++++++++++++++--
 2 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 aicodebot/agents.py

diff --git a/aicodebot/agents.py b/aicodebot/agents.py
new file mode 100644
index 0000000..f2487c7
--- /dev/null
+++ b/aicodebot/agents.py
@@ -0,0 +1,34 @@
+""" LangChain Agent Setup """
+from langchain.agents import AgentType, initialize_agent, load_tools
+from langchain.agents.agent_toolkits import FileManagementToolkit
+
+SIDEKICK_PROMPT_PREFIX = """
+You are a programming assistant named AICodeBot, acting as a sidekick to a human developer.
+You are running in the local repository, so you can interact with the files in the repository
+to get more information about the code. If you aren't sure what to do, you can ask the human.
+
+If asking the human for more clarification would produce better results, you can ask the human for more information.
+
+Before changing any local files, you should ALWAYS check with the human developer first, explaining what you are doing,
+you can give human multiple choices.
+"""
+
+
+def get_agent(name, llm, verbose):
+    """Get the agent by name"""
+    if name == "sidekick":
+        # Set up the tools. Basic local file management first
+        tools = FileManagementToolkit(selected_tools=["read_file", "write_file", "list_directory"]).get_tools()
+        tools += load_tools(["human"])  # Human input
+
+        # Set up the agent
+        return initialize_agent(
+            tools,
+            llm,
+            agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+            verbose=verbose,
+            return_intermediate_steps=True,
+            agent_kwargs={"prefix": SIDEKICK_PROMPT_PREFIX, "verbose": verbose},
+        )
+    else:
+        raise ValueError(f"Agent {name} not found")
diff --git a/aicodebot/cli.py b/aicodebot/cli.py
index 58812bc..07448ba 100644
--- a/aicodebot/cli.py
+++ b/aicodebot/cli.py
@@ -1,4 +1,5 @@
 from aicodebot import version as aicodebot_version
+from aicodebot.agents import get_agent
 from aicodebot.helpers import exec_and_get_output, get_token_length, git_diff_context
 from dotenv import load_dotenv
 from langchain.chains import LLMChain
@@ -200,7 +201,7 @@ def fun_fact(verbose):
 
 
 @cli.command
-@click.option("-c", "--commit", help="The commit hash to review.")
+@click.option("-c", "--commit", help="The commit hash to review (otherwise look at [un]staged changes).")
 @click.option("-v", "--verbose", count=True)
 def review(commit, verbose):
     """Do a code review, with [un]staged changes, or a specified commit."""
@@ -232,6 +233,22 @@ def review(commit, verbose):
         console.print(response, style=bot_style)
 
 
+@cli.command
+@click.option("--task", "-t", help="The task you want to perform - a description of what you want to do.")
+@click.option("-v", "--verbose", count=True)
+def sidekick(task, verbose):
+    """Get help with a task from your AI sidekick."""
+    setup_environment()
+
+    model = get_llm_model()
+    llm = ChatOpenAI(model=model, temperature=DEFAULT_TEMPERATURE, max_tokens=3500, verbose=verbose)
+
+    agent = get_agent("sidekick", llm, True)
+
+    response = agent({"input": task})
+    console.print(response, style=bot_style)
+
+
 # ---------------------------------------------------------------------------- #
 #                               Helper functions                               #
 # ---------------------------------------------------------------------------- #
@@ -295,7 +312,7 @@ def setup_environment():
     )
 
 
-def get_llm_model(token_size):
+def get_llm_model(token_size=0):
     # https://platform.openai.com/docs/models/gpt-3-5
     # We want to use GPT-4, if it is available for this OPENAI_API_KEY, otherwise GPT-3.5
     # We also want to use the largest model that supports the token size we need