From 3dea6eaaeaae48507947c3f0d63ba42c06d436bd Mon Sep 17 00:00:00 2001 From: jayhack Date: Sun, 16 Feb 2025 02:34:50 -0800 Subject: [PATCH 1/2] . --- src/codegen/cli/cli.py | 2 + src/codegen/cli/commands/agent/main.py | 97 ++++++++++++++++++++++++++ 2 files changed, 99 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..884a52be9 --- /dev/null +++ b/src/codegen/cli/commands/agent/main.py @@ -0,0 +1,97 @@ +import uuid + +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, +) + +console = Console() + + +@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.""" + # 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): + markdown = Markdown(message) + console.print(markdown) + + # Initialize tools + tools = [ + ViewFileTool(codebase), + ListDirectoryTool(codebase), + SearchTool(codebase), + CreateFileTool(codebase), + DeleteFileTool(codebase), + RenameFileTool(codebase), + MoveSymbolTool(codebase), + RevealSymbolTool(codebase), + EditFileTool(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 AI Agent! What can I help you with?[/bold]") + query = Prompt.ask(":speech_balloon: [bold]Your message[/bold]") + + # 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(":speech_balloon: [bold]Your message[/bold]") + 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 From d782df204e272bdb2151846e64ac976017a6bc12 Mon Sep 17 00:00:00 2001 From: jayhack Date: Sun, 16 Feb 2025 02:55:59 -0800 Subject: [PATCH 2/2] . --- src/codegen/cli/commands/agent/main.py | 34 +++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/src/codegen/cli/commands/agent/main.py b/src/codegen/cli/commands/agent/main.py index 884a52be9..acf3b4a84 100644 --- a/src/codegen/cli/commands/agent/main.py +++ b/src/codegen/cli/commands/agent/main.py @@ -1,4 +1,5 @@ import uuid +import warnings import rich_click as click from langchain_core.messages import AIMessage, HumanMessage, SystemMessage @@ -20,21 +21,42 @@ 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 = [ @@ -47,6 +69,7 @@ def say(message: str): MoveSymbolTool(codebase), RevealSymbolTool(codebase), EditFileTool(codebase), + # RunBashCommandTool(codebase), ] # Initialize chat history with system message @@ -60,8 +83,13 @@ def say(message: str): # Get initial query if not provided via command line if not query: - console.print("[bold]Welcome to the Codegen AI Agent! What can I help you with?[/bold]") - query = Prompt.ask(":speech_balloon: [bold]Your message[/bold]") + 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) @@ -69,7 +97,7 @@ def say(message: str): # Main chat loop while True: if not query: # Only prompt for subsequent messages - user_input = Prompt.ask(":speech_balloon: [bold]Your message[/bold]") + 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