From 5545c6fe41a55e590ed7f8195855fde9fa8048c6 Mon Sep 17 00:00:00 2001
From: Jay Hack <jayhack@users.noreply.github.com>
Date: Sun, 16 Feb 2025 10:28:40 -0800
Subject: [PATCH] feat: agent cli (#520)

---
 src/codegen/cli/cli.py                 |   2 +
 src/codegen/cli/commands/agent/main.py | 125 +++++++++++++++++++++++++
 2 files changed, 127 insertions(+)
 create mode 100644 src/codegen/cli/commands/agent/main.py

diff --git a/src/codegen/cli/cli.py b/src/codegen/cli/cli.py
index b78801a07..465872c36 100644
--- a/src/codegen/cli/cli.py
+++ b/src/codegen/cli/cli.py
@@ -1,6 +1,7 @@
 import rich_click as click
 from rich.traceback import install
 
+from codegen.cli.commands.agent.main import agent_command
 from codegen.cli.commands.config.main import config_command
 from codegen.cli.commands.create.main import create_command
 from codegen.cli.commands.deploy.main import deploy_command
@@ -29,6 +30,7 @@ def main():
 
 
 # Wrap commands with error handler
+main.add_command(agent_command)
 main.add_command(init_command)
 main.add_command(logout_command)
 main.add_command(login_command)
diff --git a/src/codegen/cli/commands/agent/main.py b/src/codegen/cli/commands/agent/main.py
new file mode 100644
index 000000000..acf3b4a84
--- /dev/null
+++ b/src/codegen/cli/commands/agent/main.py
@@ -0,0 +1,125 @@
+import uuid
+import warnings
+
+import rich_click as click
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.prompt import Prompt
+
+from codegen import Codebase
+from codegen.extensions.langchain.agent import create_agent_with_tools
+from codegen.extensions.langchain.tools import (
+    CreateFileTool,
+    DeleteFileTool,
+    EditFileTool,
+    ListDirectoryTool,
+    MoveSymbolTool,
+    RenameFileTool,
+    RevealSymbolTool,
+    SearchTool,
+    ViewFileTool,
+)
+
+# Suppress specific warnings
+warnings.filterwarnings("ignore", message=".*Helicone.*")
+warnings.filterwarnings("ignore", message=".*LangSmith.*")
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
+console = Console()
+
+WELCOME_ART = r"""[bold blue]
+   ____          _
+  / ___|___   __| | ___  __ _  ___ _ __
+ | |   / _ \ / _` |/ _ \/ _` |/ _ \ '_ \
+ | |__| (_) | (_| |  __/ (_| |  __/ | | |
+  \____\___/ \__,_|\___|\__, |\___|_| |_|
+                        |___/
+
+[/bold blue]
+"""
+
+
+@click.command(name="agent")
+@click.option("--query", "-q", default=None, help="Initial query for the agent.")
+def agent_command(query: str):
+    """Start an interactive chat session with the Codegen AI agent."""
+    # Show welcome message
+    console.print(WELCOME_ART)
+
+    # Initialize codebase from current directory
+    with console.status("[bold green]Initializing codebase...[/bold green]"):
+        codebase = Codebase("./")
+
+    # Helper function for agent to print messages
+    def say(message: str):
+        console.print()  # Add blank line before message
+        markdown = Markdown(message)
+        console.print(markdown)
+        console.print()  # Add blank line after message
+
+    # Initialize tools
+    tools = [
+        ViewFileTool(codebase),
+        ListDirectoryTool(codebase),
+        SearchTool(codebase),
+        CreateFileTool(codebase),
+        DeleteFileTool(codebase),
+        RenameFileTool(codebase),
+        MoveSymbolTool(codebase),
+        RevealSymbolTool(codebase),
+        EditFileTool(codebase),
+        # RunBashCommandTool(codebase),
+    ]
+
+    # Initialize chat history with system message
+    chat_history = [
+        SystemMessage(
+            content="""You are a helpful AI assistant with access to the local codebase.
+You can help with code exploration, editing, and general programming tasks.
+Always explain what you're planning to do before taking actions."""
+        )
+    ]
+
+    # Get initial query if not provided via command line
+    if not query:
+        console.print("[bold]Welcome to the Codegen CLI Agent![/bold]")
+        console.print("I'm an AI assistant that can help you explore and modify code in this repository.")
+        console.print("I can help with tasks like viewing files, searching code, making edits, and more.")
+        console.print()
+        console.print("What would you like help with today?")
+        console.print()
+        query = Prompt.ask("[bold]>[/bold]")  # Simple arrow prompt
+
+    # Create the agent
+    agent = create_agent_with_tools(codebase, tools, chat_history=chat_history)
+
+    # Main chat loop
+    while True:
+        if not query:  # Only prompt for subsequent messages
+            user_input = Prompt.ask("\n[bold]>[/bold]")  # Simple arrow prompt
+        else:
+            user_input = query
+            query = None  # Clear the initial query so we enter the prompt flow
+
+        if user_input.lower() in ["exit", "quit"]:
+            break
+
+        # Add user message to chat history
+        chat_history.append(HumanMessage(content=user_input))
+
+        # Invoke the agent
+        with console.status("[bold green]Agent is thinking...") as status:
+            try:
+                session_id = str(uuid.uuid4())
+                result = agent.invoke(
+                    {"input": user_input},
+                    config={"configurable": {"session_id": session_id}},
+                )
+                # Update chat history with AI's response
+                if result.get("output"):
+                    say(result["output"])
+                    chat_history.append(AIMessage(content=result["output"]))
+            except Exception as e:
+                console.print(f"[bold red]Error during agent execution:[/bold red] {e}")
+                break