Skip to content

Commit

Permalink
Add AI sidekick command
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
TechNickAI committed Jul 1, 2023
1 parent 3e45ebf commit 80118f1
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 2 deletions.
34 changes: 34 additions & 0 deletions aicodebot/agents.py
Original file line number Diff line number Diff line change
@@ -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")
21 changes: 19 additions & 2 deletions aicodebot/cli.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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 #
# ---------------------------------------------------------------------------- #
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 80118f1

Please sign in to comment.