Enhance AI sidekick with file context support 📁

This commit introduces the ability for the AI sidekick to use the contents of specified files as context during its session. This is achieved by adding a new function `generate_files_context` which reads the contents of the provided files and formats them into a string that can be used as context.

The `sidekick` function in `cli.py` has been updated to accept a list of files as arguments, generate the file context, and pass it along with the user's request to the `generate_sidekick_prompt` function. The token size calculations and model selection have also been adjusted to account for the additional context.

The `generate_sidekick_prompt` function now accepts an additional `files` parameter and includes the generated file context in the prompt template. The prompt template itself has been updated to include a placeholder for the file context.

The import statements in `cli.py` and `__init__.py` have been updated to include the new `generate_files_context` function.

This enhancement provides the AI sidekick with more context, potentially improving its ability to assist the user.

Author: TechNickAI

aicodebot/cli.py

@@ -1,6 +1,6 @@
 from aicodebot import version as aicodebot_version
 from aicodebot.helpers import exec_and_get_output, get_token_length, git_diff_context
-from aicodebot.prompts import generate_sidekick_prompt
+from aicodebot.prompts import generate_files_context, generate_sidekick_prompt
 from dotenv import load_dotenv
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.chains import LLMChain
@@ -242,7 +242,7 @@ def review(commit, verbose):
     prompt = load_prompt(Path(__file__).parent / "prompts" / "review.yaml")
 
     # Check the size of the diff context and adjust accordingly
-    response_token_size = DEFAULT_MAX_TOKENS / 2
+    response_token_size = DEFAULT_MAX_TOKENS
     request_token_size = get_token_length(diff_context) + get_token_length(prompt.template)
     model = get_llm_model(request_token_size)
     if verbose:
@@ -267,15 +267,29 @@ def review(commit, verbose):
 @cli.command
 @click.option("--request", "-r", help="What to ask your sidekick to do")
 @click.option("-v", "--verbose", count=True)
-def sidekick(request, verbose):
-    """EXPERIMENTAL: Coding help from your AI sidekick"""
+@click.argument("files", nargs=-1)
+def sidekick(request, verbose, files):
+    """
+    EXPERIMENTAL: Coding help from your AI sidekick\n
+    FILES: List of files to be used as context for the session
+    """
+
     console.print("This is an experimental feature. Play with it, but don't count on it.", style=warning_style)
 
     setup_environment()
 
-    # Generate the prompt with the appropriate context
-    prompt = generate_sidekick_prompt(request)
-    model = get_llm_model(get_token_length(prompt.template))
+    # Pull in context. Right now it's just the contents of files that we passed in.
+    # Soon, we could add vector embeddings of:
+    # imported code / modules / libraries
+    # Style guides/reference code
+    # git history
+    context = generate_files_context(files)
+
+    # Generate the prompt and set up the model
+    prompt = generate_sidekick_prompt(request, files)
+    model = get_llm_model(get_token_length(prompt.template) + get_token_length(context))
+    if verbose:
+        console.print(f"Context token size: {get_token_length(context)}, using model: {model}")
 
     llm = ChatOpenAI(
         model=model,
@@ -286,32 +300,30 @@ def sidekick(request, verbose):
     )
 
     # Set up the chain
-    chain = LLMChain(
-        llm=llm,
-        prompt=prompt,
-        memory=ConversationTokenBufferMemory(llm=llm, max_token_limit=DEFAULT_MAX_TOKENS),
-        verbose=verbose,
+    memory = ConversationTokenBufferMemory(
+        memory_key="chat_history", input_key="task", llm=llm, max_token_limit=DEFAULT_MAX_TOKENS
     )
+    chain = LLMChain(llm=llm, prompt=prompt, memory=memory, verbose=verbose)
 
     while True:  # continuous loop for multiple questions
         if request:
-            user_input = request
+            human_input = request
             request = None  # clear the command line request once we've handled it
         else:
-            user_input = click.prompt(
+            human_input = click.prompt(
                 "Enter a question OR (q) quit, OR (e) edit for entering a question in your editor\n>>>",
                 prompt_suffix="",
             )
-            if user_input.lower() == "q":
+            if human_input.lower() == "q":
                 break
-            elif user_input.lower() == "e":
-                user_input = click.edit()
+            elif human_input.lower() == "e":
+                human_input = click.edit()
 
         with Live(Markdown(""), auto_refresh=True) as live:
             callback = RichLiveCallbackHandler(live)
             callback.buffer = []
             llm.callbacks = [callback]
-            chain.run(user_input)
+            chain.run({"task": human_input, "context": context})
 
 
 # ---------------------------------------------------------------------------- #

aicodebot/prompts/__init__.py

@@ -1,3 +1,3 @@
-from .sidekick import generate_sidekick_prompt
+from .sidekick import generate_files_context, generate_sidekick_prompt
 
-__all__ = ["generate_sidekick_prompt"]
+__all__ = ["generate_sidekick_prompt", "generate_files_context"]

aicodebot/prompts/sidekick.py

@@ -1,18 +1,31 @@
 from langchain import PromptTemplate
+from pathlib import Path
 
-SIDEKICK_PREFIX = """You are a pair programming assistant named AICodeBot, acting as a sidekick to a human developer.
+SIDEKICK_TEMPLATE = """You are a pair programming assistant named AICodeBot, acting as a sidekick to a human developer.
 If you aren't sure what to do, you can ask the human for more clarification.
 
-Relevant history:
-{history}
-End History
+Relevant chat history:
+{chat_history}
+End chat history
+{context}
 
 Conversation with the human developer:
 Human: {task}
 AICodeBot:
 """
 
 
+def generate_files_context(files):
+    """Generate the files context for the sidekick prompt."""
+    files_context = "Here are the relevant files we are working with in this session:\n"
+    for file_name in files:
+        contents = Path(file_name).read_text()
+        files_context += f"--- START OF FILE: {file_name} ---\n"
+        files_context += contents
+        files_context += f"\n--- END OF FILE: {file_name} ---\n\n"
+    return files_context
+
+
 def generate_sidekick_prompt(task, files=None):
     """Generates a prompt for the sidekick workflow."""
-    return PromptTemplate(template=SIDEKICK_PREFIX, input_variables=["history", "task"])
+    return PromptTemplate(template=SIDEKICK_TEMPLATE, input_variables=["chat_history", "task", "context"])