import os import re import streamlit as st import requests import base64 import json import shutil from urllib.parse import urlparse from git import Repo from git.exc import GitCommandError from typing import List, Dict, Any, TypedDict, Annotated import operator import asyncio from langchain.tools import StructuredTool, Tool from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.messages import BaseMessage, HumanMessage from langchain_anthropic import ChatAnthropic from langchain_community.tools import ShellTool from langgraph.prebuilt import create_react_agent from langgraph.checkpoint.memory import MemorySaver # Show title and description. # Add a radio button for mode selection mode = st.radio("Select Mode", ["Q/A", "Task"]) st.title("Coder for NextJS Templates") st.markdown( "This chatbot connects to a Next.JS Github Repository to answer questions and modify code " "given the user's prompt. Please input your repo url and github token to allow the AI to connect, then query it by asking questions or requesting feature changes! Watch video about this app [here](https://www.youtube.com/watch?v=A3XCfAVWrH4&t=17s)" ) # Ask user for their Github Repo URL, Github Token, and Anthropic API key via `st.text_input`. os.environ["LANGCHAIN_TRACING_V2"] = "true" os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com" os.environ["LANGCHAIN_PROJECT"] = "Github-Agent" os.environ["LANGCHAIN_API_KEY"] = os.getenv("LANGCHAIN_API_KEY") github_repo_url = st.text_input("Github Repo URL (e.g., https://github.com/user/repo)") # Use st.markdown for the hyperlink text st.markdown( '[How to get your Github Token](https://docs.github.com/en/enterprise-server@3.9/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens)' ) github_token = st.text_input("Enter your Github Token", type="password") # anthropic_api_key = st.text_input("Anthropic API Key", type="password") anthropic_api_key = os.getenv("ANTHROPIC_API_KEY") graph_tools = [] if not (github_repo_url and github_token and anthropic_api_key): st.info("Please add your Github Repo URL and Github Personal Token to continue.", icon="🗝️") else: # Set environment variables os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key os.environ["GITHUB_TOKEN"] = github_token # Add the buttons after the inputs are provided if "use_sonnet" not in st.session_state: st.session_state.use_sonnet = False if "show_system_prompt" not in st.session_state: st.session_state.show_system_prompt = False col1, col2 = st.columns(2) with col1: if st.button("Show System Prompt" if not st.session_state.show_system_prompt else "Hide System Prompt"): st.session_state.show_system_prompt = not st.session_state.show_system_prompt with col2: if st.button("Use Sonnet 3.5"): st.session_state.use_sonnet = True if st.session_state.use_sonnet: sonnet_api_key = st.text_input("Input Anthropic API Key for Sonnet 3.5", type="password") if sonnet_api_key: os.environ["ANTHROPIC_API_KEY"] = sonnet_api_key # Parse the repository URL to extract user_name and REPO_NAME parsed_url = urlparse(github_repo_url) path_parts = parsed_url.path.strip('/').split('/') if len(path_parts) == 2: user_name, repo_name = path_parts else: st.error("Invalid GitHub repository URL. Please ensure it is in the format: https://github.com/user/repo") st.stop() REPO_URL = f"https://{github_token}@github.com/{user_name}/{repo_name}.git" headers = { 'Authorization': f'token {github_token}', 'Accept': 'application/vnd.github.v3+json', } def force_clone_repo(*args, **kwargs) -> str: if os.path.exists(repo_name): shutil.rmtree(repo_name) try: Repo.clone_from(REPO_URL, repo_name) return f"Repository {repo_name} forcefully cloned successfully." except GitCommandError as e: return f"Error cloning repository: {str(e)}" force_clone_tool = Tool( name="force_clone_repo", func=force_clone_repo, description="Forcefully clone the repository, removing any existing local copy." ) class WriteFileInput(BaseModel): file_path: str = Field(..., description="The path of the file to write to") content: str = Field(..., description="The content to write to the file") def write_file_content(file_path: str, content: str) -> str: full_path = os.path.join(repo_name, file_path) try: with open(full_path, 'w') as file: file.write(content) return f"Successfully wrote to {full_path}" except Exception as e: return f"Error writing to file: {str(e)}" file_write_tool = StructuredTool.from_function( func=write_file_content, name="write_file", description="Write content to a specific file in the repository.", args_schema=WriteFileInput ) def read_file_content(file_path: str) -> str: force_clone_repo() # Ensure we have the latest version before reading full_path = os.path.join(repo_name, file_path) try: with open(full_path, 'r') as file: content = file.read() return f"File content:\n{content}" except Exception as e: return f"Error reading file: {str(e)}" file_read_tool = Tool( name="read_file", func=read_file_content, description="Read content from a specific file in the repository." ) class CommitPushInput(BaseModel): commit_message: str = Field(..., description="The commit message") def commit_and_push(commit_message: str) -> str: try: repo = Repo(repo_name) repo.git.add(A=True) repo.index.commit(commit_message) origin = repo.remote(name='origin') push_info = origin.push() if push_info: if push_info[0].flags & push_info[0].ERROR: return f"Error pushing changes: {push_info[0].summary}" else: return f"Changes committed and pushed successfully with message: {commit_message}" else: return "No changes to push" except GitCommandError as e: return f"GitCommandError: {str(e)}" except Exception as e: return f"Unexpected error: {str(e)}" commit_push_tool = StructuredTool.from_function( func=commit_and_push, name="commit_and_push", description="Commit and push changes to the repository with a specific commit message.", args_schema=CommitPushInput ) tools = [force_clone_tool, file_read_tool, file_write_tool, commit_push_tool, ShellTool()] class AgentState(TypedDict): messages: Annotated[List[BaseMessage], operator.add] if st.session_state.use_sonnet and "ANTHROPIC_API_KEY" in os.environ: llm = ChatAnthropic(temperature=0, model_name="claude-3-5-sonnet-20240620") else: llm = ChatAnthropic(temperature=0, model_name="claude-3-haiku-20240307") # Modify the system prompts task_system_prompt_template = """You are an AI specialized in managing and analyzing a GitHub repository for a Next.js blog website. Your task is to answer user queries about the repository or execute tasks for modifying it. Before performing any operation, always use the force_clone_repo tool to ensure you have the latest version of the repository. Here is all of the code from the repository as well as the file paths for context of how the repo is structured: {REPO_CONTENT} Given this context, follow this prompt in completing the user's task: For user questions, provide direct answers based on the current state of the repository. For tasks given by the user, use the available tools and your knowledge of the repo to make necessary changes to the repository. When making changes, remember to force clone the repository first, make the changes, and then commit and push the changes. Available tools: 1. shell_tool: Execute shell commands 2. write_file: Write content to a specific file. Use as: write_file(file_path: str, content: str) 3. force_clone_repo: Forcefully clone the repository, removing any existing local copy 4. commit_and_push: Commit and push changes to the repository 5. read_file: Read content from a specific file in the repository When using the write_file tool, always provide both the file_path and the content as separate arguments. Respond to the human's messages and use tools when necessary to complete tasks. Take a deep breath and think through the task step by step:""" qa_system_prompt_template = """You are an AI specialized in analyzing a GitHub repository for a Next.js blog website. Your task is to answer user queries about the repository based on the provided content. Here is all of the code from the repository as well as the file paths for context of how the repo is structured: {REPO_CONTENT} Given this context, provide direct answers to user questions based on the current state of the repository. Take a deep breath and think through the question step by step before answering:""" memory = MemorySaver() def extract_repo_info(url): parts = url.split('/') if 'github.com' not in parts: raise ValueError("Not a valid GitHub URL") owner = parts[parts.index('github.com') + 1] repo = parts[parts.index('github.com') + 2] path_start_index = parts.index(repo) + 1 if path_start_index < len(parts) and parts[path_start_index] == 'tree': path_start_index += 2 path = '/'.join(parts[path_start_index:]) return owner, repo, path def get_repo_contents(owner, repo, path=''): api_url = f'https://api.github.com/repos/{owner}/{repo}/contents/{path}' response = requests.get(api_url, headers=headers) return response.json() def get_file_content_and_metadata(file_url): response = requests.get(file_url, headers=headers) content_data = response.json() content = content_data.get('content', '') if content: try: decoded_content = base64.b64decode(content) decoded_content_str = decoded_content.decode('utf-8') except (base64.binascii.Error, UnicodeDecodeError): decoded_content_str = content else: decoded_content_str = '' last_modified = content_data.get('last_modified') or response.headers.get('Last-Modified', '') return decoded_content_str, last_modified def is_valid_extension(filename): valid_extensions = ['.ipynb', '.py', '.js', '.md', '.mdx', 'tsx', 'ts', 'css', '.json'] return any(filename.endswith(ext) for ext in valid_extensions) def process_repo(repo_url): owner, repo, initial_path = extract_repo_info(repo_url) result = [] stack = [(initial_path, f'https://api.github.com/repos/{owner}/{repo}/contents/{initial_path}')] while stack: path, url = stack.pop() contents = get_repo_contents(owner, repo, path) if isinstance(contents, dict) and 'message' in contents: print(f"Error: {contents['message']}") return [] for item in contents: if item['type'] == 'file': if is_valid_extension(item['name']): file_url = item['url'] file_content, last_modified = get_file_content_and_metadata(file_url) if file_content: result.append({ 'url': item['html_url'], 'markdown': file_content, 'last_modified': last_modified }) elif item['type'] == 'dir': stack.append((item['path'], item['url'])) return result # Instead, add this block after the radio button for mode selection: if "task_system_prompt" not in st.session_state or "qa_system_prompt" not in st.session_state: st.session_state.task_system_prompt = task_system_prompt_template.format(REPO_CONTENT="") st.session_state.qa_system_prompt = qa_system_prompt_template.format(REPO_CONTENT="") # Modify the refresh_repo_data() function: def refresh_repo_data(): repo_contents = process_repo(github_repo_url) repo_contents_json = json.dumps(repo_contents, ensure_ascii=False, indent=2) st.session_state.REPO_CONTENT = repo_contents_json st.success("Repository content refreshed successfully.") # Update both system prompts with the new repo content st.session_state.task_system_prompt = task_system_prompt_template.format(REPO_CONTENT=st.session_state.REPO_CONTENT) st.session_state.qa_system_prompt = qa_system_prompt_template.format(REPO_CONTENT=st.session_state.REPO_CONTENT) # Recreate the graphs with the updated system prompts global task_graph, qa_graph if st.session_state.use_sonnet and "ANTHROPIC_API_KEY" in os.environ: new_llm = ChatAnthropic(temperature=0, model_name="claude-3-5-sonnet-20240620") else: new_llm = ChatAnthropic(temperature=0, model_name="claude-3-haiku-20240307") task_graph = create_react_agent( new_llm, tools=tools, messages_modifier=st.session_state.task_system_prompt, checkpointer=memory ) qa_graph = create_react_agent( new_llm, tools = graph_tools, messages_modifier=st.session_state.qa_system_prompt, checkpointer=memory ) if st.session_state.use_sonnet and "ANTHROPIC_API_KEY" in os.environ: refresh_repo_data() # Automatically refresh repo data when keys are provided if "REPO_CONTENT" not in st.session_state: refresh_repo_data() # Modify the code that displays the current system prompt: if st.session_state.show_system_prompt: current_prompt = st.session_state.task_system_prompt if mode == "Task" else st.session_state.qa_system_prompt st.text_area("Current System Prompt", current_prompt, height=300) # Update the graph initialization: if st.session_state.use_sonnet and "ANTHROPIC_API_KEY" in os.environ: llm = ChatAnthropic(temperature=0, model_name="claude-3-5-sonnet-20240620") else: llm = ChatAnthropic(temperature=0, model_name="claude-3-haiku-20240307") task_graph = create_react_agent( llm, tools=tools, messages_modifier=st.session_state.task_system_prompt, checkpointer=memory ) qa_graph = create_react_agent( llm, tools=graph_tools, messages_modifier=st.session_state.qa_system_prompt, checkpointer=memory ) def format_ai_response(response): # Remove custom code block formatting formatted_response = re.sub(r'```(.*?)```', r'```\1```', response, flags=re.DOTALL) # Remove custom inline code formatting formatted_response = re.sub(r'`([^`\n]+)`', r'`\1`', formatted_response) return formatted_response async def run_github_editor(query: str, thread_id: str = "default"): inputs = {"messages": [HumanMessage(content=query)]} config = { "configurable": {"thread_id": thread_id}, "recursion_limit": 50 } st.write(f"Human: {query}\n") full_response = "" response_container = st.empty() graph = task_graph if mode == "Task" else qa_graph async for event in graph.astream_events(inputs, config=config, version="v2"): kind = event["event"] if kind == "on_chat_model_start": response_container.write("AI is thinking...") elif kind == "on_chat_model_stream": data = event["data"] if data["chunk"].content: content = data["chunk"].content if isinstance(content, list) and content and isinstance(content[0], dict): text = content[0].get('text', '') full_response += text else: full_response += content response_container.markdown(format_ai_response(full_response)) elif kind == "on_tool_start" and mode == "Task": response_container.write(f"\nUsing tool: {event['name']}") elif kind == "on_tool_end" and mode == "Task": response_container.write(f"Tool result: {event['data']['output']}\n") # Update the final response using Streamlit's markdown response_container.markdown(format_ai_response(full_response)) # Create a session state variable to store the chat messages. This ensures that the # messages persist across reruns. if "messages" not in st.session_state: st.session_state.messages = [] # Display the current system prompt if show_system_prompt is True if st.session_state.show_system_prompt: st.text_area("Current System Prompt", st.session_state.system_prompt, height=300) # Display the existing chat messages via `st.chat_message`. for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # Create a chat input field to allow the user to enter a message. This will display # automatically at the bottom of the page. if prompt := st.chat_input(f"{'Ask a question' if mode == 'Q/A' else 'Give me a Task'}!"): # Store and display the current prompt. st.session_state.messages.append({"role": "user", "content": prompt}) with st.chat_message("user"): st.markdown(prompt) # Generate a response using the custom chatbot logic. asyncio.run(run_github_editor(prompt))