Spaces:
Sleeping
Sleeping
| import os | |
| import shutil | |
| import tempfile | |
| from flask import Flask, request, jsonify, render_template_string | |
| import git | |
| import json | |
| import google.generativeai as genai | |
| # --- IMPORTANT: Configure your API Key --- | |
| # For local testing, you can set this directly. | |
| # For Hugging Face Spaces, you will set this in "Secrets". | |
| # os.environ['GOOGLE_API_KEY'] = "YOUR_API_KEY_HERE" | |
| try: | |
| genai.configure(api_key=os.environ.get("GOOGLE_API_KEY")) | |
| except AttributeError: | |
| print("WARNING: GOOGLE_API_KEY secret not set. LLM functionality will fail.") | |
| app = Flask(__name__) | |
| # The HTML template remains the same | |
| HTML_TEMPLATE = """ | |
| <!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <title>AI README Generator 🧠</title> | |
| <style> | |
| @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap'); | |
| body { font-family: 'Inter', sans-serif; background-color: #f0f2f5; color: #1c1e21; margin: 0; display: flex; justify-content: center; align-items: center; min-height: 100vh; padding: 20px; box-sizing: border-box; } | |
| .container { background-color: #ffffff; padding: 40px 50px; border-radius: 12px; box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1); width: 100%; max-width: 700px; text-align: center; } | |
| h1 { font-size: 2.2em; color: #1877f2; margin-bottom: 10px; } | |
| p { color: #606770; font-size: 1.1em; margin-bottom: 30px; } | |
| .input-group { display: flex; margin-bottom: 20px; } | |
| #repo-url { flex-grow: 1; padding: 15px; border: 1px solid #dddfe2; border-radius: 6px 0 0 6px; font-size: 1em; outline: none; min-width: 0; } | |
| #repo-url:focus { border-color: #1877f2; box-shadow: 0 0 0 2px rgba(24, 119, 242, 0.2); } | |
| button { padding: 15px 25px; border: none; background-color: #1877f2; color: white; font-size: 1em; font-weight: 600; border-radius: 0 6px 6px 0; cursor: pointer; transition: background-color 0.3s; } | |
| button:hover { background-color: #166fe5; } | |
| .loader { border: 4px solid #f3f3f3; border-top: 4px solid #1877f2; border-radius: 50%; width: 40px; height: 40px; animation: spin 1s linear infinite; margin: 30px auto; display: none; } | |
| @keyframes spin { 0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); } } | |
| #result-container { display: none; margin-top: 20px; } | |
| #result { margin-top: 10px; padding: 20px; background-color: #f7f7f7; border: 1px solid #dddfe2; border-radius: 6px; text-align: left; white-space: pre-wrap; font-family: 'Courier New', Courier, monospace; max-height: 400px; overflow-y: auto; word-wrap: break-word; } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="container"> | |
| <h1>AI README Generator 🧠</h1> | |
| <p>Enter a public GitHub repository URL and let a true AI agent analyze the code and generate a README for you.</p> | |
| <form id="repo-form"><div class="input-group"><input type="url" id="repo-url" placeholder="e.g., https://github.com/user/project" required><button type="submit">Generate</button></div></form> | |
| <div class="loader" id="loader"></div> | |
| <div id="result-container"><h2>Generated README.md:</h2><pre id="result"></pre></div> | |
| </div> | |
| <script> | |
| document.getElementById('repo-form').addEventListener('submit', async function(event) { | |
| event.preventDefault(); | |
| const url = document.getElementById('repo-url').value; | |
| const loader = document.getElementById('loader'); | |
| const resultContainer = document.getElementById('result-container'); | |
| const resultDiv = document.getElementById('result'); | |
| loader.style.display = 'block'; | |
| resultContainer.style.display = 'none'; | |
| resultDiv.textContent = ''; | |
| try { | |
| const response = await fetch('/generate', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ url: url }) }); | |
| const data = await response.json(); | |
| if (response.ok) { resultDiv.textContent = data.readme; } else { resultDiv.textContent = 'Error: ' + data.error; } | |
| resultContainer.style.display = 'block'; | |
| } catch (error) { | |
| resultDiv.textContent = 'An unexpected error occurred: ' + error.toString(); | |
| resultContainer.style.display = 'block'; | |
| } finally { | |
| loader.style.display = 'none'; | |
| } | |
| }); | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| def generate_readme_with_llm(repo_path): | |
| """ | |
| Analyzes the repo and uses an LLM to generate the README content. | |
| """ | |
| # 1. Gather context: file structure and content of key files | |
| file_structure = "" | |
| file_contents = "" | |
| file_limit = 5 # Limit number of files to read to save tokens | |
| for root, _, files in os.walk(repo_path): | |
| # Ignore git directory | |
| if '.git' in root: | |
| continue | |
| level = root.replace(repo_path, '').count(os.sep) | |
| indent = ' ' * 4 * level | |
| file_structure += f"{indent}{os.path.basename(root)}/\n" | |
| sub_indent = ' ' * 4 * (level + 1) | |
| for f in files[:file_limit]: | |
| file_structure += f"{sub_indent}{f}\n" | |
| # Read content of a few key files | |
| try: | |
| with open(os.path.join(root, f), 'r', errors='ignore') as file: | |
| content = file.read(2000) # Read first 2000 chars | |
| file_contents += f"\n--- Start of {f} ---\n{content}\n--- End of {f} ---\n" | |
| except Exception: | |
| continue # Skip files we can't read | |
| # 2. Create a detailed prompt for the LLM | |
| prompt = f""" | |
| You are an expert technical writer tasked with creating a high-quality README.md for a GitHub repository. | |
| Analyze the following repository context and generate a comprehensive and user-friendly README. | |
| **Repository Context:** | |
| **1. File Structure:** | |
| ``` | |
| {file_structure} | |
| ``` | |
| **2. Content of Key Files:** | |
| ``` | |
| {file_contents} | |
| ``` | |
| **Instructions:** | |
| Based on the context provided, please generate a README.md file with the following sections: | |
| 1. **Project Title:** An appropriate title for the project. | |
| 2. **About the Project:** A concise but informative paragraph describing the project's purpose and the technology it uses. Infer this from the file contents. | |
| 3. **Getting Started:** Clear, step-by-step instructions for installation and setup. Infer the package manager (pip, npm, etc.) and run commands. | |
| 4. **Usage:** A simple example of how to run the application. | |
| **Rules:** | |
| - The output must be in valid Markdown format. | |
| - Be professional and clear. | |
| - If you cannot determine a piece of information (like a specific run command), suggest a common default and state that it might need to be verified. | |
| - Do not invent features that are not evident from the provided context. | |
| """ | |
| # 3. Call the LLM | |
| print("Sending request to Gemini API...") | |
| model = genai.GenerativeModel('gemini-1.5-flash-latest') | |
| response = model.generate_content(prompt) | |
| # Clean up the response - sometimes the model wraps it in ```markdown | |
| readme_text = response.text.strip() | |
| if readme_text.startswith("```markdown"): | |
| readme_text = readme_text[10:] | |
| if readme_text.endswith("```"): | |
| readme_text = readme_text[:-3] | |
| return readme_text | |
| def index(): | |
| return render_template_string(HTML_TEMPLATE) | |
| def generate(): | |
| repo_url = request.get_json().get('url') | |
| if not repo_url or not repo_url.startswith("[https://github.com/](https://github.com/)"): | |
| return jsonify({"error": "A valid GitHub repository URL is required."}), 400 | |
| if not os.environ.get("GOOGLE_API_KEY"): | |
| return jsonify({"error": "Server is missing the GOOGLE_API_KEY. Cannot contact the LLM."}), 500 | |
| temp_dir = tempfile.mkdtemp() | |
| try: | |
| git.Repo.clone_from(repo_url, temp_dir) | |
| cloned_folders = [d for d in os.listdir(temp_dir) if os.path.isdir(os.path.join(temp_dir, d))] | |
| if not cloned_folders: | |
| raise Exception("Cloning seems to have failed, the directory is empty.") | |
| repo_root_path = os.path.join(temp_dir, cloned_folders[0]) | |
| # --- NEW: Call the LLM-powered function --- | |
| readme_content = generate_readme_with_llm(repo_root_path) | |
| return jsonify({"readme": readme_content}) | |
| except Exception as e: | |
| return jsonify({"error": f"An error occurred: {str(e)}"}), 500 | |
| finally: | |
| shutil.rmtree(temp_dir) | |
| if __name__ == '__main__': | |
| app.run(host='0.0.0.0', port=int(os.environ.get("PORT", 7860))) |