Create app.py

app.py

@@ -0,0 +1,352 @@
+import os
+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
+
+# Initialize show_system_prompt in session state
+if "show_system_prompt" not in st.session_state:
+    st.session_state.show_system_prompt = False
+
+# Add a button to toggle the visibility of the system prompt
+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
+
+# Show title and description.
+st.title("Coder for NextJS Templates")
+st.write(
+    "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!"
+)
+
+# 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"
+
+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/[email protected]/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")
+
+if not (github_repo_url and github_token and anthropic_api_key):
+    st.info("Please add your Github Repo URL, Github Token, and Anthropic API key to continue.", icon="🗝️")
+else:
+    # Set environment variables
+    os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
+    os.environ["GITHUB_TOKEN"] = github_token
+
+    # 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]
+
+    llm = ChatAnthropic(temperature=0, model_name="claude-3-haiku-20240307")
+
+    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:"""
+
+    from langgraph.checkpoint import MemorySaver
+
+    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
+
+    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 the system prompt with the new repo content
+        st.session_state.system_prompt = system_prompt_template.format(REPO_CONTENT=st.session_state.REPO_CONTENT)
+
+        # Recreate the graph with the updated system prompt
+        global graph
+        graph = create_react_agent(
+            llm,
+            tools=tools,
+            messages_modifier=st.session_state.system_prompt,
+            checkpointer=memory
+        )
+
+    # Automatically refresh repo data when keys are provided
+    if "REPO_CONTENT" not in st.session_state:
+        refresh_repo_data()
+
+    # Initialize system_prompt in session state
+    if "system_prompt" not in st.session_state:
+        st.session_state.system_prompt = system_prompt_template.format(REPO_CONTENT=st.session_state.REPO_CONTENT)
+
+    graph = create_react_agent(
+        llm,
+        tools=tools,
+        messages_modifier=st.session_state.system_prompt,
+        checkpointer=memory
+    )
+
+    from langchain_core.messages import AIMessage, ToolMessage
+
+    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  # Add this line to set the recursion limit
+        }
+
+        st.write(f"Human: {query}\n")
+
+        current_thought = ""
+
+        async for event in graph.astream_events(inputs, config=config, version="v2"):
+            kind = event["event"]
+            if kind == "on_chat_model_start":
+                st.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', '')
+                        current_thought += text
+                        if text.endswith(('.', '?', '!')):
+                            st.write(current_thought.strip())
+                            current_thought = ""
+                    else:
+                        st.write(content, end="")
+            elif kind == "on_tool_start":
+                st.write(f"\nUsing tool: {event['name']}")
+            elif kind == "on_tool_end":
+                st.write(f"Tool result: {event['data']['output']}\n")
+
+    # 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("What is up?"):
+
+        # 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))