Create backend.py

backend.py

@@ -0,0 +1,372 @@
+import sqlite3
+import os
+import json
+import uuid
+import zipfile
+import tempfile
+import subprocess
+import re
+from huggingface_hub import whoami
+
+DB_PATH = "code_agents.db"
+PROJECT_ROOT = "./projects"
+os.makedirs(PROJECT_ROOT, exist_ok=True)
+
+def init_db():
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.executescript("""
+        CREATE TABLE IF NOT EXISTS users (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            hf_token TEXT UNIQUE NOT NULL,
+            username TEXT,
+            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
+        );
+        CREATE TABLE IF NOT EXISTS projects (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            user_id INTEGER NOT NULL,
+            title TEXT,
+            description TEXT,
+            zip_path TEXT,
+            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+            FOREIGN KEY (user_id) REFERENCES users(id)
+        );
+    """)
+    conn.commit()
+    conn.close()
+
+init_db()
+
+def get_user_by_token(hf_token):
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("SELECT id, username FROM users WHERE hf_token = ?", (hf_token,))
+    row = cursor.fetchone()
+    conn.close()
+    return row if row else None
+
+def create_user(hf_token, username=None):
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("INSERT INTO users (hf_token, username) VALUES (?, ?)", (hf_token, username))
+    user_id = cursor.lastrowid
+    conn.commit()
+    conn.close()
+    return user_id
+
+def create_project(user_id, title, description, zip_path):
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("""
+        INSERT INTO projects (user_id, title, description, zip_path)
+        VALUES (?, ?, ?, ?)
+    """, (user_id, title, description, zip_path))
+    project_id = cursor.lastrowid
+    conn.commit()
+    conn.close()
+    return project_id
+
+def get_user_projects(user_id):
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+    cursor.execute("""
+        SELECT id, title, description, zip_path, created_at
+        FROM projects WHERE user_id = ? ORDER BY created_at DESC
+    """, (user_id,))
+    projects = cursor.fetchall()
+    conn.close()
+    return projects
+
+def verify_hf_token(token):
+    try:
+        user_info = whoami(token=token)
+        username = user_info.get("name", "Anonymous")
+        user = get_user_by_token(token)
+        if not user:
+            user_id = create_user(token, username)
+            return user_id, f"Welcome to Code Agents, {username}! Your AI team is ready."
+        else:
+            return user[0], f"Welcome back, {user[1]}! Your projects are waiting."
+    except Exception as e:
+        return None, f"Invalid token. Please try again. ({str(e)})"
+
+def get_user_project_dir(user_id):
+    user_dir = os.path.join(PROJECT_ROOT, str(user_id))
+    os.makedirs(user_dir, exist_ok=True)
+    return user_dir
+
+def clear_user_project_dir(user_id):
+    user_dir = get_user_project_dir(user_id)
+    for f in os.listdir(user_dir):
+        path = os.path.join(user_dir, f)
+        if os.path.isdir(path):
+            import shutil
+            shutil.rmtree(path)
+        else:
+            os.remove(path)
+
+TOOLS = {}
+
+def create_file(user_id, path: str, content: str):
+    user_dir = get_user_project_dir(user_id)
+    full_path = os.path.join(user_dir, path.lstrip("/"))
+    os.makedirs(os.path.dirname(full_path), exist_ok=True)
+    with open(full_path, 'w') as f:
+        f.write(content)
+    return f"Created: {path}"
+
+def read_file(user_id, path: str):
+    user_dir = get_user_project_dir(user_id)
+    full_path = os.path.join(user_dir, path.lstrip("/"))
+    if not os.path.exists(full_path):
+        return f"File not found: {path}"
+    with open(full_path, 'r') as f:
+        return f.read()
+
+def list_files(user_id):
+    user_dir = get_user_project_dir(user_id)
+    files = []
+    for root, _, filenames in os.walk(user_dir):
+        for f in filenames:
+            files.append(os.path.relpath(os.path.join(root, f), user_dir))
+    return files
+
+def execute_code(user_id, code: str):
+    try:
+        user_dir = get_user_project_dir(user_id)
+        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False, dir=user_dir) as f:
+            f.write(code)
+            temp_path = f.name
+        result = subprocess.run(["python3", temp_path], capture_output=True, text=True, timeout=10)
+        os.unlink(temp_path)
+        if result.returncode == 0:
+            return result.stdout
+        else:
+            return f"Error: {result.stderr}"
+    except Exception as e:
+        return f"Timeout/Error: {str(e)}"
+
+def run_tests(user_id):
+    user_dir = get_user_project_dir(user_id)
+    test_dir = os.path.join(user_dir, "tests")
+    if not os.path.exists(test_dir):
+        return "No test directory found."
+    result = subprocess.run(["pytest", test_dir], capture_output=True, text=True, cwd=user_dir)
+    return result.stdout
+
+def lint_code(user_id):
+    user_dir = get_user_project_dir(user_id)
+    result = subprocess.run(["flake8", user_dir], capture_output=True, text=True)
+    return result.stdout if result.returncode != 0 else "No linting errors."
+
+def scan_vulns(user_id):
+    reqs = os.path.join(get_user_project_dir(user_id), "requirements.txt")
+    if not os.path.exists(reqs):
+        return "No requirements.txt found."
+    result = subprocess.run(["pip-audit", "-r", reqs], capture_output=True, text=True)
+    return result.stdout if result.returncode != 0 else "No known vulnerabilities found."
+
+def search_hf_docs(query: str):
+    return f"Searching Hugging Face docs for: '{query}' -> Found: https://huggingface.co/docs/transformers/main/en/quicktour"
+
+def write_readme(user_id, project_name: str, description: str):
+    user_dir = get_user_project_dir(user_id)
+    readme_content = f"""Project Name: {project_name}
+
+Description:
+{description}
+
+How to Run:
+1. Install dependencies: pip install -r requirements.txt
+2. Run the main file: python cli.py
+
+This project was built by Code Agents — an AI team on Hugging Face.
+"""
+    create_file(user_id, "README.md", readme_content)
+    return "README.md generated."
+
+def generate_dockerfile(user_id):
+    user_dir = get_user_project_dir(user_id)
+    dockerfile = """FROM python:3.10-slim
+COPY . /app
+WORKDIR /app
+RUN pip install --no-cache-dir -r requirements.txt
+CMD ["python", "cli.py"]
+"""
+    create_file(user_id, "Dockerfile", dockerfile)
+    return "Dockerfile generated."
+
+def zip_project(user_id, project_name: str):
+    user_dir = get_user_project_dir(user_id)
+    zip_path = os.path.join(user_dir, f"{project_name}.zip")
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
+        for root, _, files in os.walk(user_dir):
+            for f in files:
+                if f.endswith(".zip"): continue
+                full_path = os.path.join(root, f)
+                arcname = os.path.relpath(full_path, user_dir)
+                zf.write(full_path, arcname)
+    return zip_path
+
+def fix_error(error_msg: str, code: str):
+    if "ModuleNotFoundError" in error_msg and "yt_dlp" in error_msg:
+        return "Add 'yt-dlp' to requirements.txt"
+    elif "SyntaxError" in error_msg:
+        return "Fix indentation or missing colon"
+    return "Unable to auto-fix. Please review manually."
+
+TOOLS = {
+    "create_file": {"func": create_file, "desc": "Create a new file at specified path with content"},
+    "read_file": {"func": read_file, "desc": "Read content of a file"},
+    "list_files": {"func": list_files, "desc": "List all files in the project directory"},
+    "execute_code": {"func": execute_code, "desc": "Execute Python code and return stdout/stderr"},
+    "run_tests": {"func": run_tests, "desc": "Run pytest on test files in /tests directory"},
+    "lint_code": {"func": lint_code, "desc": "Run flake8 to check code style and errors"},
+    "scan_vulns": {"func": scan_vulns, "desc": "Scan dependencies for security vulnerabilities"},
+    "search_hf_docs": {"func": search_hf_docs, "desc": "Search Hugging Face documentation for model usage"},
+    "write_readme": {"func": write_readme, "desc": "Generate a README.md for the project"},
+    "generate_dockerfile": {"func": generate_dockerfile, "desc": "Generate a Dockerfile for containerization"},
+    "zip_project": {"func": zip_project, "desc": "Compress entire project into a downloadable ZIP"},
+    "fix_error": {"func": fix_error, "desc": "Given an error message and code, suggest a fix"}
+}
+
+AGENT_PROMPTS = {
+    "architect": """
+You are an expert software architect. Your job is to plan the project structure.
+Given a user request, decide:
+- File/folder structure
+- Required libraries
+- Tech stack (e.g., Flask, CLI, etc.)
+
+Respond only in JSON format:
+{
+  "plan": "Brief plan",
+  "files": [
+    {"path": "cli.py", "content": "// initial code"},
+    {"path": "requirements.txt", "content": "yt-dlp\\nflask"}
+  ]
+}
+""",
+    "coder": """
+You are a senior Python developer. Write clean, efficient, well-commented code.
+Use the tools available: create_file, read_file, execute_code, search_hf_docs.
+Never guess — use tools to verify.
+When done, write the file using create_file().
+""",
+    "tester": """
+You are a QA engineer. Write unit tests for all functions.
+Use pytest. Create test files in /tests/.
+Use execute_code() to run tests and verify output.
+If tests fail, ask coder to fix.
+""",
+    "reviewer": """
+You are a security and code review expert.
+Check for:
+- Security flaws (e.g., eval(), shell injection)
+- Poor practices
+- Missing error handling
+- Unused imports
+Use lint_code(), scan_vulns(), read_file()
+Suggest fixes using fix_error(). Update files with create_file().
+""",
+    "deployer": """
+You are a DevOps engineer.
+Generate:
+- README.md with usage instructions
+- Dockerfile
+- ZIP archive of project
+Use: write_readme(), generate_dockerfile(), zip_project()
+Return final ZIP path when done.
+"""
+}
+
+class Agent:
+    def __init__(self, name: str, prompt: str, tools: dict, user_id: int):
+        self.name = name
+        self.prompt = prompt
+        self.tools = tools
+        self.user_id = user_id
+        self.memory = []
+
+    def think(self, task: str) -> str:
+        return f"[{self.name}] {task} (simulated)"
+
+    def act(self, action: str) -> str:
+        try:
+            if "(" not in action or ")" not in action:
+                return "Invalid action format"
+            func_name = action.split("(")[0]
+            args_str = action[len(func_name)+1:-1]
+            args = {}
+            if args_str:
+                for pair in args_str.split(","):
+                    k, v = pair.split("=")
+                    args[k.strip()] = v.strip().strip('"\'')
+            if func_name in self.tools:
+                bound_func = lambda **kwargs: self.tools[func_name]["func"](self.user_id, **kwargs)
+                return bound_func(**args)
+            return f"Unknown tool: {func_name}"
+        except Exception as e:
+            return f"Error executing {action}: {str(e)}"
+
+    def run(self, task: str) -> str:
+        response = self.think(task)
+        self.memory.append(f"{self.name}: {response}")
+        
+        actions = re.findall(r'(\w+\(.*?\))', response)
+        for action in actions:
+            result = self.act(action)
+            self.memory.append(f"{self.name} → {action} → {result}")
+        
+        return "\n".join(self.memory[-5:])
+
+def run_code_agents(user_id: int, project_description: str):
+    clear_user_project_dir(user_id)
+    user_dir = get_user_project_dir(user_id)
+
+    agents = {
+        "architect": Agent("Architect", AGENT_PROMPTS["architect"], TOOLS, user_id),
+        "coder": Agent("Coder", AGENT_PROMPTS["coder"], TOOLS, user_id),
+        "tester": Agent("Tester", AGENT_PROMPTS["tester"], TOOLS, user_id),
+        "reviewer": Agent("Reviewer", AGENT_PROMPTS["reviewer"], TOOLS, user_id),
+        "deployer": Agent("Deployer", AGENT_PROMPTS["deployer"], TOOLS, user_id),
+    }
+
+    timeline = []
+    logs = ""
+
+    logs += "Architect planning...\n"
+    plan = agents["architect"].run(f"Plan a project: {project_description}")
+    timeline.append({"agent": "architect", "step": "plan", "output": plan})
+    logs += plan + "\n\n"
+
+    logs += "Coder implementing...\n"
+    coder_task = f"Implement the project based on this plan:\n{plan}"
+    coder_result = agents["coder"].run(coder_task)
+    timeline.append({"agent": "coder", "step": "implement", "output": coder_result})
+    logs += coder_result + "\n\n"
+
+    logs += "Tester adding tests...\n"
+    test_task = "Write tests for the implemented code. Save to /tests/test_main.py"
+    test_result = agents["tester"].run(test_task)
+    timeline.append({"agent": "tester", "step": "test", "output": test_result})
+    logs += test_result + "\n\n"
+
+    logs += "Reviewer checking code...\n"
+    review_task = "Review all files. Fix any bugs or security issues."
+    review_result = agents["reviewer"].run(review_task)
+    timeline.append({"agent": "reviewer", "step": "review", "output": review_result})
+    logs += review_result + "\n\n"
+
+    logs += "Deployer packaging...\n"
+    deploy_task = f"Package the project named '{project_description[:20]}...' into a ZIP with README and Dockerfile."
+    deploy_result = agents["deployer"].run(deploy_task)
+    timeline.append({"agent": "deployer", "step": "deploy", "output": deploy_result})
+
+    zip_files = [f for f in os.listdir(user_dir) if f.endswith(".zip")]
+    zip_path = os.path.join(user_dir, zip_files[0]) if zip_files else None
+
+    if zip_path:
+        project_id = create_project(user_id, project_description[:50], project_description, zip_path)
+
+    return logs, zip_path, timeline