github_loader.py

import re
import json
import requests
from typing import List, Tuple, Optional
from langchain_core.documents import Document

# Supported code & config extensions to analyze
SUPPORTED_EXTENSIONS = {
    ".py", ".ipynb", ".js", ".jsx", ".ts", ".tsx", ".java", ".cpp", ".c", ".h", ".hpp",
    ".cs", ".go", ".rs", ".rb", ".php", ".swift", ".kt", ".scala", ".sql",
    ".html", ".css", ".scss", ".md", ".json", ".yaml", ".yml", ".toml",
    ".xml", ".sh", ".bat", "Dockerfile", "Makefile"
}

# Directories and specific lock/minified files to ignore
IGNORED_DIRS = {
    "node_modules", "venv", ".venv", "env", ".env", ".git", "__pycache__",
    "dist", "build", "target", "out", "bin", "obj", ".idea", ".vscode",
    "coverage", ".next", ".nuxt", "vendor", "migrations"
}

IGNORED_FILES = {
    "package-lock.json", "yarn.lock", "pnpm-lock.yaml", "poetry.lock",
    "Cargo.lock", "composer.lock", "Gemfile.lock"
}


def parse_github_url(url: str) -> Tuple[str, str]:
    """Extracts the (owner, repo) tuple from a public GitHub URL."""
    pattern = r"github\.com/([^/]+)/([^/]+)"
    match = re.search(pattern, url.strip().rstrip("/"))
    if not match:
        raise ValueError("Invalid GitHub repository URL. Format: https://github.com/owner/repo")
    owner, repo = match.group(1), match.group(2)
    if repo.endswith(".git"):
        repo = repo[:-4]
    return owner, repo


def should_process_file(file_path: str, size: int, max_size_kb: int = 2000) -> bool:
    """Filters out binary files, massive auto-generated files, and ignored directories."""
    if size > max_size_kb * 1024:
        return False

    parts = file_path.split("/")
    if any(part in IGNORED_DIRS for part in parts):
        return False

    filename = parts[-1]
    if filename in IGNORED_FILES or filename.endswith(".min.js") or filename.endswith(".min.css"):
        return False

    if filename in {"Dockerfile", "Makefile", "requirements.txt", "package.json", "pom.xml", "go.mod", "Cargo.toml"}:
        return True

    _, ext = ("", filename) if "." not in filename else (filename[:filename.rfind(".")], filename[filename.rfind("."):])
    return ext.lower() in SUPPORTED_EXTENSIONS

def extract_notebook_content(raw_content: str) -> str:
    """Convert a Jupyter notebook's JSON into readable Markdown and code."""
    try:
        notebook = json.loads(raw_content)
    except json.JSONDecodeError:
        return raw_content

    sections = []

    for cell in notebook.get("cells", []):
        cell_type = cell.get("cell_type", "unknown")
        source = cell.get("source", [])

        if isinstance(source, list):
            source = "".join(source)

        if not source.strip():
            continue

        if cell_type == "markdown":
            sections.append(f"## Markdown Cell\n{source}")
        elif cell_type == "code":
            sections.append(f"## Code Cell\n```python\n{source}\n```")

    return "\n\n".join(sections)

def fetch_github_repo(
    repo_url: str,
    github_token: Optional[str] = None,
    max_files: int = 120,
    progress_callback=None
) -> Tuple[List[Document], str]:
    """
    Fetches repository files using the GitHub Git Trees API and raw content endpoints.
    Returns (list_of_langchain_documents, repo_tree_summary_string).
    """
    owner, repo = parse_github_url(repo_url)
    headers = {"Accept": "application/vnd.github.v3+json"}
    if github_token and github_token.strip():
        headers["Authorization"] = f"token {github_token.strip()}"

    # 1. Fetch repository metadata to identify the default branch
    repo_api_url = f"https://api.github.com/repos/{owner}/{repo}"
    repo_resp = requests.get(repo_api_url, headers=headers, timeout=15)
    if repo_resp.status_code == 404:
        raise ValueError(f"Repository '{owner}/{repo}' not found or is private.")
    elif repo_resp.status_code == 403:
        raise RuntimeError("GitHub API rate limit exceeded. Add a GitHub Personal Access Token in the sidebar.")
    repo_resp.raise_for_status()
    default_branch = repo_resp.json().get("default_branch", "main")

    # 2. Fetch the entire recursive repository tree in 1 API call
    tree_api_url = f"https://api.github.com/repos/{owner}/{repo}/git/trees/{default_branch}?recursive=1"
    tree_resp = requests.get(tree_api_url, headers=headers, timeout=15)
    tree_resp.raise_for_status()
    tree_data = tree_resp.json().get("tree", [])

    # Filter eligible files
    eligible_files = [
        item for item in tree_data
        if item.get("type") == "blob" and should_process_file(item.get("path", ""), item.get("size", 0))
    ]

    # Prioritize root configuration/manifest files & READMEs first
    def file_priority(item):
        path = item["path"]
        depth = path.count("/")
        is_priority = any(k in path.lower() for k in ["readme", "package.json", "requirements", "dockerfile", "cargo.toml", "go.mod", "pom.xml"])
        return (0 if is_priority else 1, depth, path)

    eligible_files.sort(key=file_priority)
    selected_files = eligible_files[:max_files]

    # Build a clean directory tree string so the LLM understands the global repo structure
    all_paths = [item["path"] for item in eligible_files]
    repo_tree_str = f"Repository: {owner}/{repo} (Branch: {default_branch})\nTotal indexed files: {len(selected_files)}\nFile paths:\n" + "\n".join(f"- {p}" for p in all_paths[:300])

    documents: List[Document] = []
    # Add a special architectural overview document containing the file tree
    documents.append(
        Document(
            page_content=repo_tree_str,
            metadata={"file_path": "REPO_ARCHITECTURE_TREE", "repo": f"{owner}/{repo}", "language": "text"}
        )
    )

    # 3. Download raw file contents via raw.githubusercontent.com (does not consume REST API quota)
    total = len(selected_files)
    for idx, item in enumerate(selected_files):
        file_path = item["path"]
        if progress_callback:
            progress_callback(idx + 1, total, file_path)

        raw_url = f"https://raw.githubusercontent.com/{owner}/{repo}/{default_branch}/{file_path}"
        try:
            raw_resp = requests.get(raw_url, headers=headers, timeout=10)
            if raw_resp.status_code == 200:
                content = raw_resp.text
                if file_path.lower().endswith(".ipynb"):
                    content = extract_notebook_content(content)

                if not content.strip():
                    continue
                ext = file_path[file_path.rfind("."):].lower() if "." in file_path else ""
                # Prepend file path header directly into the content so every chunk retains context
                formatted_content = f"### File: {file_path}\n\n{content}"
                documents.append(
                    Document(
                        page_content=formatted_content,
                        metadata={
                            "file_path": file_path,
                            "repo": f"{owner}/{repo}",
                            "extension": ext,
                            "url": f"https://github.com/{owner}/{repo}/blob/{default_branch}/{file_path}"
                        }
                    )
                )
        except requests.RequestException:
            continue

    return documents, repo_tree_str