From b4848ea5588dcb800a4648c11ac5f2db6de5ee7e Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 24 Apr 2026 17:25:43 +0200 Subject: [PATCH] fix: improve GitHub recipe support --- pyproject.toml | 1 + .../pyproject.toml | 1 + .../data_designer_github_repo_seed/scraper.py | 21 +++++--- .../scraper_impl/__init__.py | 2 - .../scraper_impl/gh_client.py | 21 +++++--- .../scraper_impl/scraper.py | 15 ++++-- .../single-env/data-designer-deps.txt | 3 +- .../dialogs/seed/seed-dialog.tsx | 37 +++++--------- .../recipe-studio/stores/recipe-studio.ts | 10 +++- .../utils/payload/builders-seed.ts | 38 ++++++++++++-- .../recipe-studio/utils/validation.ts | 51 ++++++++++++++++--- studio/install_python_stack.py | 39 ++++++++------ 12 files changed, 162 insertions(+), 77 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 815c6ee119..5687ea12f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ studio = [ "frontend/*.yaml", "frontend/.git*", "backend/requirements/**/*", + "backend/plugins/**/*", "backend/core/data_recipe/oxc-validator/*.json", "backend/core/data_recipe/oxc-validator/*.mjs", ] diff --git a/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml index f7f7be3001..e232adc60c 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml +++ b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml @@ -11,6 +11,7 @@ version = "0.1.0" description = "Unsloth Studio seed plugin that scrapes GitHub issues, PRs, and commits." requires-python = ">=3.11" dependencies = [ + "data-designer-engine>=0.5.4,<0.6", "requests>=2.31", ] diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py index 5ec32739e0..6aaeefbe00 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py @@ -16,11 +16,11 @@ import json import os import sys import time +import uuid from dataclasses import dataclass from pathlib import Path -# The scraper_impl gh_client reads GH_TOKEN at import time and raises if -# missing, so we defer the imports until `scrape()` runs with a resolved token. +# Defer scraper_impl imports until `scrape()` runs with a resolved token. _IMPL_DIR = Path(__file__).parent / "scraper_impl" @@ -154,9 +154,8 @@ def _flatten_commit_row(r: dict, repo: str) -> dict: def scrape(cfg: ScrapeConfig, base_dir: Path): token = _resolve_token(cfg.token) - os.environ["GH_TOKEN"] = token GitHubClient, RepoScraper = _load_impl() - client = GitHubClient() + client = GitHubClient(token = token) base_dir.mkdir(parents = True, exist_ok = True) # Per-resource trial limits. limit <= 0 means "all": use a very large cap. @@ -180,13 +179,18 @@ def scrape(cfg: ScrapeConfig, base_dir: Path): trial_limits = trial_limits, ) try: - scraper.scrape_repo_meta() + repo_meta = scraper.scrape_repo_meta() if "issues" in cfg.item_types: scraper.scrape_issues() if "pulls" in cfg.item_types: scraper.scrape_prs() if "commits" in cfg.item_types: - scraper.scrape_commits() + default_ref = repo_meta.get("defaultBranchRef") or {} + default_branch = ( + default_ref.get("name") if isinstance(default_ref, dict) else None + ) + branch = f"refs/heads/{default_branch}" if default_branch else "refs/heads/main" + scraper.scrape_commits(branch = branch) finally: scraper.close() @@ -217,9 +221,10 @@ def materialize_to_jsonl(cfg: ScrapeConfig, out_dir: Path) -> Path: out_dir.mkdir(parents = True, exist_ok = True) tag = "-".join(r.replace("/", "__") for r in cfg.repos)[:120] kinds = "-".join(cfg.item_types) - fname = f"github_{tag}__{kinds}__{cfg.limit}_{int(time.time())}.jsonl" + run_id = f"{int(time.time())}-{uuid.uuid4().hex[:12]}" + fname = f"github_{tag}__{kinds}__{cfg.limit}_{run_id}.jsonl" out = out_dir / fname - rows = scrape(cfg, out_dir / "raw") + rows = scrape(cfg, out_dir / "raw-runs" / run_id) with out.open("w", encoding = "utf-8") as f: for r in rows: f.write(json.dumps(r, ensure_ascii = False) + "\n") diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py index 86918c4135..32014236c6 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py @@ -1,4 +1,2 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 - - diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py index bd975b3f34..dd2de2f5ce 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py @@ -15,15 +15,10 @@ import requests log = logging.getLogger("gh_client") -GH_TOKEN = os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN") -if not GH_TOKEN: - raise RuntimeError("GH_TOKEN not set in environment") - GRAPHQL_URL = "https://api.github.com/graphql" REST_BASE = "https://api.github.com" -DEFAULT_HEADERS = { - "Authorization": f"Bearer {GH_TOKEN}", +BASE_HEADERS = { "Accept": "application/vnd.github+json", "X-GitHub-Api-Version": "2022-11-28", "User-Agent": "github-data-gatherer/1.0", @@ -35,9 +30,19 @@ class RateLimitError(Exception): class GitHubClient: - def __init__(self, min_remaining_graphql: int = 100, min_remaining_rest: int = 100): + def __init__( + self, + min_remaining_graphql: int = 100, + min_remaining_rest: int = 100, + token: str | None = None, + ): + token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN") + if not token: + raise RuntimeError("GH_TOKEN not set in environment") self.session = requests.Session() - self.session.headers.update(DEFAULT_HEADERS) + self.session.headers.update( + {**BASE_HEADERS, "Authorization": f"Bearer {token}"} + ) self.min_remaining_graphql = min_remaining_graphql self.min_remaining_rest = min_remaining_rest self.graphql_remaining: Optional[int] = None diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py index 3695ff1641..eaeff47da2 100644 --- a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py @@ -96,7 +96,7 @@ class RepoScraper: ) # ----- repo meta ----- - def scrape_repo_meta(self) -> None: + def scrape_repo_meta(self) -> Dict[str, Any]: data = self.client.graphql( Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name} ) @@ -104,6 +104,7 @@ class RepoScraper: repo = data.get("data", {}).get("repository") or {} repo["_fetchedAt"] = ts() self.writers["repo_meta"].write(repo) + return repo # ----- issues ----- def scrape_issues(self) -> int: @@ -674,8 +675,9 @@ def main(): owner, name = repo_spec.split("/") scraper = RepoScraper(owner, name, data_dir, client, trial_limits) try: - if not only or "meta" in only: - scraper.scrape_repo_meta() + repo_meta: Dict[str, Any] = {} + if not only or "meta" in only or "commits" in only: + repo_meta = scraper.scrape_repo_meta() if not only or "labels" in only: scraper.scrape_labels() if not only or "milestones" in only: @@ -689,7 +691,12 @@ def main(): if not only or "pulls" in only: scraper.scrape_prs() if not only or "commits" in only: - scraper.scrape_commits() + default_ref = repo_meta.get("defaultBranchRef") or {} + default_branch = ( + default_ref.get("name") if isinstance(default_ref, dict) else None + ) + branch = f"refs/heads/{default_branch}" if default_branch else "refs/heads/main" + scraper.scrape_commits(branch = branch) finally: scraper.close() finally: diff --git a/studio/backend/requirements/single-env/data-designer-deps.txt b/studio/backend/requirements/single-env/data-designer-deps.txt index fc63230922..f63c076621 100644 --- a/studio/backend/requirements/single-env/data-designer-deps.txt +++ b/studio/backend/requirements/single-env/data-designer-deps.txt @@ -19,7 +19,8 @@ ruff<1,>=0.14.10 scipy<2,>=1.11.0 sqlfluff<4,>=3.2.0 tiktoken<1,>=0.8.0 -# Unstructured-seed plugin deps (plugin installed with --no-deps) +# Local seed plugin deps (plugins installed with --no-deps) +requests>=2.31 pymupdf>=1.24.0 pymupdf4llm>=0.0.17 mammoth>=1.8.0 diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 127f81c4e0..6da2b2fb64 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -84,13 +84,12 @@ function GithubRepoSeedForm({ }): ReactElement { const [advancedOpen, setAdvancedOpen] = useState(false); const limitStr = (config.github_limit ?? "100").trim(); - const allMode = limitStr === "" || limitStr === "0"; return (