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/auth/storage.py b/studio/backend/auth/storage.py index 2b0e359d39..9a03f5f542 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -146,10 +146,18 @@ def get_connection() -> sqlite3.Connection: created_at TEXT NOT NULL, last_used_at TEXT, expires_at TEXT, - is_active INTEGER NOT NULL DEFAULT 1 + is_active INTEGER NOT NULL DEFAULT 1, + is_internal INTEGER NOT NULL DEFAULT 0 ); """ ) + api_key_columns = { + row["name"] for row in conn.execute("PRAGMA table_info(api_keys)") + } + if "is_internal" not in api_key_columns: + conn.execute( + "ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0" + ) conn.execute( """ CREATE TABLE IF NOT EXISTS app_secrets ( @@ -592,11 +600,15 @@ def create_api_key( username: str, name: str, expires_at: Optional[str] = None, + internal: bool = False, ) -> Tuple[str, dict]: """Create a new API key for *username*. Returns ``(raw_key, row_dict)`` where *raw_key* is shown to the user - exactly once. The database only stores the SHA-256 hash. + exactly once. The database only stores the PBKDF2 hash. + + Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe + runs) that should not appear in user-facing key listings. """ raw_key = API_KEY_PREFIX + secrets.token_hex(16) key_hash = _pbkdf2_api_key(raw_key) @@ -607,10 +619,18 @@ def create_api_key( try: conn.execute( """ - INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at) - VALUES (?, ?, ?, ?, ?, ?) + INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal) + VALUES (?, ?, ?, ?, ?, ?, ?) """, - (username, key_prefix, key_hash, name, now, expires_at), + ( + username, + key_prefix, + key_hash, + name, + now, + expires_at, + 1 if internal else 0, + ), ) conn.commit() cur = conn.execute("SELECT * FROM api_keys WHERE key_hash = ?", (key_hash,)) @@ -620,19 +640,33 @@ def create_api_key( conn.close() -def list_api_keys(username: str) -> list: - """Return all API keys for *username* (never exposes ``key_hash``).""" +def list_api_keys(username: str, include_internal: bool = False) -> list: + """Return API keys for *username*. Internal workflow keys are hidden + by default so they do not clutter user-facing UIs.""" conn = get_connection() try: - cur = conn.execute( - """ - SELECT id, username, key_prefix, name, created_at, last_used_at, expires_at, is_active - FROM api_keys - WHERE username = ? - ORDER BY created_at DESC - """, - (username,), - ) + if include_internal: + cur = conn.execute( + """ + SELECT id, username, key_prefix, name, created_at, last_used_at, + expires_at, is_active, is_internal + FROM api_keys + WHERE username = ? + ORDER BY created_at DESC + """, + (username,), + ) + else: + cur = conn.execute( + """ + SELECT id, username, key_prefix, name, created_at, last_used_at, + expires_at, is_active, is_internal + FROM api_keys + WHERE username = ? AND is_internal = 0 + ORDER BY created_at DESC + """, + (username,), + ) return [dict(row) for row in cur.fetchall()] finally: conn.close() @@ -652,6 +686,24 @@ def revoke_api_key(username: str, key_id: int) -> bool: conn.close() +def revoke_internal_api_key(key_id: int) -> bool: + """Revoke an internal workflow-minted key without requiring a username. + + Used by the recipe runner to retire its sk-unsloth-* key once the job + terminates, shrinking the window a leaked key could be abused. + """ + conn = get_connection() + try: + cursor = conn.execute( + "UPDATE api_keys SET is_active = 0 WHERE id = ? AND is_internal = 1", + (key_id,), + ) + conn.commit() + return cursor.rowcount > 0 + finally: + conn.close() + + def validate_api_key(raw_key: str) -> Optional[str]: """Validate *raw_key* and return the owning username, or ``None``. diff --git a/studio/backend/core/data_recipe/jobs/constants.py b/studio/backend/core/data_recipe/jobs/constants.py index 08237326f8..0045276e20 100644 --- a/studio/backend/core/data_recipe/jobs/constants.py +++ b/studio/backend/core/data_recipe/jobs/constants.py @@ -9,6 +9,7 @@ STAGE_PREVIEW = "preview" STAGE_DAG = "dag" STAGE_HEALTHCHECK = "healthcheck" STAGE_SAMPLING = "sampling" +STAGE_SOURCE = "source" STAGE_COLUMN_CONFIG = "column_config" STAGE_GENERATING = "generating" STAGE_BATCH = "batch" diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 3d7cf2dbe6..db59d573a9 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -33,6 +33,60 @@ from .worker import run_job_process _CTX = mp.get_context("spawn") +def _github_source_estimated_total(recipe: dict) -> int | None: + seed_config = recipe.get("seed_config") + if not isinstance(seed_config, dict): + return None + source = seed_config.get("source") + if not isinstance(source, dict) or source.get("seed_type") != "github_repo": + return None + + repos_raw = source.get("repos") + repos = ( + [repo for repo in repos_raw if isinstance(repo, str) and repo.strip()] + if isinstance(repos_raw, list) + else [] + ) + item_types_raw = source.get("item_types") + item_types = ( + [ + item + for item in item_types_raw + if isinstance(item, str) and item in {"issues", "pulls", "commits"} + ] + if isinstance(item_types_raw, list) + else [] + ) + try: + limit = int(source.get("limit") or 0) + except (TypeError, ValueError): + return None + if not repos or not item_types or limit <= 0: + return None + return len(repos) * len(item_types) * limit + + +def _source_progress_status(job: Job) -> dict[str, Any] | None: + progress = job.source_progress + if progress is None: + return None + return { + "source": progress.source, + "status": progress.status, + "repo": progress.repo, + "resource": progress.resource, + "page": progress.page, + "page_items": progress.page_items, + "fetched_items": progress.fetched_items, + "estimated_total": progress.estimated_total, + "percent": progress.percent, + "rate_remaining": progress.rate_remaining, + "retry_after_sec": progress.retry_after_sec, + "message": progress.message, + "updated_at": progress.updated_at, + } + + @dataclass class Subscription: replay: list[dict] @@ -71,8 +125,20 @@ class JobManager: self._pump_thread: threading.Thread | None = None self._seq: int = 0 - def start(self, *, recipe: dict, run: dict) -> str: - """Spawn the job subprocess (one at a time, no cap).""" + def start( + self, + *, + recipe: dict, + run: dict, + internal_api_key_id: int | None = None, + ) -> str: + """Spawn the job subprocess (one at a time, no cap). + + ``internal_api_key_id`` is the row id of a workflow-scoped + sk-unsloth-* key minted by the route layer for local providers. + JobManager revokes it when the job reaches a terminal state so the + key's live window is no longer than the run. + """ llm_columns = recipe.get("columns") or [] llm_column_count = 0 if isinstance(llm_columns, list): @@ -92,6 +158,10 @@ class JobManager: job_id = uuid.uuid4().hex self._job = Job(job_id = job_id, status = "pending", started_at = time.time()) self._job.progress_columns_total = llm_column_count + self._job.source_progress_estimated_total = _github_source_estimated_total( + recipe + ) + self._job.internal_api_key_id = internal_api_key_id self._events.clear() self._seq = 0 @@ -163,6 +233,7 @@ class JobManager: "ok": job.column_progress.ok, "failed": job.column_progress.failed, }, + "source_progress": _source_progress_status(job), "model_usage": { name: { "model": usage.model, @@ -405,6 +476,7 @@ class JobManager: for e in self._drain_queue(mp_q): self._handle_event(job, e) + retired_job: Job | None = None with self._lock: if self._job and self._job.status in { "pending", @@ -429,6 +501,9 @@ class JobManager: "job_id": self._job.job_id, } ) + retired_job = self._job + if retired_job is not None: + self._retire_workflow_key(retired_job) return def _handle_event(self, job: Job, event: dict) -> None: @@ -436,6 +511,7 @@ class JobManager: et = event.get("type") msg = event.get("message") if et == "log" else None + terminal = False with self._lock: if self._job is None or self._job.job_id != job.job_id: return @@ -452,18 +528,43 @@ class JobManager: if self._job.progress.total and self._job.progress.total > 0: self._job.progress.done = self._job.progress.total self._job.progress.percent = 100.0 + terminal = True if et == EVENT_JOB_ERROR: self._job.status = "error" self._job.finished_at = time.time() self._job.error = event.get("error") or "error" + terminal = True + if et == EVENT_JOB_CANCELLED: + terminal = True if msg: upd = parse_log_message(msg) if upd: apply_update(self._job, upd) + if terminal: + self._retire_workflow_key(job) + self._emit(event) + def _retire_workflow_key(self, job: Job) -> None: + """Revoke the workflow-scoped sk-unsloth-* key, if one was minted. + + Best-effort: revocation failures are swallowed. The key would + expire on its own after 24h, so a missed revoke is a latency + concern, not a correctness one. + """ + key_id = getattr(job, "internal_api_key_id", None) + if not key_id: + return + try: + from auth import storage # deferred: avoids circular import + + storage.revoke_internal_api_key(int(key_id)) + except Exception: + pass + job.internal_api_key_id = None + _JOB_MANAGER: JobManager | None = None diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index 324b62a92e..cea6d8ea64 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -4,6 +4,7 @@ from __future__ import annotations import re +import time from dataclasses import dataclass from typing import Any @@ -17,9 +18,10 @@ from .constants import ( STAGE_PREVIEW, STAGE_PROFILING, STAGE_SAMPLING, + STAGE_SOURCE, USAGE_RESET_STAGES, ) -from .types import Job, ModelUsage, Progress +from .types import Job, ModelUsage, Progress, SourceProgress @dataclass(frozen = True) @@ -41,6 +43,7 @@ class ParsedUpdate: usage_requests_total: int | None = None usage_rpm: float | None = None usage_section_start: bool | None = None + source_progress: SourceProgress | None = None # kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI. @@ -61,9 +64,165 @@ _RE_USAGE_TOKENS = re.compile( _RE_USAGE_REQUESTS = re.compile( r"requests:\s*success=(?P\d+),\s*failed=(?P\d+),\s*total=(?P\d+),\s*rpm=(?P[0-9.]+)" ) +_RE_GITHUB_PAGE = re.compile( + r"^\[(?P[^\]\s]+/[^\]\s]+)\]\s+" + r"(?Pissues|PRs|commits)\s+page\s+(?P\d+)\s+" + r"\(\+(?P\d+)\).*?\bremaining=(?P\d+)", + re.IGNORECASE, +) +_RE_GITHUB_RATE_LIMIT = re.compile( + r"Rate limit hit\. Sleeping (?P\d+)s until reset\.", + re.IGNORECASE, +) +_RE_GITHUB_SECONDARY_RATE_LIMIT = re.compile( + r"Secondary rate limit(?: on REST)?\. Sleep (?P\d+)s\.", + re.IGNORECASE, +) +_RE_GITHUB_REST_RATE_LIMIT = re.compile( + r"REST 403/429, sleep (?P\d+)", + re.IGNORECASE, +) +_RE_GITHUB_TRANSIENT = re.compile( + r"^(?PGraphQL|REST) (?P\d{3}) transient, retrying", + re.IGNORECASE, +) +_RE_GITHUB_NETWORK_RETRY = re.compile( + r"^(?PGraphQL|REST) network error: .* Retry\.", + re.IGNORECASE, +) +_RE_GITHUB_TRIAL_LIMIT = re.compile( + r"Trial limit reached for (?Pissues|PRs|commits) \((?P\d+)\)", + re.IGNORECASE, +) +_RE_GITHUB_COMPLETE = re.compile( + r"Scraper complete\. GraphQL calls=\d+ REST calls=\d+", + re.IGNORECASE, +) def parse_log_message(msg: str) -> ParsedUpdate | None: + m = _RE_GITHUB_PAGE.search(msg) + if m: + resource_raw = m.group("resource") + resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower() + repo = m.group("repo") + page = int(m.group("page")) + page_items = int(m.group("items")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "fetching", + repo = repo, + resource = resource, + page = page, + page_items = page_items, + rate_remaining = int(m.group("remaining")), + message = ( + f"Scraping GitHub source: {repo} " + f"{resource} page {page} (+{page_items})" + ), + ), + ) + + m = _RE_GITHUB_RATE_LIMIT.search(msg) + if m: + seconds = int(m.group("seconds")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "rate_limited", + retry_after_sec = seconds, + message = ( + "Waiting for GitHub rate limit. " + "Studio will resume automatically." + ), + ), + ) + + m = _RE_GITHUB_SECONDARY_RATE_LIMIT.search(msg) + if m: + seconds = int(m.group("seconds")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "rate_limited", + retry_after_sec = seconds, + message = ( + "Waiting for GitHub secondary rate limit. " + "Studio will resume automatically." + ), + ), + ) + + m = _RE_GITHUB_REST_RATE_LIMIT.search(msg) + if m: + seconds = int(m.group("seconds")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "rate_limited", + retry_after_sec = seconds, + message = ( + "Waiting for GitHub rate limit. " + "Studio will resume automatically." + ), + ), + ) + + m = _RE_GITHUB_TRIAL_LIMIT.search(msg) + if m: + resource_raw = m.group("resource") + resource = "pulls" if resource_raw.lower() == "prs" else resource_raw.lower() + items = int(m.group("items")) + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "fetching", + resource = resource, + message = f"GitHub {resource} trial limit reached ({items}).", + ), + ) + + m = _RE_GITHUB_TRANSIENT.search(msg) + if m: + api = m.group("api") + code = m.group("code") + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "retrying", + message = f"GitHub {api} returned {code}; retrying automatically.", + ), + ) + + m = _RE_GITHUB_NETWORK_RETRY.search(msg) + if m: + api = m.group("api") + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "retrying", + message = f"GitHub {api} request failed; retrying automatically.", + ), + ) + + if _RE_GITHUB_COMPLETE.search(msg): + return ParsedUpdate( + stage = STAGE_SOURCE, + source_progress = SourceProgress( + source = "github", + status = "completed", + message = "GitHub source scrape complete.", + ), + ) + m = _RE_SAMPLERS.search(msg) if m: return ParsedUpdate( @@ -172,6 +331,8 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: job.batch.idx = update.batch_idx if update.batch_total is not None: job.batch.total = update.batch_total + if update.source_progress is not None: + _apply_source_progress(job, update.source_progress) if update.stage in USAGE_RESET_STAGES: # usage summary is a short block so we reset once we move into the next stage. @@ -216,6 +377,67 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: usage.rpm = update.usage_rpm +def _apply_source_progress(job: Job, progress: SourceProgress) -> None: + previous = job.source_progress + now = time.time() + + page_items = progress.page_items + if progress.repo and progress.resource and progress.page is not None: + page_key = f"{progress.repo}:{progress.resource}:{progress.page}" + count_key = f"{progress.repo}:{progress.resource}" + if page_key not in job._source_seen_pages: + job._source_seen_pages.add(page_key) + job._source_counts[count_key] = int( + job._source_counts.get(count_key, 0) + ) + int(page_items or 0) + + fetched_items = sum(job._source_counts.values()) + if fetched_items <= 0: + fetched_items = progress.fetched_items or ( + previous.fetched_items if previous else None + ) + + estimated_total = ( + progress.estimated_total + or job.source_progress_estimated_total + or (previous.estimated_total if previous else None) + ) + percent: float | None = progress.percent + if percent is None and estimated_total and fetched_items is not None: + raw_percent = (float(fetched_items) / float(max(1, estimated_total))) * 100.0 + percent = 100.0 if progress.status == "completed" else min(99.0, raw_percent) + if percent is None and previous is not None: + percent = previous.percent + + job.source_progress = SourceProgress( + source = "github", + status = progress.status or (previous.status if previous else None), + repo = progress.repo or (previous.repo if previous else None), + resource = progress.resource or (previous.resource if previous else None), + page = ( + progress.page + if progress.page is not None + else (previous.page if previous else None) + ), + page_items = ( + page_items + if page_items is not None + else (previous.page_items if previous else None) + ), + fetched_items = fetched_items, + estimated_total = estimated_total, + percent = percent, + rate_remaining = ( + progress.rate_remaining + if progress.rate_remaining is not None + else (previous.rate_remaining if previous else None) + ), + retry_after_sec = progress.retry_after_sec, + message = progress.message or (previous.message if previous else None), + updated_at = now, + ) + + def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress: if not job.rows: return column_progress diff --git a/studio/backend/core/data_recipe/jobs/types.py b/studio/backend/core/data_recipe/jobs/types.py index 8d77903238..3d3ddb974e 100644 --- a/studio/backend/core/data_recipe/jobs/types.py +++ b/studio/backend/core/data_recipe/jobs/types.py @@ -35,6 +35,23 @@ class BatchProgress: total: int | None = None +@dataclass +class SourceProgress: + source: str = "github" + status: str | None = None + repo: str | None = None + resource: str | None = None + page: int | None = None + page_items: int | None = None + fetched_items: int | None = None + estimated_total: int | None = None + percent: float | None = None + rate_remaining: int | None = None + retry_after_sec: int | None = None + message: str | None = None + updated_at: float | None = None + + @dataclass class ModelUsage: model: str @@ -57,6 +74,7 @@ class Job: progress: Progress = field(default_factory = Progress) column_progress: Progress = field(default_factory = Progress) batch: BatchProgress = field(default_factory = BatchProgress) + source_progress: SourceProgress | None = None rows: int | None = None cols: int | None = None error: str | None = None @@ -70,8 +88,15 @@ class Job: processor_artifacts: dict[str, Any] | None = None model_usage: dict[str, ModelUsage] = field(default_factory = dict) progress_columns_total: int | None = None + source_progress_estimated_total: int | None = None completed_columns: list[str] = field(default_factory = list) + # Id of the internal sk-unsloth-* API key minted for a local-model + # workflow. Revoked when the job terminates so the key's live window + # matches the run rather than its 24h TTL. + internal_api_key_id: int | None = None _current_usage_model: str | None = None _in_usage_summary: bool = False _seen_generation_columns: list[str] = field(default_factory = list) _column_done: dict[str, int] = field(default_factory = dict) + _source_counts: dict[str, int] = field(default_factory = dict) + _source_seen_pages: set[str] = field(default_factory = set) diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index 63e38bd18d..8c5c7fe657 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -21,6 +21,15 @@ from ..service import build_config_builder, create_data_designer from utils.paths import ensure_dir, recipe_datasets_root _ARTIFACT_ROOT = recipe_datasets_root() +_RE_GITHUB_CURSOR = re.compile(r"\bcursor=[^\s,]+") +_RE_SECRET_TOKEN = re.compile( + r"\b(?:(?:ghp|gho|ghu|ghs|ghr|github_pat)_[A-Za-z0-9_]+|sk-unsloth-[A-Za-z0-9]+)" +) + + +def _sanitize_log_message(message: str) -> str: + message = _RE_GITHUB_CURSOR.sub("cursor=", message) + return _RE_SECRET_TOKEN.sub("", message) class _QueueLogHandler(logging.Handler): @@ -35,7 +44,7 @@ class _QueueLogHandler(logging.Handler): "ts": record.created, "level": record.levelname, "logger": record.name, - "message": record.getMessage(), + "message": _sanitize_log_message(record.getMessage()), } self._q.put(event) except (OSError, RuntimeError, ValueError): @@ -119,10 +128,16 @@ def run_job_process( # Attach queue logger directly to `data_designer` so parser events survive root resets. handler = _QueueLogHandler(event_queue) handler.setLevel(logging.INFO) - data_designer_logger = logging.getLogger("data_designer") - data_designer_logger.addHandler(handler) - data_designer_logger.setLevel(logging.INFO) - data_designer_logger.propagate = True + for logger_name in ( + "data_designer", + "scraper", + "gh_client", + "data_designer_github_repo_seed", + ): + logger = logging.getLogger(logger_name) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = True if run_config_raw: designer.set_run_config(RunConfig.model_validate(run_config_raw)) @@ -180,8 +195,8 @@ def run_job_process( { "type": EVENT_JOB_ERROR, "ts": time.time(), - "error": str(exc), - "stack": traceback.format_exc(limit = 20), + "error": _sanitize_log_message(str(exc)), + "stack": _sanitize_log_message(traceback.format_exc(limit = 20)), } ) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 81ff9ae4e5..33722f103c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1618,8 +1618,10 @@ class LlamaCppBackend: # Model fits on selected GPU(s) -- offload all layers cmd.extend(["-ngl", "-1"]) - if n_threads is not None: - cmd.extend(["--threads", str(n_threads)]) + # -1 = llama.cpp auto-detect (physical cores). Pass explicitly so we + # do not inherit llama-server's internal default, which has historically + # varied (hardware concurrency incl. hyperthreads on some builds). + cmd.extend(["--threads", str(n_threads if n_threads is not None else -1)]) # Always enable Jinja chat template rendering for proper template support cmd.extend(["--jinja"]) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/README.md b/studio/backend/plugins/data-designer-github-repo-seed/README.md new file mode 100644 index 0000000000..346d94b305 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/README.md @@ -0,0 +1,73 @@ +# data-designer-github-repo-seed + +A Data Designer seed-reader plugin for **Unsloth Studio** that scrapes real +GitHub data (issues, pull requests, commits) from one or more repositories +and hands it to the recipe pipeline as a seed dataset. + +Designed to ship with Studio as a default seed source so any user with a +GitHub token can build training datasets straight from live repos. + +## What it does + +Given a list of `owner/name` repos, a GitHub token, and a per-resource +`limit`, the plugin uses GitHub's GraphQL API to fetch issues, pull +requests, and/or commits, with labels, state, authors, and the first N +comments of each item, and materialises a single JSONL with uniform +columns so the rest of the recipe (LLM text / LLM structured / processors) +can treat it like any other seed table. + +| Column | Description | +|---------------|------------------------------------------------| +| `item_type` | `issue` / `pull` / `commit` | +| `repo` | `owner/name` | +| `number` | Issue/PR number, or commit SHA | +| `title` | Title (or commit message headline) | +| `body` | Issue/PR body (or full commit message) | +| `state` | `OPEN` / `CLOSED` / `MERGED` (empty for commit)| +| `author` | GitHub login of the author | +| `created_at` | ISO8601 | +| `closed_at` | ISO8601 (empty for commits) | +| `url` | Permalink | +| `labels` | List of label names | +| `comments` | First N comments concatenated | + +## Usage in a recipe + +```json +{ + "seed_config": { + "source": { + "seed_type": "github_repo", + "repos": ["unslothai/unsloth", "unslothai/unsloth-zoo"], + "token": "", + "item_types": ["issues", "pulls"], + "limit": 100, + "include_comments": true, + "max_comments_per_item": 30 + }, + "sampling_strategy": "shuffle", + "selection_strategy": null + } +} +``` + +Leave `token` empty to fall back to the server's `GH_TOKEN` / `GITHUB_TOKEN` +environment variable, useful when the recipe is published and shouldn't +carry a secret. + +## Auth + +A GitHub personal access token with `public_repo` scope is enough for public +repositories; `repo` scope is required for private ones. GraphQL requests +are rate-limit aware: the client inspects `x-ratelimit-*` headers and +sleeps until reset when the budget drops below a safety threshold. + +## Install + +Shipped as a default Studio plugin. For development: + +```bash +pip install -e . +``` + +Registered automatically via the `data_designer.plugins` entry point. diff --git a/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml new file mode 100644 index 0000000000..e232adc60c --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/pyproject.toml @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "data-designer-github-repo-seed" +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", +] + +[project.entry-points."data_designer.plugins"] +github_repo_seed = "data_designer_github_repo_seed.plugin:github_repo_seed_plugin" + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py new file mode 100644 index 0000000000..f57af4c6c3 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +# Intentionally empty. Data-designer loads submodules lazily via qualified names +# (impl_qualified_name / config_qualified_name in plugin.py), so importing this +# package must NOT touch modules that depend on data_designer.engine.* during +# Studio's bootstrap (circular import). diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py new file mode 100644 index 0000000000..6b347c4f83 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/config.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +from typing import Literal + +from pydantic import Field, field_validator, model_validator + +from data_designer.config.seed_source import SeedSource + + +class GitHubRepoSeedSource(SeedSource): + seed_type: Literal["github_repo"] = "github_repo" + + repos: list[str] = Field( + default_factory = list, + description = "List of GitHub repositories to scrape, each in `owner/name` form.", + ) + token: str = Field( + default = "", + description = "Personal access token. Leave blank to read GH_TOKEN / GITHUB_TOKEN from env at run time.", + ) + item_types: list[Literal["issues", "pulls", "commits"]] = Field( + default = ["issues", "pulls"], + description = "Which GitHub item types to fetch per repo.", + ) + limit: int = Field( + default = 100, + ge = 1, + le = 5000, + description = "Maximum items per repo per item type (e.g. limit=100 + ['issues','pulls'] => up to 200 items per repo).", + ) + include_comments: bool = Field( + default = True, + description = "Fetch the first N comments of each issue/PR and include them in the `comments` column.", + ) + max_comments_per_item: int = Field(default = 30, ge = 0, le = 200) + + @field_validator("repos") + @classmethod + def _validate_repos(cls, v: list[str]) -> list[str]: + out: list[str] = [] + for r in v or []: + r = r.strip() + if not r: + continue + if r.count("/") != 1 or not all(r.split("/")): + raise ValueError(f"Each repo must be `owner/name`; got {r!r}") + out.append(r) + return out + + @field_validator("item_types") + @classmethod + def _validate_item_types(cls, v: list[str]) -> list[str]: + if not v: + raise ValueError("item_types must not be empty") + return list(dict.fromkeys(v)) + + @model_validator(mode = "after") + def _ensure_repos(self) -> "GitHubRepoSeedSource": + if not self.repos: + raise ValueError("At least one repo is required") + return self diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py new file mode 100644 index 0000000000..5a38e26d6b --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/impl.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import hashlib +import tempfile +import threading +from pathlib import Path +from typing import Optional + +import data_designer.lazy_heavy_imports as lazy +from data_designer.engine.resources.seed_reader import SeedReader + +from .config import GitHubRepoSeedSource +from .scraper import ScrapeConfig, materialize_to_jsonl + + +# In-process cache mapping a stable config signature to the JSONL materialization +# path. A single recipe job invokes the seed reader multiple times (validation, +# preview, per-column sampling), and the default flow re-scrapes the repo on +# every call: for a 2-repo preview that is ~15s of redundant GitHub GraphQL +# traffic before any generation fires. Memoize the materialization so the second +# and third passes reuse the file the first pass wrote. Cache key excludes the +# raw token and uses a short SHA-256 digest so token values never hit memory +# twice and token rotation invalidates cleanly. +_SCRAPE_CACHE: dict[tuple, str] = {} +_SCRAPE_CACHE_LOCK = threading.Lock() + + +def _scrape_cache_key(cfg: ScrapeConfig) -> tuple: + token_digest = hashlib.sha256( + (cfg.token or "").encode("utf-8"), + ).hexdigest()[:16] + return ( + tuple(cfg.repos), + tuple(cfg.item_types), + cfg.limit, + bool(cfg.include_comments), + cfg.max_comments_per_item, + token_digest, + ) + + +def _lookup_cached_scrape(key: tuple) -> Optional[str]: + with _SCRAPE_CACHE_LOCK: + path = _SCRAPE_CACHE.get(key) + if path and Path(path).exists(): + return path + # Stale entry (tmp cleanup, user restarted, ...); drop it so the caller + # materializes a fresh file rather than returning a dangling path. + if path: + with _SCRAPE_CACHE_LOCK: + _SCRAPE_CACHE.pop(key, None) + return None + + +def _store_cached_scrape(key: tuple, path: str) -> None: + with _SCRAPE_CACHE_LOCK: + _SCRAPE_CACHE[key] = path + + +class GitHubRepoSeedReader(SeedReader[GitHubRepoSeedSource]): + def create_duckdb_connection(self): + return lazy.duckdb.connect() + + def get_dataset_uri(self) -> str: + out_dir = Path(tempfile.gettempdir()) / "studio-github-repo-seed" + cfg = ScrapeConfig( + repos = list(self.source.repos), + token = self.source.token, + item_types = list(self.source.item_types), + limit = self.source.limit, + include_comments = self.source.include_comments, + max_comments_per_item = self.source.max_comments_per_item, + ) + cache_key = _scrape_cache_key(cfg) + cached_path = _lookup_cached_scrape(cache_key) + if cached_path is not None: + return cached_path + path = materialize_to_jsonl(cfg, out_dir) + _store_cached_scrape(cache_key, str(path)) + return str(path) diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py new file mode 100644 index 0000000000..f87dbd0507 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/plugin.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from data_designer.plugins.plugin import Plugin, PluginType + +github_repo_seed_plugin = Plugin( + impl_qualified_name = "data_designer_github_repo_seed.impl.GitHubRepoSeedReader", + config_qualified_name = "data_designer_github_repo_seed.config.GitHubRepoSeedSource", + plugin_type = PluginType.SEED_READER, +) 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 new file mode 100644 index 0000000000..d768fe37be --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Multi-repo GitHub scraper for the Studio seed plugin. + +Drives the GraphQL-based scraper in `scraper_impl/` per repo. Each repo is +scraped with a trial_limits cap so we stop at `limit` items per resource. +After scraping, we read the per-resource JSONL shards and flatten them into +a single unified JSONL with stable columns (`item_type`, `repo`, `number`, +`title`, `body`, ...). +""" + +from __future__ import annotations + +import json +import os +import sys +import time +import uuid +from dataclasses import dataclass +from pathlib import Path + +# Defer scraper_impl imports until `scrape()` runs with a resolved token. +_IMPL_DIR = Path(__file__).parent / "scraper_impl" + + +def _ensure_impl_on_path() -> None: + if str(_IMPL_DIR) not in sys.path: + sys.path.insert(0, str(_IMPL_DIR)) + + +def _load_impl(): + _ensure_impl_on_path() + import importlib + + gh_client = importlib.import_module("gh_client") # type: ignore + scraper_mod = importlib.import_module("scraper") # type: ignore + return gh_client.GitHubClient, scraper_mod.RepoScraper + + +@dataclass +class ScrapeConfig: + repos: list[str] + token: str + item_types: list[str] + limit: int + include_comments: bool + max_comments_per_item: int + + +def _resolve_token(token: str) -> str: + tok = token or os.environ.get("GH_TOKEN", "") or os.environ.get("GITHUB_TOKEN", "") + if not tok: + raise ValueError( + "GitHub token is required. Set it in the recipe config or the GH_TOKEN / GITHUB_TOKEN env var." + ) + return tok + + +def _read_jsonl(path: Path, max_rows: int | None = None): + if not path.exists(): + return + with path.open(encoding = "utf-8") as f: + for i, line in enumerate(f): + if not line.strip(): + continue + if max_rows is not None and i >= max_rows: + return + try: + yield json.loads(line) + except json.JSONDecodeError: + continue + + +def _flatten_issue_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] + comments_nodes = (r.get("comments") or {}).get("nodes") or [] + comments_text = "" + if include_comments and comments_nodes: + kept = comments_nodes[:max_c] + comments_text = "\n\n".join( + f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" + for c in kept + ) + return { + "item_type": "issue", + "repo": repo, + "number": r.get("number"), + "title": r.get("title") or "", + "body": r.get("body") or "", + "state": r.get("state") or "", + "author": (r.get("author") or {}).get("login", ""), + "created_at": r.get("createdAt") or "", + "closed_at": r.get("closedAt") or "", + "url": r.get("url") or r.get("permalink") or "", + "labels": labels, + "comments": comments_text, + } + + +def _flatten_pr_row(r: dict, repo: str, include_comments: bool, max_c: int) -> dict: + labels = [ + l.get("name") + for l in (r.get("labels", {}) or {}).get("nodes", []) + if l.get("name") + ] + comments_nodes = (r.get("comments") or {}).get("nodes") or [] + comments_text = "" + if include_comments and comments_nodes: + kept = comments_nodes[:max_c] + comments_text = "\n\n".join( + f"[{(c.get('author') or {}).get('login', '?')}]: {c.get('body') or ''}" + for c in kept + ) + return { + "item_type": "pull", + "repo": repo, + "number": r.get("number"), + "title": r.get("title") or "", + "body": r.get("body") or "", + "state": r.get("state") or "", + "author": (r.get("author") or {}).get("login", ""), + "created_at": r.get("createdAt") or "", + "closed_at": r.get("closedAt") or "", + "url": r.get("url") or r.get("permalink") or "", + "labels": labels, + "comments": comments_text, + } + + +def _flatten_commit_row(r: dict, repo: str) -> dict: + msg = r.get("messageHeadline") or r.get("message") or "" + body = r.get("messageBody") or r.get("message") or msg + author = r.get("author") or {} + return { + "item_type": "commit", + "repo": repo, + "number": r.get("oid") or r.get("sha") or "", + "title": msg, + "body": body, + "state": "", + "author": (author.get("user") or {}).get("login") or author.get("name", ""), + "created_at": (author.get("date") or r.get("committedDate") or ""), + "closed_at": "", + "url": r.get("url") or "", + "labels": [], + "comments": "", + } + + +def scrape(cfg: ScrapeConfig, base_dir: Path): + token = _resolve_token(cfg.token) + GitHubClient, RepoScraper = _load_impl() + 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. + effective_limit = cfg.limit if cfg.limit and cfg.limit > 0 else 1_000_000 + trial_limits: dict[str, int] = {} + if "issues" in cfg.item_types: + trial_limits["issues"] = effective_limit + if "pulls" in cfg.item_types: + trial_limits["pull_requests"] = effective_limit + if "commits" in cfg.item_types: + trial_limits["commits"] = effective_limit + + all_rows: list[dict] = [] + for repo in cfg.repos: + owner, name = repo.split("/", 1) + scraper = RepoScraper( + owner = owner, + name = name, + base_dir = base_dir, + client = client, + trial_limits = trial_limits, + light = True, + ) + try: + 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: + 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() + + read_cap = cfg.limit if cfg.limit and cfg.limit > 0 else None + repo_dir = base_dir / f"{owner}__{name}" + if "issues" in cfg.item_types: + for row in _read_jsonl(repo_dir / "issues.jsonl", read_cap): + all_rows.append( + _flatten_issue_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) + ) + if "pulls" in cfg.item_types: + for row in _read_jsonl(repo_dir / "pull_requests.jsonl", read_cap): + all_rows.append( + _flatten_pr_row( + row, repo, cfg.include_comments, cfg.max_comments_per_item + ) + ) + if "commits" in cfg.item_types: + for row in _read_jsonl(repo_dir / "commits.jsonl", read_cap): + all_rows.append(_flatten_commit_row(row, repo)) + + return all_rows + + +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) + 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-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") + return out 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 new file mode 100644 index 0000000000..32014236c6 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/__init__.py @@ -0,0 +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 new file mode 100644 index 0000000000..dd2de2f5ce --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/gh_client.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GitHub API client with rate-limit awareness, retry, and dual REST/GraphQL support.""" + +from __future__ import annotations + +import json +import os +import time +import logging +from typing import Any, Dict, Iterable, Iterator, List, Optional + +import requests + +log = logging.getLogger("gh_client") + +GRAPHQL_URL = "https://api.github.com/graphql" +REST_BASE = "https://api.github.com" + +BASE_HEADERS = { + "Accept": "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + "User-Agent": "github-data-gatherer/1.0", +} + + +class RateLimitError(Exception): + pass + + +class GitHubClient: + 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( + {**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 + self.graphql_reset: Optional[int] = None + self.rest_remaining: Optional[int] = None + self.rest_reset: Optional[int] = None + self.calls_graphql = 0 + self.calls_rest = 0 + self.retry_count = 0 + + def _sleep_until(self, reset_ts: int, buffer_s: int = 10) -> None: + now = int(time.time()) + wait = max(0, reset_ts - now) + buffer_s + log.warning("Rate limit hit. Sleeping %ds until reset.", wait) + time.sleep(wait) + + def _check_rate_and_wait(self, kind: str) -> None: + if kind == "graphql": + remaining = self.graphql_remaining + reset = self.graphql_reset + min_remaining = self.min_remaining_graphql + else: + remaining = self.rest_remaining + reset = self.rest_reset + min_remaining = self.min_remaining_rest + if remaining is not None and remaining < min_remaining: + if reset: + self._sleep_until(reset) + # Reset remaining so we don't spin + if kind == "graphql": + self.graphql_remaining = None + else: + self.rest_remaining = None + + def graphql( + self, + query: str, + variables: Optional[Dict[str, Any]] = None, + max_retries: int = 20, + ) -> Dict[str, Any]: + self._check_rate_and_wait("graphql") + backoff = 2 + last_err = None + for attempt in range(max_retries): + try: + r = self.session.post( + GRAPHQL_URL, + json = {"query": query, "variables": variables or {}}, + timeout = 120, + ) + self.calls_graphql += 1 + # Update rate info from response headers + rem = r.headers.get("X-RateLimit-Remaining") + rst = r.headers.get("X-RateLimit-Reset") + if rem is not None: + try: + self.graphql_remaining = int(rem) + except ValueError: + pass + if rst is not None: + try: + self.graphql_reset = int(rst) + except ValueError: + pass + if r.status_code in (502, 503, 504): + log.warning("GraphQL %s transient, retrying", r.status_code) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + continue + if r.status_code == 403 or r.status_code == 429: + # Check for secondary/abuse + retry_after = r.headers.get("Retry-After") + if retry_after: + t = int(retry_after) + log.warning("Secondary rate limit. Sleep %ds.", t) + time.sleep(t + 2) + continue + if self.graphql_reset: + self._sleep_until(self.graphql_reset) + continue + time.sleep(60) + continue + r.raise_for_status() + data = r.json() + if "errors" in data and data["errors"]: + # Surface errors but allow partial data + errs = data["errors"] + # Retry on RATE_LIMITED + for e in errs: + if e.get("type") == "RATE_LIMITED": + self._sleep_until( + (self.graphql_reset or int(time.time()) + 60) + ) + break + else: + # No rate-limit error, log and return partial + log.warning("GraphQL errors: %s", json.dumps(errs)[:400]) + return data + continue + return data + except requests.RequestException as e: + last_err = e + log.warning("GraphQL network error: %s. Retry.", e) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + raise RuntimeError(f"GraphQL failed after {max_retries} retries: {last_err}") + + def rest( + self, + method: str, + path: str, + params: Optional[Dict[str, Any]] = None, + json_body: Optional[Dict[str, Any]] = None, + max_retries: int = 6, + ) -> requests.Response: + self._check_rate_and_wait("rest") + if path.startswith("http"): + url = path + else: + url = REST_BASE + path + backoff = 2 + last_err = None + for attempt in range(max_retries): + try: + r = self.session.request( + method, url, params = params, json = json_body, timeout = 120 + ) + self.calls_rest += 1 + rem = r.headers.get("X-RateLimit-Remaining") + rst = r.headers.get("X-RateLimit-Reset") + if rem is not None: + try: + self.rest_remaining = int(rem) + except ValueError: + pass + if rst is not None: + try: + self.rest_reset = int(rst) + except ValueError: + pass + if r.status_code in (502, 503, 504): + log.warning("REST %s transient, retrying", r.status_code) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + continue + if r.status_code in (403, 429): + retry_after = r.headers.get("Retry-After") + if retry_after: + t = int(retry_after) + log.warning("Secondary rate limit on REST. Sleep %ds.", t) + time.sleep(t + 2) + continue + # Check if primary rate + if self.rest_remaining == 0 and self.rest_reset: + self._sleep_until(self.rest_reset) + continue + log.warning("REST 403/429, sleep 60") + time.sleep(60) + continue + return r + except requests.RequestException as e: + last_err = e + log.warning("REST network error: %s. Retry.", e) + time.sleep(backoff) + backoff = min(backoff * 2, 60) + raise RuntimeError(f"REST failed after {max_retries} retries: {last_err}") + + def rest_paginate( + self, path: str, params: Optional[Dict[str, Any]] = None, per_page: int = 100 + ) -> Iterator[dict]: + params = dict(params or {}) + params.setdefault("per_page", per_page) + url = path + while True: + r = self.rest("GET", url, params = params if url == path else None) + if r.status_code != 200: + log.error( + "REST paginate got %s at %s: %s", r.status_code, url, r.text[:200] + ) + return + items = r.json() + if isinstance(items, dict): + # Some endpoints return dict with list field + items = items.get("items", []) + for it in items: + yield it + # Follow link header + link = r.headers.get("Link", "") + nxt = None + for part in link.split(","): + if 'rel="next"' in part: + nxt = part.split(";")[0].strip().strip("<>") + break + if not nxt: + return + url = nxt + params = None + + def rate_snapshot(self) -> Dict[str, Any]: + r = self.rest("GET", "/rate_limit") + if r.status_code == 200: + return r.json() + return {} diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py new file mode 100644 index 0000000000..9dc7613db5 --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/queries.py @@ -0,0 +1,685 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GraphQL queries for GitHub data scraping. + +GitHub's GraphQL rejects queries that define unused fragments, so each query +only includes the fragments it actually references. +""" + +# ---- Fragments (kept as raw strings, composed per query) ---- +F_ACTOR = """ +fragment ActorFields on Actor { + __typename + login + url + avatarUrl + ... on User { id databaseId name } + ... on Bot { id databaseId } + ... on Organization { id databaseId name } +} +""" + +F_LABEL = """ +fragment LabelFields on Label { + id + name + color + description + createdAt +} +""" + +F_TIMELINE = """ +fragment TimelineItem on IssueTimelineItems { + __typename + ... on Node { id } + ... on AddedToProjectEvent { createdAt actor { ...ActorFields } } + ... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on ClosedEvent { createdAt actor { ...ActorFields } stateReason closer { __typename ... on Commit { oid url } ... on PullRequest { number url } } } + ... on CommentDeletedEvent { createdAt actor { ...ActorFields } } + ... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url repository { nameWithOwner } } ... on PullRequest { number url repository { nameWithOwner } } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on ConvertedNoteToIssueEvent { createdAt actor { ...ActorFields } } + ... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } } + ... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } } + ... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on LockedEvent { createdAt actor { ...ActorFields } lockReason } + ... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on MentionedEvent { createdAt actor { ...ActorFields } } + ... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } } + ... on PinnedEvent { createdAt actor { ...ActorFields } } + ... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } } + ... on RemovedFromProjectEvent { createdAt actor { ...ActorFields } } + ... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle } + ... on ReopenedEvent { createdAt actor { ...ActorFields } } + ... on SubscribedEvent { createdAt actor { ...ActorFields } } + ... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } } + ... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on UnlockedEvent { createdAt actor { ...ActorFields } } + ... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } } + ... on UnpinnedEvent { createdAt actor { ...ActorFields } } + ... on UnsubscribedEvent { createdAt actor { ...ActorFields } } + ... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration } +} +""" + +F_PR_TIMELINE = """ +fragment PRTimelineItem on PullRequestTimelineItems { + __typename + ... on Node { id } + ... on AssignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on AutoMergeDisabledEvent { createdAt actor { ...ActorFields } reason } + ... on AutoMergeEnabledEvent { createdAt actor { ...ActorFields } } + ... on AutoRebaseEnabledEvent { createdAt actor { ...ActorFields } } + ... on AutoSquashEnabledEvent { createdAt actor { ...ActorFields } } + ... on AutomaticBaseChangeFailedEvent { createdAt actor { ...ActorFields } oldBase newBase } + ... on AutomaticBaseChangeSucceededEvent { createdAt actor { ...ActorFields } oldBase newBase } + ... on BaseRefChangedEvent { createdAt actor { ...ActorFields } previousRefName currentRefName } + ... on BaseRefDeletedEvent { createdAt actor { ...ActorFields } baseRefName } + ... on BaseRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } } + ... on ClosedEvent { createdAt actor { ...ActorFields } stateReason } + ... on CommentDeletedEvent { createdAt actor { ...ActorFields } } + ... on ConnectedEvent { createdAt actor { ...ActorFields } source { __typename ... on Issue { number url } ... on PullRequest { number url } } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on ConvertToDraftEvent { createdAt actor { ...ActorFields } } + ... on CrossReferencedEvent { createdAt actor { ...ActorFields } isCrossRepository willCloseTarget source { __typename ... on Issue { number url repository { nameWithOwner } title } ... on PullRequest { number url repository { nameWithOwner } title } } } + ... on DemilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on DeployedEvent { createdAt actor { ...ActorFields } } + ... on DeploymentEnvironmentChangedEvent { createdAt actor { ...ActorFields } } + ... on DisconnectedEvent { createdAt actor { ...ActorFields } subject { __typename ... on Issue { number url } ... on PullRequest { number url } } source { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on HeadRefDeletedEvent { createdAt actor { ...ActorFields } headRefName } + ... on HeadRefForcePushedEvent { createdAt actor { ...ActorFields } beforeCommit { oid } afterCommit { oid } ref { name } } + ... on HeadRefRestoredEvent { createdAt actor { ...ActorFields } } + ... on IssueComment { id databaseId createdAt updatedAt author { ...ActorFields } body url reactionGroups { content reactors { totalCount } } } + ... on LabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on LockedEvent { createdAt actor { ...ActorFields } lockReason } + ... on MarkedAsDuplicateEvent { createdAt actor { ...ActorFields } canonical { __typename ... on Issue { number url } ... on PullRequest { number url } } } + ... on MentionedEvent { createdAt actor { ...ActorFields } } + ... on MergedEvent { createdAt actor { ...ActorFields } commit { oid url } mergeRefName } + ... on MilestonedEvent { createdAt actor { ...ActorFields } milestoneTitle } + ... on MovedColumnsInProjectEvent { createdAt actor { ...ActorFields } } + ... on PinnedEvent { createdAt actor { ...ActorFields } } + ... on PullRequestCommit { commit { oid url message author { user { login } date } committedDate } } + ... on PullRequestCommitCommentThread { commit { oid } } + ... on PullRequestReview { id databaseId createdAt submittedAt author { ...ActorFields } body state url reactionGroups { content reactors { totalCount } } } + ... on PullRequestReviewThread { id isResolved isOutdated path line diffSide } + ... on PullRequestRevisionMarker { createdAt lastSeenCommit { oid } } + ... on ReadyForReviewEvent { createdAt actor { ...ActorFields } } + ... on ReferencedEvent { createdAt actor { ...ActorFields } commit { oid url } commitRepository { nameWithOwner } } + ... on RenamedTitleEvent { createdAt actor { ...ActorFields } previousTitle currentTitle } + ... on ReopenedEvent { createdAt actor { ...ActorFields } } + ... on ReviewDismissedEvent { createdAt actor { ...ActorFields } dismissalMessage previousReviewState } + ... on ReviewRequestRemovedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } } + ... on ReviewRequestedEvent { createdAt actor { ...ActorFields } requestedReviewer { __typename ... on User { login } ... on Team { name } } } + ... on SubscribedEvent { createdAt actor { ...ActorFields } } + ... on TransferredEvent { createdAt actor { ...ActorFields } fromRepository { nameWithOwner } } + ... on UnassignedEvent { createdAt actor { ...ActorFields } assignee { __typename ... on User { login } ... on Bot { login } } } + ... on UnlabeledEvent { createdAt actor { ...ActorFields } label { name color } } + ... on UnlockedEvent { createdAt actor { ...ActorFields } } + ... on UnmarkedAsDuplicateEvent { createdAt actor { ...ActorFields } } + ... on UnpinnedEvent { createdAt actor { ...ActorFields } } + ... on UnsubscribedEvent { createdAt actor { ...ActorFields } } + ... on UserBlockedEvent { createdAt actor { ...ActorFields } blockDuration } +} +""" + + +def _q(parts: list[str], body: str) -> str: + return "\n".join(parts + [body]) + + +ISSUES_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL, F_TIMELINE], + """ +query IssuesPage($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state stateReason + createdAt updatedAt closedAt + url + author { ...ActorFields } + editor { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + assignees(first: 20) { nodes { login id } } + milestone { title number state dueOn } + reactionGroups { content reactors { totalCount } } + comments(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + timelineItems(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { ...TimelineItem } + } + trackedInIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } } + trackedIssues(first: 20) { totalCount nodes { number url repository { nameWithOwner } } } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PRS_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL, F_PR_TIMELINE], + """ +query PRsPage($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state isDraft + createdAt updatedAt closedAt mergedAt + url + headRefName headRefOid + baseRefName baseRefOid + additions deletions changedFiles + mergeable merged mergeStateStatus + author { ...ActorFields } + editor { ...ActorFields } + mergedBy { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + assignees(first: 20) { nodes { login id } } + milestone { title number state dueOn } + reactionGroups { content reactors { totalCount } } + closingIssuesReferences(first: 20) { totalCount nodes { number url repository { nameWithOwner } title } } + comments(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + reviewThreads(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id isResolved isOutdated path line diffSide + comments(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body path diffHunk + author { ...ActorFields } + editor { ...ActorFields } + position originalPosition line originalLine + commit { oid } + reactionGroups { content reactors { totalCount } } + } + } + } + } + reviews(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId state createdAt submittedAt body url + author { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + commits(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + commit { + oid + message + messageHeadline + committedDate + authoredDate + author { name email user { login } date } + committer { name email user { login } date } + additions deletions changedFilesIfAvailable + parents(first: 3) { nodes { oid } } + } + } + } + files(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + path additions deletions changeType + } + } + timelineItems(first: 100) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { ...PRTimelineItem } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PRS_PAGE_QUERY_LIGHT = _q( + [F_ACTOR, F_LABEL], + """ +query PRsPageLight($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequests(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state isDraft + createdAt updatedAt closedAt mergedAt + url + author { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + comments(first: 30) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +ISSUES_PAGE_QUERY_LIGHT = _q( + [F_ACTOR, F_LABEL], + """ +query IssuesPageLight($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issues(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body state + createdAt updatedAt closedAt + url + author { ...ActorFields } + labels(first: 50) { nodes { ...LabelFields } } + comments(first: 30) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +ISSUE_COMMENTS_QUERY = _q( + [F_ACTOR], + """ +query IssueComments($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issueOrPullRequest(number: $number) { + __typename + ... on Issue { + comments(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + ... on PullRequest { + comments(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId createdAt updatedAt url body + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +ISSUE_TIMELINE_QUERY = _q( + [F_ACTOR, F_TIMELINE], + """ +query IssueTimeline($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + issue(number: $number) { + timelineItems(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { ...TimelineItem } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PR_TIMELINE_QUERY = _q( + [F_ACTOR, F_PR_TIMELINE], + """ +query PRTimeline($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + timelineItems(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { ...PRTimelineItem } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +PR_COMMITS_QUERY = """ +query PRCommits($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + commits(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + commit { + oid message messageHeadline committedDate authoredDate + author { name email user { login } date } + committer { name email user { login } date } + additions deletions changedFilesIfAvailable + parents(first: 3) { nodes { oid } } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +PR_FILES_QUERY = """ +query PRFiles($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + files(first: 100, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { path additions deletions changeType } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +PR_REVIEW_THREADS_QUERY = _q( + [F_ACTOR], + """ +query PRReviewThreads($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + pullRequest(number: $number) { + reviewThreads(first: 50, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id isResolved isOutdated path line diffSide + comments(first: 50) { + totalCount + nodes { + id databaseId createdAt updatedAt url body path diffHunk + author { ...ActorFields } + editor { ...ActorFields } + position originalPosition line originalLine + commit { oid } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +DISCUSSIONS_PAGE_QUERY = _q( + [F_ACTOR, F_LABEL], + """ +query DiscussionsPage($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + discussions(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + id databaseId number title body + createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + locked + answerChosenAt + closed closedAt + category { id name emoji description isAnswerable } + labels(first: 30) { nodes { ...LabelFields } } + upvoteCount + answer { id databaseId body author { ...ActorFields } createdAt url } + reactionGroups { content reactors { totalCount } } + comments(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + upvoteCount + isAnswer + reactionGroups { content reactors { totalCount } } + replies(first: 50) { + totalCount + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +DISCUSSION_COMMENTS_QUERY = _q( + [F_ACTOR], + """ +query DiscussionComments($owner: String!, $name: String!, $number: Int!, $after: String) { + repository(owner: $owner, name: $name) { + discussion(number: $number) { + comments(first: 50, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + upvoteCount + isAnswer + reactionGroups { content reactors { totalCount } } + replies(first: 50) { + totalCount + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +DISCUSSION_REPLIES_QUERY = _q( + [F_ACTOR], + """ +query DiscussionReplies($commentId: ID!, $after: String) { + node(id: $commentId) { + ... on DiscussionComment { + replies(first: 50, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId body createdAt updatedAt url + author { ...ActorFields } + editor { ...ActorFields } + reactionGroups { content reactors { totalCount } } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +COMMITS_PAGE_QUERY = """ +query CommitsPage($owner: String!, $name: String!, $first: Int!, $after: String, $branch: String!) { + repository(owner: $owner, name: $name) { + ref(qualifiedName: $branch) { + target { + ... on Commit { + history(first: $first, after: $after) { + pageInfo { hasNextPage endCursor } + totalCount + nodes { + oid + message + messageHeadline + committedDate + authoredDate + url + additions deletions changedFilesIfAvailable + author { name email date user { login id } } + committer { name email date user { login id } } + parents(first: 3) { nodes { oid } } + associatedPullRequests(first: 5) { nodes { number url state } } + } + } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +RELEASES_QUERY = _q( + [F_ACTOR], + """ +query Releases($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + releases(first: $first, after: $after, orderBy: {field: CREATED_AT, direction: ASC}) { + pageInfo { hasNextPage endCursor } + nodes { + id databaseId name tagName description + createdAt publishedAt updatedAt + isDraft isPrerelease isLatest + url + author { ...ActorFields } + tagCommit { oid url } + reactionGroups { content reactors { totalCount } } + releaseAssets(first: 50) { + nodes { name contentType size downloadUrl createdAt updatedAt } + } + } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +LABELS_QUERY = _q( + [F_LABEL], + """ +query LabelsList($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + labels(first: $first, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { ...LabelFields } + } + } + rateLimit { cost remaining resetAt } +} +""", +) + +MILESTONES_QUERY = """ +query Milestones($owner: String!, $name: String!, $first: Int!, $after: String) { + repository(owner: $owner, name: $name) { + milestones(first: $first, after: $after) { + pageInfo { hasNextPage endCursor } + nodes { + id number title description state + createdAt updatedAt closedAt dueOn + creator { login } + } + } + } + rateLimit { cost remaining resetAt } +} +""" + +REPO_META_QUERY = """ +query RepoMeta($owner: String!, $name: String!) { + repository(owner: $owner, name: $name) { + id databaseId name nameWithOwner description url + createdAt updatedAt pushedAt + isArchived isDisabled isFork isPrivate + primaryLanguage { name } + languages(first: 20, orderBy: {field: SIZE, direction: DESC}) { + edges { size node { name } } + totalSize + } + stargazerCount forkCount watchers { totalCount } + diskUsage + licenseInfo { key name } + homepageUrl + defaultBranchRef { name } + } + rateLimit { cost remaining resetAt } +} +""" 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 new file mode 100644 index 0000000000..127129e18b --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/scraper.py @@ -0,0 +1,756 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Main scraper orchestration. Collects issues, PRs, discussions, commits, releases, etc. + +Resumable via state file. Writes JSONL shards under data/{repo}/{resource}.jsonl. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple + +# Allow running as a module or script +THIS_DIR = Path(__file__).resolve().parent +if str(THIS_DIR) not in sys.path: + sys.path.insert(0, str(THIS_DIR)) + +from gh_client import GitHubClient +from state_store import JsonlWriter, StateStore +import queries as Q + +log = logging.getLogger("scraper") + + +def ts() -> str: + return time.strftime("%Y-%m-%d %H:%M:%S") + + +class RepoScraper: + def __init__( + self, + owner: str, + name: str, + base_dir: Path, + client: GitHubClient, + trial_limits: Optional[Dict[str, int]] = None, + light: bool = False, + ): + self.owner = owner + self.name = name + self.base_dir = base_dir + self.client = client + self.trial_limits = trial_limits or {} + # When light=True, use trimmed GraphQL queries (no reviewThreads, + # reviews, commits, timelineItems, files) so PR pages can be much + # larger without blowing GitHub's node-count ceiling. + self.light = light + self.repo_dir = base_dir / f"{owner}__{name}" + self.repo_dir.mkdir(parents = True, exist_ok = True) + self.state = StateStore(base_dir / "state" / f"{owner}__{name}.json") + + # Writers + self.writers: Dict[str, JsonlWriter] = {} + for key in ( + "issues", + "pull_requests", + "discussions", + "commits", + "releases", + "labels", + "milestones", + "pr_extra_comments", + "pr_extra_timeline", + "pr_extra_reviews", + "issue_extra_comments", + "issue_extra_timeline", + "discussion_extra_comments", + "discussion_extra_replies", + "repo_meta", + ): + self.writers[key] = JsonlWriter(self.repo_dir / f"{key}.jsonl") + + # ----- helpers ----- + def _trial_stop(self, key: str, counter: int) -> bool: + lim = self.trial_limits.get(key) + if lim is None: + return False + return counter >= lim + + def _log_rate(self, where: str, data: Dict[str, Any]) -> None: + rl = ( + data.get("data", {}).get("rateLimit") + if isinstance(data.get("data"), dict) + else None + ) + if rl: + log.debug( + "[%s] rate cost=%s remaining=%s resetAt=%s", + where, + rl.get("cost"), + rl.get("remaining"), + rl.get("resetAt"), + ) + + # ----- repo meta ----- + def scrape_repo_meta(self) -> Dict[str, Any]: + data = self.client.graphql( + Q.REPO_META_QUERY, {"owner": self.owner, "name": self.name} + ) + self._log_rate("repo_meta", data) + repo = data.get("data", {}).get("repository") or {} + repo["_fetchedAt"] = ts() + self.writers["repo_meta"].write(repo) + return repo + + # ----- issues ----- + def scrape_issues(self) -> int: + key = "issues" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + log.info("%s/%s issues already complete", self.owner, self.name) + return 0 + total_new = 0 + page = 0 + # Light query skips heavy nested fields; safe at 50 per page. + # Clamp by trial_limit so e.g. limit=1 asks GitHub for first:1 + # instead of fetching a full 50-item page and discarding 49. + page_cap = 50 if self.light else 15 + trial_cap = self.trial_limits.get(key) + per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } + query = Q.ISSUES_PAGE_QUERY_LIGHT if self.light else Q.ISSUES_PAGE_QUERY + data = self.client.graphql(query, vars_) + self._log_rate("issues", data) + repo = (data.get("data") or {}).get("repository") or {} + issues = repo.get("issues") or {} + nodes = issues.get("nodes") or [] + for it in nodes: + it["_owner"] = self.owner + it["_repo"] = self.name + it["_fetchedAt"] = ts() + if not self.light: + if it.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_issue_comments( + it["number"], it["comments"]["pageInfo"]["endCursor"] + ) + if ( + it.get("timelineItems", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): + self._paginate_issue_timeline( + it["number"], + it["timelineItems"]["pageInfo"]["endCursor"], + ) + if self.writers[key].write(it): + total_new += 1 + info = issues.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] issues page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + log.info("Trial limit reached for issues (%d)", total_new) + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + def _paginate_issue_comments(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "issueOrPullRequest" + ) or {} + comments = item.get("comments") or {} + for c in comments.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_issueNumber"] = number + self.writers["issue_extra_comments"].write(c) + info = comments.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_issue_timeline(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.ISSUE_TIMELINE_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get("issue") or {} + tl = item.get("timelineItems") or {} + for ev in tl.get("nodes") or []: + ev["_owner"] = self.owner + ev["_repo"] = self.name + ev["_issueNumber"] = number + self.writers["issue_extra_timeline"].write(ev) + info = tl.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + # ----- PRs ----- + def scrape_prs(self) -> int: + key = "pull_requests" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + log.info("%s/%s PRs already complete", self.owner, self.name) + return 0 + total_new = 0 + page = 0 + # Heavy nested PR query is capped at 3 per page (GitHub node-count + # ceiling); light query skips reviewThreads/reviews/commits/etc and + # can safely go to 25 per page. Clamp by trial_limit for small + # previews so limit=1 does not fetch a whole 25-item page. + page_cap = 25 if self.light else 3 + trial_cap = self.trial_limits.get(key) + per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } + query = Q.PRS_PAGE_QUERY_LIGHT if self.light else Q.PRS_PAGE_QUERY + data = self.client.graphql(query, vars_) + self._log_rate("prs", data) + repo = (data.get("data") or {}).get("repository") or {} + prs = repo.get("pullRequests") or {} + nodes = prs.get("nodes") or [] + for pr in nodes: + pr["_owner"] = self.owner + pr["_repo"] = self.name + pr["_fetchedAt"] = ts() + num = pr["number"] + if not self.light: + if pr.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_comments( + num, pr["comments"]["pageInfo"]["endCursor"] + ) + if ( + pr.get("timelineItems", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): + self._paginate_pr_timeline( + num, pr["timelineItems"]["pageInfo"]["endCursor"] + ) + if pr.get("commits", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_commits( + num, pr["commits"]["pageInfo"]["endCursor"] + ) + if pr.get("files", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_pr_files( + num, pr["files"]["pageInfo"]["endCursor"] + ) + if ( + pr.get("reviewThreads", {}) + .get("pageInfo", {}) + .get("hasNextPage") + ): + self._paginate_pr_review_threads( + num, pr["reviewThreads"]["pageInfo"]["endCursor"] + ) + if self.writers[key].write(pr): + total_new += 1 + info = prs.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] PRs page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + log.info("Trial limit reached for PRs (%d)", total_new) + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + def _paginate_pr_comments(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.ISSUE_COMMENTS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "issueOrPullRequest" + ) or {} + comments = item.get("comments") or {} + for c in comments.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_prNumber"] = number + self.writers["pr_extra_comments"].write(c) + info = comments.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_timeline(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_TIMELINE_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + tl = item.get("timelineItems") or {} + for ev in tl.get("nodes") or []: + ev["_owner"] = self.owner + ev["_repo"] = self.name + ev["_prNumber"] = number + self.writers["pr_extra_timeline"].write(ev) + info = tl.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_commits(self, number: int, after: str) -> None: + cur = after + out_key = "pr_extra_commits" + if out_key not in self.writers: + self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_COMMITS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + cc = item.get("commits") or {} + for c in cc.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_prNumber"] = number + self.writers[out_key].write(c) + info = cc.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_files(self, number: int, after: str) -> None: + cur = after + out_key = "pr_extra_files" + if out_key not in self.writers: + self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_FILES_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + ff = item.get("files") or {} + for f in ff.get("nodes") or []: + f["_owner"] = self.owner + f["_repo"] = self.name + f["_prNumber"] = number + # files don't have id, synthesize one + f["_syntheticId"] = f"{self.owner}/{self.name}#{number}:{f.get('path')}" + self.writers[out_key].write(f) + info = ff.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_pr_review_threads(self, number: int, after: str) -> None: + cur = after + out_key = "pr_extra_review_threads" + if out_key not in self.writers: + self.writers[out_key] = JsonlWriter(self.repo_dir / f"{out_key}.jsonl") + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.PR_REVIEW_THREADS_QUERY, vars_) + item = ((data.get("data") or {}).get("repository") or {}).get( + "pullRequest" + ) or {} + rt = item.get("reviewThreads") or {} + for th in rt.get("nodes") or []: + th["_owner"] = self.owner + th["_repo"] = self.name + th["_prNumber"] = number + self.writers[out_key].write(th) + info = rt.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + # ----- Discussions ----- + def scrape_discussions(self) -> int: + key = "discussions" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + log.info("%s/%s discussions already complete", self.owner, self.name) + return 0 + total_new = 0 + page = 0 + per_page = 15 + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + } + data = self.client.graphql(Q.DISCUSSIONS_PAGE_QUERY, vars_) + self._log_rate("discussions", data) + repo = (data.get("data") or {}).get("repository") or {} + dd = repo.get("discussions") or {} + nodes = dd.get("nodes") or [] + for d in nodes: + d["_owner"] = self.owner + d["_repo"] = self.name + d["_fetchedAt"] = ts() + num = d["number"] + if d.get("comments", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_discussion_comments( + num, d["comments"]["pageInfo"]["endCursor"] + ) + # paginate replies per comment if needed + for c in d.get("comments", {}).get("nodes", []) or []: + if c.get("replies", {}).get("pageInfo", {}).get("hasNextPage"): + self._paginate_discussion_replies( + c["id"], c["replies"]["pageInfo"]["endCursor"], num + ) + if self.writers[key].write(d): + total_new += 1 + info = dd.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] discussions page %d (+%d) cursor=%s remaining=%s", + self.owner, + self.name, + page, + len(nodes), + str(cursor)[:20], + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + def _paginate_discussion_comments(self, number: int, after: str) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "number": number, + "after": cur, + } + data = self.client.graphql(Q.DISCUSSION_COMMENTS_QUERY, vars_) + disc = ((data.get("data") or {}).get("repository") or {}).get( + "discussion" + ) or {} + cc = disc.get("comments") or {} + for c in cc.get("nodes") or []: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_discussionNumber"] = number + self.writers["discussion_extra_comments"].write(c) + info = cc.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + def _paginate_discussion_replies( + self, comment_id: str, after: str, disc_number: int + ) -> None: + cur = after + while cur: + vars_ = { + "owner": self.owner, + "name": self.name, + "commentId": comment_id, + "after": cur, + } + data = self.client.graphql(Q.DISCUSSION_REPLIES_QUERY, vars_) + node = (data.get("data") or {}).get("node") or {} + replies = node.get("replies") or {} + for r in replies.get("nodes") or []: + r["_owner"] = self.owner + r["_repo"] = self.name + r["_discussionNumber"] = disc_number + r["_commentId"] = comment_id + self.writers["discussion_extra_replies"].write(r) + info = replies.get("pageInfo") or {} + cur = info.get("endCursor") if info.get("hasNextPage") else None + + # ----- Commits ----- + def scrape_commits(self, branch: str = "refs/heads/main") -> int: + key = "commits" + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + return 0 + total_new = 0 + page = 0 + page_cap = 100 + trial_cap = self.trial_limits.get(key) + per_page = min(page_cap, trial_cap) if trial_cap and trial_cap > 0 else page_cap + while True: + page += 1 + vars_ = { + "owner": self.owner, + "name": self.name, + "first": per_page, + "after": cursor, + "branch": branch, + } + data = self.client.graphql(Q.COMMITS_PAGE_QUERY, vars_) + self._log_rate("commits", data) + ref = ((data.get("data") or {}).get("repository") or {}).get("ref") or {} + tgt = ref.get("target") or {} + hist = tgt.get("history") or {} + nodes = hist.get("nodes") or [] + for c in nodes: + c["_owner"] = self.owner + c["_repo"] = self.name + c["_fetchedAt"] = ts() + if self.writers[key].write(c): + total_new += 1 + info = hist.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + log.info( + "[%s/%s] commits page %d (+%d) remaining=%s", + self.owner, + self.name, + page, + len(nodes), + self.client.graphql_remaining, + ) + if self._trial_stop(key, total_new): + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + return total_new + + # ----- Releases/Labels/Milestones ----- + def scrape_releases(self) -> int: + return self._scrape_simple("releases", Q.RELEASES_QUERY, "releases") + + def scrape_labels(self) -> int: + return self._scrape_simple("labels", Q.LABELS_QUERY, "labels") + + def scrape_milestones(self) -> int: + return self._scrape_simple("milestones", Q.MILESTONES_QUERY, "milestones") + + def _scrape_simple(self, key: str, query: str, field: str) -> int: + cursor = self.state.get(f"{key}_cursor") + done = self.state.get(f"{key}_done", False) + if done: + return 0 + total_new = 0 + while True: + vars_ = { + "owner": self.owner, + "name": self.name, + "first": 50, + "after": cursor, + } + data = self.client.graphql(query, vars_) + repo = (data.get("data") or {}).get("repository") or {} + col = repo.get(field) or {} + for it in col.get("nodes") or []: + it["_owner"] = self.owner + it["_repo"] = self.name + it["_fetchedAt"] = ts() + if self.writers[key].write(it): + total_new += 1 + info = col.get("pageInfo") or {} + cursor = info.get("endCursor") + self.state.set(f"{key}_cursor", cursor) + if self._trial_stop(key, total_new): + return total_new + if not info.get("hasNextPage"): + self.state.set(f"{key}_done", True) + break + log.info("[%s/%s] %s done +%d", self.owner, self.name, key, total_new) + return total_new + + def close(self) -> None: + for w in self.writers.values(): + try: + w.close() + except Exception: + pass + + +def setup_logging(log_file: Path) -> None: + log_file.parent.mkdir(parents = True, exist_ok = True) + fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s" + handlers = [ + logging.StreamHandler(sys.stdout), + logging.FileHandler(log_file, mode = "a", encoding = "utf-8"), + ] + logging.basicConfig(level = logging.INFO, format = fmt, handlers = handlers, force = True) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--base-dir", default = "/mnt/disks/unslothai/ubuntu/workspace_34/github_scraper" + ) + ap.add_argument( + "--repos", nargs = "+", default = ["unslothai/unsloth", "unslothai/unsloth-zoo"] + ) + ap.add_argument("--trial", action = "store_true", help = "Small trial run") + ap.add_argument( + "--only", + nargs = "+", + default = None, + help = "Only run these resource keys: issues,pulls,discussions,commits,releases,labels,milestones,meta", + ) + ap.add_argument( + "--hf-upload-interval", + type = int, + default = 900, + help = "Seconds between HF uploads (0 to disable)", + ) + args = ap.parse_args() + + base = Path(args.base_dir) + data_dir = base / "data" + data_dir.mkdir(parents = True, exist_ok = True) + setup_logging(base / "logs" / f"scraper_{time.strftime('%Y%m%d_%H%M%S')}.log") + log.info("Scraper starting: repos=%s trial=%s", args.repos, args.trial) + + client = GitHubClient(min_remaining_graphql = 80, min_remaining_rest = 80) + rl = client.rate_snapshot() + log.info( + "Rate limit snapshot: %s", + json.dumps(rl.get("resources", {}), default = str)[:400], + ) + + # Start HF uploader in background if requested + uploader = None + if args.hf_upload_interval > 0: + from hf_uploader import HFUploader + + uploader = HFUploader(data_dir, interval_s = args.hf_upload_interval) + uploader.start() + + trial_limits = None + if args.trial: + trial_limits = { + "issues": 5, + "pull_requests": 5, + "discussions": 3, + "commits": 20, + "releases": 3, + "labels": 20, + "milestones": 20, + } + + only = set(args.only or []) + + try: + for repo_spec in args.repos: + owner, name = repo_spec.split("/") + scraper = RepoScraper(owner, name, data_dir, client, trial_limits) + try: + 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: + scraper.scrape_milestones() + if not only or "releases" in only: + scraper.scrape_releases() + if not only or "discussions" in only: + scraper.scrape_discussions() + if not only or "issues" in only: + scraper.scrape_issues() + if not only or "pulls" in only: + scraper.scrape_prs() + if not only or "commits" in only: + 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: + if uploader: + log.info("Stopping uploader and final sync...") + uploader.stop(final_upload = True) + log.info( + "Scraper complete. GraphQL calls=%d REST calls=%d", + client.calls_graphql, + client.calls_rest, + ) + + +if __name__ == "__main__": + main() diff --git a/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py new file mode 100644 index 0000000000..efa663db2f --- /dev/null +++ b/studio/backend/plugins/data-designer-github-repo-seed/src/data_designer_github_repo_seed/scraper_impl/state_store.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Checkpoint state management for resumable scraping.""" + +from __future__ import annotations + +import json +import os +import threading +from pathlib import Path +from typing import Any, Dict + + +class StateStore: + def __init__(self, path: str | Path): + self.path = Path(path) + self.path.parent.mkdir(parents = True, exist_ok = True) + self._lock = threading.Lock() + self._data: Dict[str, Any] = {} + if self.path.exists(): + try: + with self.path.open() as f: + self._data = json.load(f) + except Exception: + self._data = {} + + def get(self, key: str, default: Any = None) -> Any: + with self._lock: + return self._data.get(key, default) + + def set(self, key: str, value: Any) -> None: + with self._lock: + self._data[key] = value + self._flush() + + def update(self, key: str, **kwargs) -> None: + with self._lock: + sub = dict(self._data.get(key, {})) + sub.update(kwargs) + self._data[key] = sub + self._flush() + + def all(self) -> Dict[str, Any]: + with self._lock: + return dict(self._data) + + def _flush(self) -> None: + tmp = self.path.with_suffix(self.path.suffix + ".tmp") + with tmp.open("w") as f: + json.dump(self._data, f, indent = 2, default = str) + os.replace(tmp, self.path) + + +class JsonlWriter: + """Append-only JSONL writer, thread-safe, with line buffering.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + self.path.parent.mkdir(parents = True, exist_ok = True) + self._lock = threading.Lock() + self._fh = self.path.open("a", buffering = 1) + self._count_seen_keys: set[str] = set() + # Preload seen keys if file exists (for dedup across resumes) + if self.path.exists() and self.path.stat().st_size > 0: + try: + with self.path.open() as f: + for line in f: + try: + obj = json.loads(line) + k = self._key(obj) + if k is not None: + self._count_seen_keys.add(k) + except Exception: + pass + except Exception: + pass + + def _key(self, obj: dict) -> str | None: + for k in ("id", "node_id", "number", "sha", "url"): + if k in obj: + return f"{k}:{obj[k]}" + return None + + def has(self, key: str) -> bool: + return key in self._count_seen_keys + + def write(self, obj: dict) -> bool: + """Return True if newly written, False if already present.""" + k = self._key(obj) + with self._lock: + if k is not None and k in self._count_seen_keys: + return False + if k is not None: + self._count_seen_keys.add(k) + self._fh.write(json.dumps(obj, default = str, ensure_ascii = False)) + self._fh.write("\n") + self._fh.flush() + return True + + def close(self) -> None: + try: + self._fh.close() + except Exception: + pass 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/backend/routes/data_recipe/jobs.py b/studio/backend/routes/data_recipe/jobs.py index 606ef1832c..da6416e324 100644 --- a/studio/backend/routes/data_recipe/jobs.py +++ b/studio/backend/routes/data_recipe/jobs.py @@ -5,8 +5,9 @@ from __future__ import annotations -from datetime import timedelta -from typing import Any +import copy +from datetime import datetime, timedelta, timezone +from typing import Any, Optional from urllib.parse import urlparse from fastapi import APIRouter, HTTPException, Query, Request @@ -94,14 +95,111 @@ def _used_llm_model_aliases(recipe: dict[str, Any]) -> set[str]: return aliases -def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: +def _inject_local_structured_response_format( + recipe: dict[str, Any], local_provider_names: set[str] +) -> None: + """For each llm-structured column that targets a local-provider model_config, + clone the model_config and inject an OpenAI ``response_format`` with the + column's ``output_format`` JSON schema. The column is rewritten to point at + the clone so llm-text / llm-judge columns that share the same alias keep + free-form sampling. + + Without this, data_designer only injects a prompt-level "return JSON in a + ```json fence" instruction. Small GGUF models frequently break format, + wasting the full ``max_tokens`` budget per row and then failing to parse. + Forwarding ``response_format`` lets llama-server apply grammar-constrained + sampling from the JSON schema, which guarantees a parseable response and + terminates early. + """ + columns = recipe.get("columns") + model_configs = recipe.get("model_configs") + if not isinstance(columns, list) or not isinstance(model_configs, list): + return + + # alias -> model_config (only configs referencing a local provider qualify). + alias_to_local_mc: dict[str, dict[str, Any]] = {} + for mc in model_configs: + if not isinstance(mc, dict): + continue + if mc.get("provider") in local_provider_names and isinstance( + mc.get("alias"), str + ): + alias_to_local_mc[mc["alias"]] = mc + + if not alias_to_local_mc: + return + + # Clone per (alias, column) so each llm-structured column gets its own + # schema without leaking response_format onto other columns that share the + # same base alias. + seen_clone_aliases: set[str] = { + mc.get("alias") for mc in model_configs if isinstance(mc.get("alias"), str) + } + new_configs: list[dict[str, Any]] = [] + for column in columns: + if not isinstance(column, dict): + continue + if column.get("column_type") != "llm-structured": + continue + alias = column.get("model_alias") + if not isinstance(alias, str) or alias not in alias_to_local_mc: + continue + output_format = column.get("output_format") + if not isinstance(output_format, dict) or not output_format: + continue + base_mc = alias_to_local_mc[alias] + column_name = column.get("name") or "structured" + clone_alias_base = f"{alias}__{column_name}_structured" + clone_alias = clone_alias_base + counter = 1 + while clone_alias in seen_clone_aliases: + counter += 1 + clone_alias = f"{clone_alias_base}_{counter}" + seen_clone_aliases.add(clone_alias) + + clone = copy.deepcopy(base_mc) + clone["alias"] = clone_alias + params = clone.get("inference_parameters") + if not isinstance(params, dict): + params = {} + clone["inference_parameters"] = params + # data_designer's BaseInferenceParams is a pydantic model with + # extra="forbid", so response_format cannot sit at the top level of + # inference_parameters. It does expose an `extra_body: dict` pass- + # through that the OpenAI client spreads into the request body at the + # top level, which is where llama-server reads response_format from. + # llama.cpp server shape (tools/server/README.md): the schema sits + # directly under response_format, not nested in a json_schema object + # the way OpenAI's Chat Completions API expects. llama-server converts + # the schema to a GBNF grammar and applies it during sampling. + extra_body = params.get("extra_body") + if not isinstance(extra_body, dict): + extra_body = {} + extra_body["response_format"] = { + "type": "json_schema", + "schema": output_format, + } + params["extra_body"] = extra_body + new_configs.append(clone) + column["model_alias"] = clone_alias + + if new_configs: + model_configs.extend(new_configs) + + +def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optional[int]: """ Mutate recipe dict in-place: for any provider with is_local=True, - generate a JWT and fill in the endpoint pointing at this server. + fill in the endpoint pointing at this server and inject a short-lived + internal sk-unsloth-* API key for workflow auth. + + Returns the row id of the minted internal key (so the caller can + revoke it on job completion) or ``None`` when no local provider is + actually reachable from an LLM column. """ providers = recipe.get("model_providers") if not providers: - return + return None # Collect local providers and pop is_local from ALL dicts unconditionally. # Strict `is True` guard so malformed payloads (is_local: 1, @@ -115,7 +213,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: local_indices.append(i) if not local_indices: - return + return None endpoint = _resolve_local_v1_endpoint(request) @@ -138,6 +236,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: } token = "" + internal_key_id: Optional[int] = None if local_names & referenced_providers: # Verify a model is loaded. # NOTE: This is a point-in-time check (TOCTOU). The model could be unloaded @@ -158,18 +257,21 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: "No model loaded in Chat. Load a model first, then run the recipe." ) - from auth.authentication import ( - create_access_token, - ) # deferred: avoids circular import + from auth import storage # deferred: avoids circular import - # Uses the "unsloth" admin subject. If the user changes their password, - # the JWT secret rotates and this token becomes invalid mid-run. - # Acceptable for v1 - recipes typically finish well within one session. - token = create_access_token( - subject = "unsloth", - expires_delta = timedelta(hours = 24), - desktop = _request_has_desktop_access_token(request), + # Mint an internal sk-unsloth-* key scoped to this workflow run. + # Uses the unified API-key issuance path (one mint/revoke/verify + # surface instead of a second JWT code path). The key is marked + # internal so it is hidden from the user's API-key list, and the + # caller revokes it when the job terminates. + expires_at = (datetime.now(timezone.utc) + timedelta(hours = 24)).isoformat() + token, row = storage.create_api_key( + username = "unsloth", + name = "data-recipe workflow", + expires_at = expires_at, + internal = True, ) + internal_key_id = int(row["id"]) # Defensively strip any stale "external"-only fields the frontend may # have left on the dict (extra_headers/extra_body/api_key_env). The UI @@ -196,6 +298,37 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> None: continue if mc.get("provider") in local_names: mc["skip_health_check"] = True + # Disable thinking for data-recipe inference on local providers. + # Reasoning models emit a ... preamble before the + # answer, which roughly doubles generated token count per row and + # pushes the visible answer past data_designer's json-fence + # regex. Forward chat_template_kwargs={enable_thinking: False} + # through the OpenAI SDK's extra_body passthrough so llama-server + # renders the template without the reasoning preamble. Free-form + # llm-text columns benefit from the latency cut, and structured + # columns also stop leaking think tags into the grammar- + # constrained JSON (llama-server's GBNF path still enforces the + # schema either way). + params = mc.get("inference_parameters") + if not isinstance(params, dict): + params = {} + mc["inference_parameters"] = params + extra_body = params.get("extra_body") + if not isinstance(extra_body, dict): + extra_body = {} + tpl_kwargs = extra_body.get("chat_template_kwargs") + if not isinstance(tpl_kwargs, dict): + tpl_kwargs = {} + tpl_kwargs.setdefault("enable_thinking", False) + extra_body["chat_template_kwargs"] = tpl_kwargs + params["extra_body"] = extra_body + + # Forward each llm-structured column's output_format as an OpenAI + # response_format so llama-server uses grammar-constrained sampling and + # small GGUFs stop wasting the full max_tokens budget on broken JSON. + _inject_local_structured_response_format(recipe, local_names) + + return internal_key_id def _normalize_run_name(value: Any) -> str | None: @@ -240,21 +373,49 @@ def create_job(payload: RecipePayload, request: Request): ) from exc try: - _inject_local_providers(recipe, request) + internal_api_key_id = _inject_local_providers(recipe, request) except ValueError as exc: raise HTTPException(status_code = 400, detail = str(exc)) from exc - mgr = get_job_manager() + # Single try block covers get_job_manager() AND mgr.start() so a workflow + # key minted above never outlives the request even when an unexpected + # exception type (TypeError from a stale kwarg, OSError from a queue + # write, etc.) bubbles up. Without the bare except, such exceptions let + # the sk-unsloth-* key live until its 24h TTL. try: - job_id = mgr.start(recipe = recipe, run = run) + mgr = get_job_manager() + job_id = mgr.start( + recipe = recipe, + run = run, + internal_api_key_id = internal_api_key_id, + ) except RuntimeError as exc: + if internal_api_key_id is not None: + _revoke_internal_api_key_safe(internal_api_key_id) raise HTTPException(status_code = 409, detail = str(exc)) from exc except ValueError as exc: + if internal_api_key_id is not None: + _revoke_internal_api_key_safe(internal_api_key_id) raise HTTPException(status_code = 400, detail = str(exc)) from exc + except Exception: + if internal_api_key_id is not None: + _revoke_internal_api_key_safe(internal_api_key_id) + raise return {"job_id": job_id} +def _revoke_internal_api_key_safe(key_id: int) -> None: + """Best-effort revoke of a workflow-minted key; swallow any error so + that revocation failures never mask the caller's own error path.""" + try: + from auth import storage # deferred: avoids circular import + + storage.revoke_internal_api_key(key_id) + except Exception: + pass + + @router.get("/jobs/{job_id}/status") def job_status(job_id: str): mgr = get_job_manager() diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index e9cf828610..91cf718e6e 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -8,6 +8,7 @@ from __future__ import annotations import base64 import binascii import json +import os import re from itertools import islice from pathlib import Path @@ -627,3 +628,14 @@ def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectRespons split = None, subset = None, ) + + +@router.get("/seed/github/env-token") +def get_github_env_token_status() -> dict: + """Report whether the server has a GH_TOKEN / GITHUB_TOKEN env var. + + The value is never returned; the UI uses this to tell the user they + can leave the token field blank. + """ + has_token = bool(os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")) + return {"has_token": has_token} diff --git a/studio/backend/routes/data_recipe/validate.py b/studio/backend/routes/data_recipe/validate.py index 555e3eaa06..87eef939b4 100644 --- a/studio/backend/routes/data_recipe/validate.py +++ b/studio/backend/routes/data_recipe/validate.py @@ -18,6 +18,57 @@ from models.data_recipe import RecipePayload, ValidateError, ValidateResponse router = APIRouter() +_GITHUB_VALIDATE_NOTE = "Recipe shape is valid. GitHub access and rate limits are checked when the run starts." +_GITHUB_ITEM_TYPES = {"issues", "pulls", "commits"} + + +def _github_seed_source(recipe: dict[str, Any]) -> dict[str, Any] | None: + seed_config = recipe.get("seed_config") + if not isinstance(seed_config, dict): + return None + source = seed_config.get("source") + if not isinstance(source, dict) or source.get("seed_type") != "github_repo": + return None + return source + + +def _validate_github_seed_static(source: dict[str, Any]) -> list[ValidateError]: + errors: list[ValidateError] = [] + + repos = source.get("repos") + if not isinstance(repos, list) or not repos: + errors.append(ValidateError(message = "GitHub seed requires at least one repo.")) + else: + for repo in repos: + if not isinstance(repo, str) or not repo.strip() or "/" not in repo: + errors.append( + ValidateError(message = "GitHub repos must be owner/name strings.") + ) + break + + item_types = source.get("item_types") + if not isinstance(item_types, list) or not item_types: + errors.append( + ValidateError(message = "GitHub seed requires at least one item type.") + ) + else: + invalid_items = [item for item in item_types if item not in _GITHUB_ITEM_TYPES] + if invalid_items: + errors.append( + ValidateError( + message = "GitHub item types must be issues, pulls, or commits." + ) + ) + + try: + limit = int(source.get("limit")) + except (TypeError, ValueError): + limit = 0 + if limit < 1 or limit > 5000: + errors.append(ValidateError(message = "GitHub limit must be from 1 to 5000.")) + + return errors + def _collect_validation_errors(recipe: dict[str, Any]) -> list[ValidateError]: try: @@ -93,6 +144,22 @@ def validate(payload: RecipePayload) -> ValidateResponse: _patch_local_providers(recipe) + github_source = _github_seed_source(recipe) + if github_source is not None: + static_errors = _validate_github_seed_static(github_source) + if static_errors: + return ValidateResponse(valid = False, errors = static_errors) + try: + build_config_builder(recipe) + except Exception as exc: + detail = str(exc).strip() or "Validation failed." + return ValidateResponse( + valid = False, + errors = [ValidateError(message = detail)], + raw_detail = detail, + ) + return ValidateResponse(valid = True, raw_detail = _GITHUB_VALIDATE_NOTE) + try: validate_recipe(recipe) except RuntimeError as exc: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index cf3b37a2fd..b13eb08967 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1286,6 +1286,20 @@ async def openai_chat_completions( llama_backend = get_llama_cpp_backend() using_gguf = llama_backend.is_loaded + # OpenAI-SDK clients send ``chat_template_kwargs`` via ``extra_body``, + # which the SDK spreads into the request body at the top level. Studio's + # ChatCompletionRequest has ``extra="allow"`` so pydantic stashes them in + # ``model_extra``, but the typed ``payload.enable_thinking`` path is what + # downstream generators actually consume. Lift ``enable_thinking`` from + # the extra-body chat_template_kwargs onto the typed field so clients + # that only know the OpenAI shape (data_designer recipe runs, etc.) + # can still control the reasoning preamble. + _extra = getattr(payload, "model_extra", None) + if payload.enable_thinking is None and isinstance(_extra, dict): + _tpl_kw = _extra.get("chat_template_kwargs") + if isinstance(_tpl_kw, dict) and "enable_thinking" in _tpl_kw: + payload.enable_thinking = bool(_tpl_kw["enable_thinking"]) + # ── Determine which backend is active ───────────────────── if using_gguf: model_name = llama_backend.model_identifier or payload.model @@ -1440,11 +1454,22 @@ async def openai_chat_completions( # carry `tool_calls` (content=None) — both of which are valid in # multi-turn client-side tool loops. _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) + # Route guided-decoding requests through the verbatim passthrough so + # ``response_format`` (JSON schema) actually reaches llama-server and + # the model's GBNF-constrained output comes back unmodified. The + # non-passthrough GGUF path below calls ``generate_chat_completion`` + # which has no response_format kwarg, so the schema gets silently + # dropped and data_designer falls back to free-form sampling. Guided + # decoding does not require ``supports_tools`` - the grammar machinery + # is independent of tool-call parsing. + _has_response_format = _extract_response_format(payload) is not None + _tools_passthrough = llama_backend.supports_tools and ( + (payload.tools and len(payload.tools) > 0) or _has_tool_messages + ) if ( using_gguf - and llama_backend.supports_tools and not payload.enable_tools - and ((payload.tools and len(payload.tools) > 0) or _has_tool_messages) + and (_tools_passthrough or _has_response_format) ): # Preserve the vision guard that would otherwise run in the # non-passthrough path below: text-only tool-capable GGUFs @@ -3652,6 +3677,8 @@ def _build_passthrough_payload( repetition_penalty = None, presence_penalty = None, tool_choice = "auto", + response_format = None, + chat_template_kwargs = None, backend_ctx = None, ): body = { @@ -3680,6 +3707,17 @@ def _build_passthrough_payload( body["repeat_penalty"] = repetition_penalty if presence_penalty is not None: body["presence_penalty"] = presence_penalty + if response_format is not None: + # llama-server applies a GBNF grammar derived from the JSON schema + # when response_format is present. Field is documented flat at the + # request root (tools/server/README.md), which is also what the + # OpenAI SDK produces by spreading extra_body into the body top. + body["response_format"] = response_format + if chat_template_kwargs is not None: + # Propagate reasoning / template overrides (e.g. enable_thinking) + # so llama-server renders the Jinja template in the mode the caller + # asked for instead of whatever default the model was loaded with. + body["chat_template_kwargs"] = chat_template_kwargs return body @@ -3990,6 +4028,20 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages +def _extract_response_format(payload): + """Return the ``response_format`` field on an incoming ChatCompletionRequest + (or None). The model is declared with ``extra="allow"`` so pydantic stashes + unknown top-level fields in ``model_extra``; OpenAI-SDK clients spread + ``extra_body`` into the request body top level, which is where guided- + decoding recipes park their JSON-schema response_format. + """ + extra = getattr(payload, "model_extra", None) + if not isinstance(extra, dict): + return None + rf = extra.get("response_format") + return rf if isinstance(rf, dict) else None + + def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: """Assemble the llama-server request body from a ChatCompletionRequest. @@ -3999,6 +4051,12 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: """ messages = _openai_messages_for_passthrough(payload) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" + # When the caller asked for a specific reasoning mode, forward it to + # llama-server via chat_template_kwargs so the Jinja template renders + # with (or without) the reasoning preamble. + tpl_kwargs = None + if payload.enable_thinking is not None: + tpl_kwargs = {"enable_thinking": bool(payload.enable_thinking)} return _build_passthrough_payload( messages, payload.tools, @@ -4012,6 +4070,8 @@ def _build_openai_passthrough_body(payload, backend_ctx = None) -> dict: repetition_penalty = payload.repetition_penalty, presence_penalty = payload.presence_penalty, tool_choice = tool_choice, + response_format = _extract_response_format(payload), + chat_template_kwargs = tpl_kwargs, backend_ctx = backend_ctx, ) @@ -4214,6 +4274,41 @@ async def _openai_passthrough_non_streaming( detail = f"llama-server error: {resp.text[:500]}", ) + # Guided-decoding fence wrap. llama-server returns raw JSON that matches + # the schema (no surrounding markdown) because the GBNF grammar only + # emits the JSON object itself. data_designer's llm-structured parser + # looks for a ```json ... ``` markdown fence and discards unfenced + # output, which collapses a 100%-valid guided-decoding run to 0/N. + # Wrap each choice's content in the expected fence when the caller + # asked for guided decoding, leaving already-fenced content alone. + if _extract_response_format(payload) is not None: + try: + data = resp.json() + changed = False + for choice in data.get("choices", []): + if not isinstance(choice, dict): + continue + msg = choice.get("message") + if not isinstance(msg, dict): + continue + content = msg.get("content") + if not isinstance(content, str): + continue + stripped = content.strip() + if not stripped or stripped.startswith("```"): + continue + msg["content"] = f"```json\n{stripped}\n```" + changed = True + if changed: + return JSONResponse(content = data) + except Exception as exc: + # Wrap is best-effort; fall through to the verbatim body if + # the response is not JSON-shaped or the structure is unusual. + logger.warning( + "response_format fence wrap skipped: %s", + exc, + ) + # Pass the upstream body through as raw bytes — skips a redundant # parse+re-serialize round-trip and keeps the response truly # verbatim (matches the docstring). Status is guaranteed 200 by diff --git a/studio/backend/tests/test_data_recipe_github_progress.py b/studio/backend/tests/test_data_recipe_github_progress.py new file mode 100644 index 0000000000..8e8c3995f4 --- /dev/null +++ b/studio/backend/tests/test_data_recipe_github_progress.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from core.data_recipe.jobs.parse import apply_update, parse_log_message +from core.data_recipe.jobs.types import Job +from routes.data_recipe.validate import _GITHUB_VALIDATE_NOTE, validate +from models.data_recipe import RecipePayload + + +def test_github_page_log_updates_source_progress_without_cursor(): + job = Job(job_id = "job-1") + job.source_progress_estimated_total = 200 + + update = parse_log_message( + "[unslothai/unsloth] issues page 2 (+15) cursor=abc123 remaining=2960" + ) + + assert update is not None + apply_update(job, update) + + progress = job.source_progress + assert progress is not None + assert progress.source == "github" + assert progress.status == "fetching" + assert progress.repo == "unslothai/unsloth" + assert progress.resource == "issues" + assert progress.page == 2 + assert progress.page_items == 15 + assert progress.fetched_items == 15 + assert progress.estimated_total == 200 + assert progress.rate_remaining == 2960 + assert progress.message is not None + assert "cursor" not in progress.message + assert "abc123" not in progress.message + + +def test_github_rate_limit_log_updates_source_progress(): + job = Job(job_id = "job-1") + + update = parse_log_message("Rate limit hit. Sleeping 123s until reset.") + + assert update is not None + apply_update(job, update) + + progress = job.source_progress + assert progress is not None + assert progress.status == "rate_limited" + assert progress.retry_after_sec == 123 + assert "resume automatically" in (progress.message or "") + + +def test_github_real_sample_prs_and_trial_limit_are_parsed(): + job = Job(job_id = "job-1") + + for message in ( + "[unslothai/unsloth] PRs page 4 (+25) cursor=abc123 remaining=4983", + "Trial limit reached for PRs (100)", + ): + update = parse_log_message(message) + assert update is not None + apply_update(job, update) + + progress = job.source_progress + assert progress is not None + assert progress.repo == "unslothai/unsloth" + assert progress.resource == "pulls" + assert progress.page == 4 + assert progress.fetched_items == 25 + assert progress.rate_remaining == 4983 + assert progress.message == "GitHub pulls trial limit reached (100)." + + +def test_github_validate_skips_live_access_with_honest_note(): + response = validate( + RecipePayload( + recipe = { + "seed_config": { + "source": { + "seed_type": "github_repo", + "repos": ["unslothai/unsloth"], + "item_types": ["issues"], + "limit": 1, + } + }, + "columns": [{"column_type": "expression", "name": "x", "expr": "1"}], + } + ) + ) + + assert response.valid is True + assert response.raw_detail == _GITHUB_VALIDATE_NOTE diff --git a/studio/frontend/src/features/data-recipes/learning-recipes/github-support-bot.json b/studio/frontend/src/features/data-recipes/learning-recipes/github-support-bot.json new file mode 100644 index 0000000000..b984f84d2f --- /dev/null +++ b/studio/frontend/src/features/data-recipes/learning-recipes/github-support-bot.json @@ -0,0 +1,238 @@ +{ + "recipe": { + "model_providers": [ + { + "name": "Local Model", + "endpoint": "", + "provider_type": "openai", + "extra_headers": {}, + "extra_body": {}, + "is_local": true + } + ], + "mcp_providers": [], + "model_configs": [ + { + "alias": "model_1", + "model": "unsloth/gemma-4-E2B-it-GGUF", + "provider": "Local Model", + "inference_parameters": { + "temperature": 0.4, + "max_tokens": 800 + } + } + ], + "seed_config": { + "source": { + "seed_type": "github_repo", + "repos": [ + "unslothai/unsloth", + "unslothai/unsloth-zoo" + ], + "item_types": [ + "issues", + "pulls" + ], + "limit": 100, + "include_comments": true, + "max_comments_per_item": 20 + }, + "sampling_strategy": "shuffle", + "selection_strategy": null + }, + "tool_configs": [], + "columns": [ + { + "column_type": "llm-text", + "name": "User", + "drop": false, + "model_alias": "model_1", + "prompt": "Read the GitHub {{ item_type }} below and write ONE realistic user request that could have produced it. Imagine a developer asking a GitHub co-author model to either file this {{ item_type }} or draft a PR that resolves it. Use first-person imperative phrasing (\"Open an issue...\", \"Draft a PR that...\", \"Investigate why...\"). Preserve concrete technical details (model names, flags, file paths, tracebacks) that appear in the thread. Keep it 1-3 sentences. Output ONLY the user request, no preamble.\n\n--- INPUT ---\nRepo: {{ repo }}\nType: {{ item_type }}\nTitle: {{ title }}\nBody:\n{{ body }}\n\nFirst comments:\n{{ comments }}", + "system_prompt": "You invert real GitHub threads into the user request that would have produced them. Faithful to the thread, no invented facts, no em-dashes, no emojis.", + "with_trace": "none", + "extract_reasoning_content": false + }, + { + "column_type": "llm-structured", + "name": "Assistant", + "drop": false, + "model_alias": "model_1", + "prompt": "You are generating one training row for an Unsloth GitHub co-author model. Given the real GitHub thread and a synthesized user request, produce a grounded structured response.\n\nSource thread:\n- Repo: {{ repo }}\n- Type: {{ item_type }}\n- Title: {{ title }}\n- URL: {{ url }}\n- State: {{ state }}\n- Labels: {{ labels }}\n- Body: {{ body }}\n- First comments: {{ comments }}\n\nUser request:\n{{ User }}\n\nRules:\n- `response`: 100-250 words of Markdown grounded in the thread. If the thread is a closed / resolved issue, follow the `issue_fix_plan` shape: brief diagnosis, numbered fix steps, and a short repro. If the thread is a PR, follow the `explain_pr` shape: what changed, why, and which files or symbols were touched. If the thread is open / unresolved, answer honestly and ask for the missing info.\n- Cite the source URL at least once inline as `[source: {{ url }}]`.\n- Name at least one concrete symbol (function, class, flag, env var, or file path) from the thread when available.\n- Include a short ```bash or ```python code block ONLY if the thread itself contains that code or command.\n- Never recommend `rm -rf`, force push, or other destructive commands without an explicit warning.\n- No em-dashes, no emojis, no AI-disclaimer phrases. Only cite URLs / paths that appear in the thread.\n- `followups`: 0-4 follow-up questions when the thread is missing info (versions, GPU, traceback). Empty list if the response is complete.\n- `cites`: URLs / file paths actually used. Always include `{{ url }}`.\n- `task`: one of `explain_pr`, `issue_fix_plan`, `issue_solution`, `discussion_qa`. Pick the closest match.\n- `confidence`: `high` / `medium` / `low`. Use `low` when ambiguous or out of scope.", + "system_prompt": "You write grounded GitHub co-author responses for Unsloth. Faithful to the thread, no invented facts, no em-dashes, no emojis, no AI-disclaimer phrases.", + "with_trace": "none", + "extract_reasoning_content": false, + "output_format": { + "type": "object", + "properties": { + "response": { + "type": "string", + "minLength": 1 + }, + "followups": { + "type": "array", + "items": { + "type": "string" + }, + "maxItems": 4 + }, + "cites": { + "type": "array", + "items": { + "type": "string" + }, + "maxItems": 6 + }, + "task": { + "type": "string", + "enum": [ + "explain_pr", + "issue_fix_plan", + "issue_solution", + "discussion_qa" + ] + }, + "confidence": { + "type": "string", + "enum": [ + "high", + "medium", + "low" + ] + } + }, + "required": [ + "response", + "followups", + "cites", + "task", + "confidence" + ], + "additionalProperties": false + } + } + ], + "processors": [] + }, + "run": { + "rows": 5, + "preview": true, + "output_formats": [ + "jsonl" + ] + }, + "ui": { + "nodes": [ + { + "id": "note_1", + "x": 460.4272962489311, + "y": 123.10114554082236, + "width": 400, + "node_type": "markdown_note", + "name": "note_1", + "markdown": "### GitHub Crawler\nReal GitHub issues and PRs turned into `{User, Assistant}` training pairs. Mirrors two of the eleven canonical enrichment tasks in the `github_data_gatherer` dataset: `pr_requests_20` / `issue_requests_20` for the input side, and `explain_pr` / `issue_fix_plan` / `issue_solution` for the output side.\n\n**Click `Run` below for 10 sample rows.** Defaults point at `unslothai/unsloth` + `unslothai/unsloth-zoo`, use the server's `GH_TOKEN` / `GITHUB_TOKEN` env var, and run the bundled local model. Paste a PAT only when you need private repos.\n\n**Configure source data**\n- Paste `owner/name` values or GitHub URLs; the editor normalizes and dedupes them.\n- Keep the seed `limit` around 100 for previews. Increase toward 5000 per item type for larger backfills.\n- Comments and large repos can make Check or Run take minutes; watch logs for GitHub page and rate-limit messages.\n\n**Upgrade to production**\n- Swap `unsloth/gemma-4-E2B-it-GGUF` for a stronger model (`gpt-5.4-mini` with `reasoning_effort=medium` is what the reference dataset uses).\n- Replace the demo prompts with the task-specific prompts from the reference dataset (see Note 3).\n- Raise `max_parallel_requests` to 4 once the inference server can handle it.", + "note_color": "#E0F2FE", + "note_opacity": "35" + }, + { + "id": "note_2", + "x": -514.4586648521598, + "y": 869.2785862193363, + "width": 400, + "node_type": "markdown_note", + "name": "note_2", + "markdown": "The **User** column inverts each GitHub thread into a realistic request a developer would give a co-author model (`\"Draft a PR that...\"`, `\"Investigate why...\"`). Same shape as the `pr_requests_20` and `issue_requests_20` enrichments.\n\nTweak the prompt to:\n- always keep the traceback verbatim\n- vary persona (newcomer, maintainer, ops)\n- split one thread into multiple alternative phrasings for data augmentation.", + "note_color": "#E0F2FE", + "note_opacity": "35" + }, + { + "id": "note_3", + "x": -536.3163554442993, + "y": -84.18780704666585, + "width": 400, + "node_type": "markdown_note", + "name": "note_3", + "markdown": "The **Assistant** block emits `{response, followups, cites, task, confidence}` and branches on thread type: closed issues become `issue_fix_plan` rows, PRs become `explain_pr` rows, everything else becomes `issue_solution` or `discussion_qa`.\n\n**Demo default**: 100-250 word response, one inline `[source: ]` cite, `max_parallel_requests=1` so a small local model stays stable.\n\n**Production prompt (paste in):**\n- Match the reference dataset's per-task prompts (`explain_pr`, `issue_fix_plan`, `pr_review_critique`, `pr_test_plan`, etc.).\n- Require 300 words for explanations, 6-12 bullets for test plans.\n- Enforce named symbols (function / flag / env var / file path).\n- Code fences only for content already in the thread.\n- Reject rows with em-dashes, emojis, or AI-disclaimer phrases.\n\nSee the `github_data_gatherer` dataset card for the full task catalog and the codex prompts used to train the reference GitHub model.", + "note_color": "#E0F2FE", + "note_opacity": "35" + }, + { + "id": "seed", + "x": 0, + "y": 140, + "width": 400 + }, + { + "id": "Local Model", + "x": -1056, + "y": 520, + "width": 400 + }, + { + "id": "model_1", + "x": -544, + "y": 488, + "width": 400 + }, + { + "id": "User", + "x": 0, + "y": 440, + "width": 400 + }, + { + "id": "Assistant", + "x": 0, + "y": 740, + "width": 400 + } + ], + "edges": [ + { + "from": "seed", + "to": "User", + "type": "canvas", + "source_handle": "data-out-bottom", + "target_handle": "data-in-top" + }, + { + "from": "User", + "to": "Assistant", + "type": "canvas", + "source_handle": "data-out-bottom", + "target_handle": "data-in-top" + }, + { + "from": "Local Model", + "to": "model_1", + "type": "semantic", + "source_handle": "semantic-out", + "target_handle": "semantic-in" + }, + { + "from": "model_1", + "to": "User", + "type": "semantic", + "source_handle": "semantic-out", + "target_handle": "data-in" + }, + { + "from": "model_1", + "to": "Assistant", + "type": "semantic", + "source_handle": "semantic-out-bottom", + "target_handle": "data-in" + } + ], + "layout_direction": "LR", + "seed_source_type": "github_repo", + "seed_columns": [], + "seed_drop_columns": [], + "seed_preview_rows": [], + "local_file_name": "", + "unstructured_file_ids": [], + "unstructured_file_names": [], + "unstructured_file_sizes": [], + "unstructured_chunk_size": "1200", + "unstructured_chunk_overlap": "200" + } +} \ No newline at end of file diff --git a/studio/frontend/src/features/data-recipes/learning-recipes/index.ts b/studio/frontend/src/features/data-recipes/learning-recipes/index.ts index d7a7e66e0a..8607d3faeb 100644 --- a/studio/frontend/src/features/data-recipes/learning-recipes/index.ts +++ b/studio/frontend/src/features/data-recipes/learning-recipes/index.ts @@ -19,6 +19,10 @@ const ocrDocumentExtractionUrl = new URL( "./ocr-document-extraction.json", import.meta.url, ).href; +const githubSupportBotUrl = new URL( + "./github-support-bot.json", + import.meta.url, +).href; function isRecord(value: unknown): value is Record { return !!value && typeof value === "object" && !Array.isArray(value); @@ -137,4 +141,11 @@ export const LEARNING_RECIPES: LearningRecipeDef[] = [ "Use image context to generate OCR-style document extraction output.", loadPayload: () => loadPayloadFromUrl(ocrDocumentExtractionUrl), }, + { + id: "github-support-bot", + title: "GitHub Crawler", + description: + "Crawl real GitHub issues and PRs and turn each thread into a {User, Assistant} training pair.", + loadPayload: () => loadPayloadFromUrl(githubSupportBotUrl), + }, ]; diff --git a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx index 9148b9e0da..0a51f9df70 100644 --- a/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx +++ b/studio/frontend/src/features/data-recipes/pages/data-recipes-page.tsx @@ -35,6 +35,7 @@ import { Delete02Icon, DocumentAttachmentIcon, FunctionIcon, + GithubIcon, Plant01Icon, PlusSignIcon, } from "@hugeicons/core-free-icons"; @@ -162,6 +163,22 @@ const TEMPLATE_CARDS: TemplateCard[] = [ ], learningRecipeId: "structured-outputs-jinja", }, + { + title: "GitHub Crawler", + description: + "Crawl real GitHub issues and PRs and invert each thread into a {User, Assistant} training pair.", + icon: GithubIcon, + difficulty: "Intermediate", + learningBadges: ["GitHub", "LLM Text", "Structured LLM"], + surfaceClassName: + "from-slate-500/15 via-zinc-500/5 to-transparent dark:from-slate-400/30 dark:via-zinc-400/14 dark:to-slate-950/16", + shineColor: [ + "rgb(71 85 105 / 0.45)", + "rgb(100 116 139 / 0.4)", + "rgb(148 163 184 / 0.45)", + ], + learningRecipeId: "github-support-bot", + }, ]; const LEARNING_RECIPE_BY_ID = new Map( diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 9212e4db9b..273d4aea8d 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -27,6 +27,28 @@ export type PublishRecipeJobResponse = { message: string; }; +export type SourceProgressResponse = { + source?: string | null; + status?: string | null; + repo?: string | null; + resource?: string | null; + page?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + page_items?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + fetched_items?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + estimated_total?: number | null; + percent?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + rate_remaining?: number | null; + // biome-ignore lint/style/useNamingConvention: api schema + retry_after_sec?: number | null; + message?: string | null; + // biome-ignore lint/style/useNamingConvention: api schema + updated_at?: number | null; +}; + export type JobStatusResponse = { // biome-ignore lint/style/useNamingConvention: api schema job_id: string; @@ -62,6 +84,8 @@ export type JobStatusResponse = { failed?: number | null; }; // biome-ignore lint/style/useNamingConvention: api schema + source_progress?: SourceProgressResponse | null; + // biome-ignore lint/style/useNamingConvention: api schema model_usage?: Record; rows?: number | null; cols?: number | null; @@ -183,12 +207,7 @@ async function parseErrorResponse(response: Response): Promise { // biome-ignore lint/style/useNamingConvention: api schema raw_detail?: string; }; - return ( - parsed.detail ?? - parsed.message ?? - parsed.raw_detail ?? - text - ); + return parsed.detail ?? parsed.message ?? parsed.raw_detail ?? text; } catch { return text; } @@ -264,11 +283,15 @@ export async function validateRecipe( return postJson("/validate", payload); } -export async function createRecipeJob(payload: unknown): Promise { +export async function createRecipeJob( + payload: unknown, +): Promise { return postJson("/jobs", payload); } -export async function getRecipeJobStatus(jobId: string): Promise { +export async function getRecipeJobStatus( + jobId: string, +): Promise { return getJson(`/jobs/${jobId}/status`); } @@ -292,7 +315,9 @@ export async function getRecipeJobDataset( ); } -export async function cancelRecipeJob(jobId: string): Promise { +export async function cancelRecipeJob( + jobId: string, +): Promise { return postJson(`/jobs/${jobId}/cancel`, {}); } @@ -315,6 +340,13 @@ export async function inspectSeedUpload( return postJson("/seed/inspect-upload", payload); } +// biome-ignore lint/style/useNamingConvention: api schema +export type GithubEnvTokenStatus = { has_token: boolean }; + +export async function getGithubEnvTokenStatus(): Promise { + return getJson("/seed/github/env-token"); +} + export async function listMcpTools( payload: McpToolsListRequest, ): Promise { @@ -407,11 +439,14 @@ export async function uploadUnstructuredFile( formData.append("existing_file_ids", existingFileIds.join(",")); } - const res = await authFetch(`${DATA_DESIGNER_API_BASE}/seed/upload-unstructured-file`, { - method: "POST", - body: formData, - signal, - }); + const res = await authFetch( + `${DATA_DESIGNER_API_BASE}/seed/upload-unstructured-file`, + { + method: "POST", + body: formData, + signal, + }, + ); if (res.status === 413) { const detail = await res.json().catch(() => ({ detail: "File too large" })); @@ -420,13 +455,16 @@ export async function uploadUnstructuredFile( filename: file.name, size_bytes: file.size, status: "error", - error: typeof detail.detail === "string" ? detail.detail : "File too large", + error: + typeof detail.detail === "string" ? detail.detail : "File too large", }; } if (!res.ok) { const detail = await res.json().catch(() => ({ detail: "Upload failed" })); - throw new Error(typeof detail.detail === "string" ? detail.detail : "Upload failed"); + throw new Error( + typeof detail.detail === "string" ? detail.detail : "Upload failed", + ); } return res.json(); diff --git a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts index 363cd85be5..ed58d1418d 100644 --- a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts +++ b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts @@ -12,6 +12,7 @@ import { EqualSignIcon, FingerPrintIcon, FunctionIcon, + GithubIcon, Plug01Icon, Parabola02Icon, PencilEdit02Icon, @@ -58,11 +59,16 @@ export type BlockType = | "seed_hf" | "seed_local" | "seed_unstructured" + | "seed_github" | "model_provider" | "model_config" | "tool_config"; -export type SeedBlockType = "seed_hf" | "seed_local" | "seed_unstructured"; +export type SeedBlockType = + | "seed_hf" + | "seed_local" + | "seed_unstructured" + | "seed_github"; type IconType = typeof CodeIcon; @@ -169,6 +175,15 @@ const BLOCK_DEFINITIONS: BlockDefinition[] = [ dialogKey: "seed", createConfig: (id, existing) => makeSeedConfig(id, existing, "unstructured"), }, + { + kind: "seed", + type: "seed_github", + title: "GitHub repositories", + description: "Crawl issues, pull requests, and commits from one or more GitHub repos.", + icon: GithubIcon, + dialogKey: "seed", + createConfig: (id, existing) => makeSeedConfig(id, existing, "github_repo"), + }, { kind: "sampler", type: "category", @@ -388,6 +403,7 @@ export function getBlockDefinitionForConfig( hf: "seed_hf", local: "seed_local", unstructured: "seed_unstructured", + github_repo: "seed_github", }; return getBlockDefinition("seed", seedType[config.seed_source_type ?? "hf"]); } diff --git a/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx b/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx index a8e12ac1ca..a3e5dbe962 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/execution-data-tab.tsx @@ -2,6 +2,8 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { ReactElement } from "react"; +import { GithubIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import type { ColumnDef } from "@tanstack/react-table"; import { Button } from "@/components/ui/button"; import { DataTable } from "@/components/ui/data-table"; @@ -12,11 +14,136 @@ import { DropdownMenuLabel, DropdownMenuTrigger, } from "@/components/ui/dropdown-menu"; +import { Progress } from "@/components/ui/progress"; import { Spinner } from "@/components/ui/spinner"; -import { cn } from "@/lib/utils"; import { isExecutionInProgress } from "../../executions/execution-helpers"; import type { RecipeExecutionRecord } from "../../execution-types"; -import { hasExpandableTextCell } from "./executions-view-helpers"; +import { formatMetricValue } from "./executions-view-helpers"; + +function formatSourceResource(value: string | null | undefined): string { + if (value === "pulls") { + return "PRs"; + } + return value ?? "--"; +} + +function formatFetchedValue(execution: RecipeExecutionRecord): string { + const source = execution.source_progress; + if (!source) { + return "--"; + } + const fetched = formatMetricValue(source.fetched_items); + if (typeof source.estimated_total !== "number" || source.estimated_total <= 0) { + return fetched; + } + return `${fetched} / ${formatMetricValue(source.estimated_total)}`; +} + +function formatGitHubSourceMessage(execution: RecipeExecutionRecord): string { + const source = execution.source_progress; + if (!source) { + return "Collecting repository threads before rows are available."; + } + if (source.status === "rate_limited") { + return source.message ?? "Waiting for GitHub rate limit. Studio will resume automatically."; + } + return source.message ?? "Collecting repository threads before rows are available."; +} + +function RunningDatasetEmptyState({ + execution, + onOpenOverview, +}: { + execution: RecipeExecutionRecord; + onOpenOverview: () => void; +}): ReactElement { + const source = execution.source_progress; + if (source?.source === "github") { + const title = + source.status === "rate_limited" + ? "Waiting for GitHub rate limit" + : "Crawling GitHub source"; + const showProgress = typeof source.percent === "number"; + + return ( +
+
+
+
+ {showProgress ? ( + + ) : ( + + )} +
+
+

{title}

+

+ {formatGitHubSourceMessage(execution)} +

+
+
+ +
+ {showProgress && } +
+

+ Repo + {source.repo ?? "--"} +

+

+ Resource + + {formatSourceResource(source.resource)} + {typeof source.page === "number" ? ` page ${source.page}` : ""} + +

+

+ Fetched + {formatFetchedValue(execution)} +

+

+ Rate remaining + {formatMetricValue(source.rate_remaining)} +

+

+ Retry wait + + {typeof source.retry_after_sec === "number" + ? `${formatMetricValue(source.retry_after_sec)}s` + : "--"} + +

+
+
+ ); + } + + return ( +
+
+ +
+

+ Generating data… +

+

+ {execution.current_column + ? `Current column: ${execution.current_column}` + : "Rows will appear here once the run produces a dataset sample."} +

+
+ +
+
+ ); +} type ExecutionDataTabProps = { execution: RecipeExecutionRecord; @@ -27,13 +154,10 @@ type ExecutionDataTabProps = { totalPages: number; tableColumns: ColumnDef>[]; datasetRowsForTable: Record[]; - visibleDatasetColumnNames: string[]; - expandedDatasetRows: Record; - selectedExecutionIdSafe: string | null; + onOpenOverview: () => void; onSetHiddenColumns: (updater: (current: string[]) => string[]) => void; onPrevPage: () => void; onNextPage: () => void; - onToggleRowExpanded: (rowId: string) => void; }; export function ExecutionDataTab({ @@ -45,13 +169,10 @@ export function ExecutionDataTab({ totalPages, tableColumns, datasetRowsForTable, - visibleDatasetColumnNames, - expandedDatasetRows, - selectedExecutionIdSafe, + onOpenOverview, onSetHiddenColumns, onPrevPage, onNextPage, - onToggleRowExpanded, }: ExecutionDataTabProps): ReactElement { return (
@@ -123,17 +244,10 @@ export function ExecutionDataTab({
{execution.dataset.length === 0 ? ( isExecutionInProgress(execution.status) ? ( -
- -
-

- Generating data… -

-

- Check the Overview tab for live terminal logs. -

-
-
+ ) : (

No rows returned.

) @@ -143,27 +257,7 @@ export function ExecutionDataTab({

) : (
- { - const canExpand = hasExpandableTextCell(row, visibleDatasetColumnNames); - if (!canExpand) { - return undefined; - } - return cn( - "cursor-pointer", - expandedDatasetRows[rowId] ? "bg-primary/[0.05]" : "hover:bg-primary/[0.06]", - ); - }} - onRowClick={(row, _rowIndex, rowId) => { - const canExpand = hasExpandableTextCell(row, visibleDatasetColumnNames); - if (!canExpand || !selectedExecutionIdSafe) { - return; - } - onToggleRowExpanded(rowId); - }} - /> +
)} diff --git a/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx b/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx index 8e2e710299..72624738fc 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/execution-overview-tab.tsx @@ -6,6 +6,7 @@ import { Database01Icon, Database02Icon, Flag02Icon, + GithubIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { Badge } from "@/components/ui/badge"; @@ -23,6 +24,28 @@ import type { RecipeExecutionRecord } from "../../execution-types"; import type { ModelUsageRow } from "./executions-view-helpers"; import { formatMetricValue } from "./executions-view-helpers"; +function formatSourceResource(value: string | null | undefined): string { + if (value === "pulls") { + return "PRs"; + } + return value ?? "--"; +} + +function formatSourceMessage(execution: RecipeExecutionRecord): string { + const source = execution.source_progress; + if (!source) { + return "No source progress captured."; + } + if (source.status === "rate_limited") { + const wait = + typeof source.retry_after_sec === "number" && source.retry_after_sec > 0 + ? ` Waiting ~${formatMetricValue(source.retry_after_sec)}s.` + : ""; + return `Waiting for GitHub rate limit. Studio will resume automatically.${wait}`; + } + return source.message ?? "Crawling GitHub source."; +} + type ExecutionOverviewTabProps = { execution: RecipeExecutionRecord; showSummaryCards: boolean; @@ -60,6 +83,8 @@ export function ExecutionOverviewTab({ canPublish, onOpenPublish, }: ExecutionOverviewTabProps): ReactElement { + const sourceProgress = execution.source_progress; + return (
{showSummaryCards && ( @@ -204,6 +229,47 @@ export function ExecutionOverviewTab({ )}
)} + {sourceProgress?.source === "github" && ( +
+
+

Source data

+ +
+

+ {sourceProgress.status === "completed" + ? "GitHub source complete" + : "Crawling GitHub source"} +

+

+ {formatSourceMessage(execution)} +

+
+

+ Repo + {sourceProgress.repo ?? "--"} +

+

+ Resource + + {formatSourceResource(sourceProgress.resource)} + {typeof sourceProgress.page === "number" ? ` page ${sourceProgress.page}` : ""} + +

+

+ Fetched + + {formatMetricValue(sourceProgress.fetched_items)} + +

+

+ Rate remaining + + {formatMetricValue(sourceProgress.rate_remaining)} + +

+
+
+ )}

Terminal output

diff --git a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx index f442d8f98e..cfc0e61338 100644 --- a/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx +++ b/studio/frontend/src/features/recipe-studio/components/executions/executions-view.tsx @@ -32,10 +32,9 @@ import { formatCellValue, formatDuration, formatPercent, - hasExpandableTextCell, + isExpandableCellValue, parseAnalysisColumns, parseModelUsageRows, - truncateCellValue, } from "./executions-view-helpers"; type ExecutionsViewProps = { @@ -63,9 +62,6 @@ export function ExecutionsView({ const [hiddenDatasetColumnsByExecution, setHiddenDatasetColumnsByExecution] = useState< Record >({}); - const [expandedDatasetRowsByExecution, setExpandedDatasetRowsByExecution] = useState< - Record> - >({}); const [previewDatasetPageByExecution, setPreviewDatasetPageByExecution] = useState< Record >({}); @@ -91,13 +87,6 @@ export function ExecutionsView({ } return hiddenDatasetColumnsByExecution[selectedExecutionIdSafe] ?? []; }, [hiddenDatasetColumnsByExecution, selectedExecutionIdSafe]); - const expandedDatasetRows = useMemo(() => { - if (!selectedExecutionIdSafe) { - return {}; - } - return expandedDatasetRowsByExecution[selectedExecutionIdSafe] ?? {}; - }, [expandedDatasetRowsByExecution, selectedExecutionIdSafe]); - const datasetColumnNames = useMemo(() => { if (!selectedExecution) { return []; @@ -119,6 +108,36 @@ export function ExecutionsView({ [datasetColumnNames, hiddenDatasetColumns], ); + // Columns where at least one row has text long enough that it would wrap at + // the default narrow width. We give those columns a wider min-width so the + // text is readable without clicking anything. The table's wrapper already + // scrolls horizontally, so a few wide columns just add a horizontal + // scrollbar instead of squeezing everything into the viewport. + const wideColumns = useMemo(() => { + const result = new Set(); + if (!selectedExecution) { + return result; + } + for (const row of selectedExecution.dataset) { + for (const name of visibleDatasetColumnNames) { + if (result.has(name)) { + continue; + } + const raw = row[name]; + if (resolveImagePreview(raw)) { + continue; + } + if (isExpandableCellValue(formatCellValue(raw))) { + result.add(name); + } + } + if (result.size === visibleDatasetColumnNames.length) { + break; + } + } + return result; + }, [selectedExecution, visibleDatasetColumnNames]); + const tableColumns = useMemo>[]>(() => { if (!selectedExecution) { return []; @@ -126,12 +145,12 @@ export function ExecutionsView({ return visibleDatasetColumnNames.map((name) => ({ accessorKey: name, header: name, - cell: ({ getValue, row }) => { + cell: ({ getValue }) => { const rawValue = getValue(); const imagePreview = resolveImagePreview(rawValue); if (imagePreview?.kind === "ready") { return ( -
+
{`${name} -

- Image too large to preview -

-
+

+ Image too large to preview +

); } const value = formatCellValue(rawValue); - const rowExpanded = Boolean(expandedDatasetRows[row.id]); - const rowHasExpandableCell = hasExpandableTextCell( - row.original, - visibleDatasetColumnNames, - ); - const showTruncated = rowHasExpandableCell && !rowExpanded; - + const isWide = wideColumns.has(name); return ( -
-

- {showTruncated ? truncateCellValue(value) : value} -

+
+

{value}

); }, })); - }, [expandedDatasetRows, selectedExecution, visibleDatasetColumnNames]); + }, [selectedExecution, visibleDatasetColumnNames, wideColumns]); const analysisColumns = useMemo( () => parseAnalysisColumns(selectedExecution?.analysis ?? null), @@ -512,9 +521,7 @@ export function ExecutionsView({ totalPages={totalPages} tableColumns={tableColumns} datasetRowsForTable={datasetRowsForTable} - visibleDatasetColumnNames={visibleDatasetColumnNames} - expandedDatasetRows={expandedDatasetRows} - selectedExecutionIdSafe={selectedExecutionIdSafe} + onOpenOverview={() => setDetailTab("overview")} onSetHiddenColumns={(updater) => { const selectedId = selectedExecution.id; setHiddenDatasetColumnsByExecution((current) => { @@ -547,18 +554,6 @@ export function ExecutionsView({ } onLoadDatasetPage(selectedExecution.id, currentDatasetPage + 1); }} - onToggleRowExpanded={(rowId) => { - setExpandedDatasetRowsByExecution((current) => { - const rows = current[selectedExecution.id] ?? {}; - return { - ...current, - [selectedExecution.id]: { - ...rows, - [rowId]: !rows[rowId], - }, - }; - }); - }} /> diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx index bf022fa028..82ec2f1de8 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-seed.tsx @@ -1,7 +1,12 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -import { DocumentAttachmentIcon, DocumentCodeIcon, Plant01Icon } from "@hugeicons/core-free-icons"; +import { + DocumentAttachmentIcon, + DocumentCodeIcon, + GithubIcon, + Plant01Icon, +} from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import type { ReactElement } from "react"; import type { SeedConfig } from "../../types"; @@ -13,9 +18,75 @@ type InlineSeedProps = { onUpdate: (patch: Partial) => void; }; -export function InlineSeed({ config, onUpdate }: InlineSeedProps): ReactElement { +const GITHUB_REPO_LIST_SPLIT_RE = /[\n,]+/; + +function parseGithubRepos(value: string | undefined): string[] { + return (value ?? "") + .split(GITHUB_REPO_LIST_SPLIT_RE) + .map((repo) => repo.trim()) + .filter(Boolean); +} + +function hasInvalidGithubRepo(repo: string): boolean { + const parts = repo.split("/"); + return parts.length !== 2 || parts.some((part) => !part.trim()); +} + +export function InlineSeed({ + config, + onUpdate, +}: InlineSeedProps): ReactElement { const mode = config.seed_source_type ?? "hf"; + if (mode === "github_repo") { + const repos = parseGithubRepos(config.github_repo_slug); + const invalidCount = repos.filter(hasInvalidGithubRepo).length; + const summary = + repos.length === 0 + ? "Add GitHub repos" + : repos.length === 1 + ? repos[0] + : `${repos.length} repositories`; + const configuredItems = config.github_item_types; + const items = + configuredItems && configuredItems.length > 0 + ? configuredItems + : ["issues", "pulls"]; + const itemsLabel = items.join(" · "); + const limit = (config.github_limit ?? "100").trim() || "100"; + const commentsLabel = + config.github_include_comments === false + ? "comments off" + : `comments ≤ ${config.github_max_comments_per_item ?? "30"}`; + const tokenLabel = config.github_token?.trim() ? "PAT set" : "server token"; + const warning = + repos.length === 0 + ? "No repos configured" + : invalidCount > 0 + ? `${invalidCount} invalid repo${invalidCount === 1 ? "" : "s"}` + : null; + return ( +
+
+ +
+
+

{summary}

+

+ {warning ?? + `${itemsLabel} · limit ${limit} · ${commentsLabel} · ${tokenLabel}`} +

+
+ +
+ ); + } + if (mode === "hf") { return (
@@ -62,10 +133,14 @@ export function InlineSeed({ config, onUpdate }: InlineSeedProps): ReactElement {fileName || "No file selected"}

- {isLocal ? "Structured file" : "Unstructured document"} · configure in dialog + {isLocal ? "Structured file" : "Unstructured document"} · configure in + dialog

- +
); } diff --git a/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx b/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx index acbda2d7d8..385bef2b77 100644 --- a/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx +++ b/studio/frontend/src/features/recipe-studio/components/recipe-studio-header.tsx @@ -34,6 +34,7 @@ type RecipeStudioHeaderProps = { savedAtLabel: string; workflowName: string; warnings?: GraphWarning[]; + supportsEasyMode?: boolean; onWorkflowNameChange: (value: string) => void; onViewChange: (view: RecipeStudioView) => void; onSaveRecipe: () => void; @@ -51,6 +52,7 @@ export function RecipeStudioHeader({ savedAtLabel, workflowName, warnings = [], + supportsEasyMode = false, onWorkflowNameChange, onViewChange, onSaveRecipe, @@ -58,7 +60,7 @@ export function RecipeStudioHeader({ const [editingWorkflowName, setEditingWorkflowName] = useState(false); function handleViewValueChange(value: string): void { - if (value === "editor" || value === "executions") { + if (value === "easy" || value === "editor" || value === "executions") { onViewChange(value); } } @@ -130,7 +132,12 @@ export function RecipeStudioHeader({
- Editor + {supportsEasyMode && ( + Easy + )} + + {supportsEasyMode ? "Advanced" : "Editor"} + Runs diff --git a/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx b/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx index 43f7437848..563280c55f 100644 --- a/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx +++ b/studio/frontend/src/features/recipe-studio/components/runtime/execution-progress-island.tsx @@ -49,6 +49,52 @@ function statusLabel(input: { return "Run status"; } +function formatSourceResource(value: string | null | undefined): string { + if (value === "pulls") { + return "PRs"; + } + return value ?? "source"; +} + +function formatGitHubSourceSummary( + execution: RecipeExecutionRecord, +): string | null { + const source = execution.source_progress; + if (!source || source.source !== "github") { + return null; + } + if (source.status === "rate_limited") { + const wait = + typeof source.retry_after_sec === "number" && source.retry_after_sec > 0 + ? ` ~${formatMetricValue(source.retry_after_sec)}s` + : ""; + return `Waiting for GitHub rate limit${wait}. Studio will resume automatically.`; + } + if (source.status === "retrying") { + return source.message ?? "GitHub request failed; retrying automatically."; + } + + const parts = [ + source.repo, + source.resource + ? `${formatSourceResource(source.resource)}${ + typeof source.page === "number" ? ` page ${source.page}` : "" + }` + : null, + typeof source.fetched_items === "number" + ? `${formatMetricValue(source.fetched_items)} fetched` + : null, + typeof source.rate_remaining === "number" + ? `remaining ${formatMetricValue(source.rate_remaining)}` + : null, + ].filter(Boolean); + + if (parts.length === 0) { + return source.message ?? "Crawling GitHub source"; + } + return parts.join(" · "); +} + export function ExecutionProgressIsland({ execution, currentColumnIcon, @@ -58,14 +104,24 @@ export function ExecutionProgressIsland({ }: ExecutionProgressIslandProps): ReactElement { const complete = execution.status === "completed"; const inProgress = isExecutionInProgress(execution.status); - const progressPercent = execution.progress?.percent ?? (complete ? 100 : 0); + const sourceSummary = formatGitHubSourceSummary(execution); + const showSourceProgress = Boolean( + sourceSummary && + inProgress && + (execution.stage === "source" || + execution.source_progress?.status === "rate_limited" || + execution.source_progress?.status === "retrying"), + ); + const sourcePercent = showSourceProgress ? execution.source_progress?.percent : null; + const progressPercent = sourcePercent ?? execution.progress?.percent ?? (complete ? 100 : 0); const hasProgressSignal = Boolean( - execution.progress && + (execution.progress && (typeof execution.progress.done === "number" || typeof execution.progress.total === "number" || typeof execution.progress.percent === "number" || typeof execution.progress.rate === "number" || - typeof execution.progress.eta_sec === "number"), + typeof execution.progress.eta_sec === "number")) || + typeof sourcePercent === "number", ); const showLoadingSpinner = inProgress && !hasProgressSignal; const batchTotal = execution.batch?.total ?? null; @@ -136,18 +192,30 @@ export function ExecutionProgressIsland({ ETA: {formatEta(execution.progress?.eta_sec)}

-
- -

- Column: {execution.current_column ?? "--"} -

-
+ {showSourceProgress ? ( +
+ +

+ GitHub source: {sourceSummary} +

+
+ ) : ( +
+ +

+ Column: {execution.current_column ?? "--"} +

+
+ )} {showBatch && (
)} - {!validateResult.valid && validateResult.rawDetail && ( -

+ {validateResult.rawDetail && ( +

{validateResult.rawDetail}

)} 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 868974acb5..54eae08f7c 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 @@ -31,18 +31,27 @@ import { TableHeader, TableRow, } from "@/components/ui/table"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { - Tabs, - TabsContent, - TabsList, - TabsTrigger, -} from "@/components/ui/tabs"; -import { type ReactElement, useCallback, useEffect, useMemo, useRef, useState } from "react"; + type KeyboardEvent, + type ReactElement, + useCallback, + useEffect, + useId, + useMemo, + useRef, + useState, +} from "react"; import { cn } from "@/lib/utils"; import { UnstructuredDropZone, type FileEntry } from "./unstructured-drop-zone"; -import { inspectSeedDataset, inspectSeedUpload } from "../../api"; +import { + getGithubEnvTokenStatus, + inspectSeedDataset, + inspectSeedUpload, +} from "../../api"; import { resolveImagePreview } from "../../utils/image-preview"; import type { + GithubItemType, SeedConfig, SeedSamplingStrategy, SeedSelectionType, @@ -51,10 +60,11 @@ import { CollapsibleSectionTriggerButton } from "../shared/collapsible-section-t import { HfDatasetCombobox } from "../../components/shared/hf-dataset-combobox"; import { FieldLabel } from "../shared/field-label"; -const SAMPLING_OPTIONS: Array<{ value: SeedSamplingStrategy; label: string }> = [ - { value: "ordered", label: "Ordered" }, - { value: "shuffle", label: "Shuffle" }, -]; +const SAMPLING_OPTIONS: Array<{ value: SeedSamplingStrategy; label: string }> = + [ + { value: "ordered", label: "Ordered" }, + { value: "shuffle", label: "Shuffle" }, + ]; const SELECTION_OPTIONS: Array<{ value: SeedSelectionType; label: string }> = [ { value: "none", label: "None" }, @@ -68,6 +78,14 @@ const DEFAULT_CHUNK_SIZE = 1200; const DEFAULT_CHUNK_OVERLAP = 200; const MAX_CHUNK_SIZE = 20000; const PREVIEW_TRUNCATE_AT = 320; +const GITHUB_TRAILING_PUNCTUATION_RE = /[),.;]+$/g; +const GITHUB_SSH_PREFIX_RE = /^git@github\.com:/i; +const URL_PROTOCOL_RE = /^https?:\/\//i; +const WWW_PREFIX_RE = /^www\./i; +const URL_QUERY_OR_HASH_RE = /[?#]/; +const GIT_SUFFIX_RE = /\.git$/i; +const GITHUB_REPO_INPUT_SPLIT_RE = /[\s,]+/; +const GITHUB_REPO_PART_RE = /^[A-Za-z0-9_.-]+$/; type SeedDialogProps = { config: SeedConfig; @@ -75,6 +93,367 @@ type SeedDialogProps = { open: boolean; }; +function normalizeGithubRepoInput(value: string): string { + let raw = value.trim().replace(GITHUB_TRAILING_PUNCTUATION_RE, ""); + if (!raw) { + return ""; + } + raw = raw.replace(GITHUB_SSH_PREFIX_RE, "https://github.com/"); + raw = raw.replace(URL_PROTOCOL_RE, ""); + raw = raw.replace(WWW_PREFIX_RE, ""); + if (raw.toLowerCase().startsWith("github.com/")) { + raw = raw.slice("github.com/".length); + } + raw = raw.split(URL_QUERY_OR_HASH_RE)[0] ?? ""; + raw = raw.replace(GIT_SUFFIX_RE, ""); + const parts = raw.split("/").filter(Boolean); + if (parts.length >= 2) { + return `${parts[0]}/${parts[1]}`; + } + return raw; +} + +function splitGithubRepoInput(value: string): string[] { + return value + .split(GITHUB_REPO_INPUT_SPLIT_RE) + .map(normalizeGithubRepoInput) + .filter(Boolean); +} + +function getGithubRepoError(repo: string): string | null { + const parts = repo.split("/"); + if (parts.length !== 2 || parts.some((part) => !part.trim())) { + return "Use owner/name."; + } + if (parts.some((part) => !GITHUB_REPO_PART_RE.test(part))) { + return "Use only GitHub repo characters."; + } + return null; +} + +function githubLimitString(value: string | undefined): string { + return (value ?? "100").trim(); +} + +export function GithubRepoSeedForm({ + config, + onUpdate, +}: { + config: SeedConfig; + onUpdate: (patch: Partial) => void; +}): ReactElement { + const [repoDraft, setRepoDraft] = useState(""); + const [serverHasEnvToken, setServerHasEnvToken] = useState( + null, + ); + const repoInputId = useId(); + const repoHelpId = useId(); + const repoErrorId = useId(); + const tokenId = useId(); + const tokenHelpId = useId(); + const limitId = useId(); + const limitHelpId = useId(); + const commentsId = useId(); + const includeCommentsId = useId(); + const commentsHelpId = useId(); + const repos = useMemo( + () => splitGithubRepoInput(config.github_repo_slug ?? ""), + [config.github_repo_slug], + ); + const repoErrors = repos + .map((repo, index) => ({ repo, index, error: getGithubRepoError(repo) })) + .filter((item) => item.error); + const hasRepoErrors = repoErrors.length > 0; + const configuredItemTypes = config.github_item_types; + const itemTypes: GithubItemType[] = + configuredItemTypes && configuredItemTypes.length > 0 + ? configuredItemTypes + : ["issues", "pulls"]; + const limitNum = Number.parseInt(githubLimitString(config.github_limit), 10); + const boundedLimit = Number.isFinite(limitNum) + ? Math.min(5000, Math.max(1, limitNum)) + : 100; + const estimatedItems = repos.length * itemTypes.length * boundedLimit; + const includeComments = config.github_include_comments ?? true; + const hasToken = Boolean(config.github_token?.trim()); + const usingEnvToken = !hasToken && serverHasEnvToken === true; + + useEffect(() => { + let cancelled = false; + void getGithubEnvTokenStatus() + .then((status) => { + if (!cancelled) setServerHasEnvToken(status.has_token); + }) + .catch(() => { + if (!cancelled) setServerHasEnvToken(false); + }); + return () => { + cancelled = true; + }; + }, []); + + function updateRepos(nextRepos: string[]): void { + const seen = new Set(); + const deduped = nextRepos.filter((repo) => { + const key = repo.toLowerCase(); + if (seen.has(key)) { + return false; + } + seen.add(key); + return true; + }); + onUpdate({ github_repo_slug: deduped.join("\n") }); + } + + function addRepos(raw: string): void { + const next = splitGithubRepoInput(raw); + if (next.length === 0) { + return; + } + updateRepos([...repos, ...next]); + setRepoDraft(""); + } + + function removeRepo(index: number): void { + updateRepos(repos.filter((_, i) => i !== index)); + } + + function handleRepoKeyDown(event: KeyboardEvent): void { + if (["Enter", ",", " "].includes(event.key)) { + event.preventDefault(); + addRepos(repoDraft); + } else if (event.key === "Backspace" && !repoDraft && repos.length > 0) { + event.preventDefault(); + removeRepo(repos.length - 1); + } + } + + return ( +
+
+ +
+ {repos.map((repo, index) => { + const error = getGithubRepoError(repo); + return ( + + {repo} + + + ); + })} + setRepoDraft(event.target.value)} + onKeyDown={handleRepoKeyDown} + onBlur={() => addRepos(repoDraft)} + onPaste={(event) => { + const pasted = event.clipboardData.getData("text"); + if (pasted) { + event.preventDefault(); + addRepos(pasted); + } + }} + placeholder={ + repos.length === 0 + ? "unslothai/unsloth or GitHub URL" + : "Add repo…" + } + aria-invalid={hasRepoErrors} + aria-describedby={`${repoHelpId}${hasRepoErrors ? ` ${repoErrorId}` : ""}`} + /> +
+

+ Stored as one owner/name repo per line in the recipe. +

+ {hasRepoErrors && ( +
    + {repoErrors.map((item) => ( +
  • + Row {item.index + 1}: {item.error} +
    • + ))} +
+ )} +
+ +
+
+ + {usingEnvToken && ( + + Using server env var + + )} +
+ onUpdate({ github_token: e.target.value })} + placeholder={ + usingEnvToken + ? "Using server GH_TOKEN / GITHUB_TOKEN" + : "Leave blank to use server GH_TOKEN" + } + aria-describedby={tokenHelpId} + /> +

+ {usingEnvToken + ? "Studio detected a server env token, so saved/shared recipes can leave this blank." + : "Blank is safest for saved/shared recipes because Studio will read the server environment at run time."} +

+ {hasToken && ( +

+ Personal access tokens are sensitive. Prefer server env vars when + possible, and avoid sharing recipes that contain a PAT. +

+ )} +
+ +
+ + Fetch scope + +
+ + onUpdate({ github_limit: e.target.value })} + placeholder="100" + aria-describedby={limitHelpId} + /> +

+ Estimate before comments: up to {estimatedItems.toLocaleString()}{" "} + rows ({repos.length || 0} repos × {itemTypes.length} item types ×{" "} + {boundedLimit} limit). Commit crawling uses each repo's default + branch. +

+
+ +
+ + Item types + +
+ {(["issues", "pulls", "commits"] as const).map((kind) => { + const checked = itemTypes.includes(kind); + const itemTypeId = `${repoInputId}-${kind}`; + return ( + + ); + })} +
+
+ +
+ +
+ + + onUpdate({ github_max_comments_per_item: e.target.value }) + } + aria-describedby={commentsHelpId} + /> +

+ Comments increase GraphQL cost and can make Check/Run look quiet + while GitHub pages and rate-limit waits stream in logs. +

+
+
+
+ +

+ Backed by Studio's built-in github_repo seed reader. Large + repos can take minutes, so start with small limits for previews. +

+
+ ); +} + function getErrorMessage(error: unknown, fallback: string): string { if (error instanceof Error && error.message) { return error.message; @@ -85,7 +464,8 @@ function getErrorMessage(error: unknown, fallback: string): string { function stringifyCell(value: unknown): string { if (value === null || value === undefined) return ""; if (typeof value === "string") return value; - if (typeof value === "number" || typeof value === "boolean") return String(value); + if (typeof value === "number" || typeof value === "boolean") + return String(value); try { return JSON.stringify(value); } catch { @@ -111,7 +491,8 @@ function getPreviewEmptyStateCopy(mode: SeedConfig["seed_source_type"]): { if (mode === "local") { return { title: "No preview yet", - description: "Upload a CSV, JSON, or JSONL file and click Load to see a sample.", + description: + "Upload a CSV, JSON, or JSONL file and click Load to see a sample.", }; } if (mode === "unstructured") { @@ -121,9 +502,17 @@ function getPreviewEmptyStateCopy(mode: SeedConfig["seed_source_type"]): { "Upload your documents and the preview will appear once processing is done.", }; } + if (mode === "github_repo") { + return { + title: "GitHub data loads during Check or Run", + description: + "Configure repos, item types, and limits above. GitHub crawling can take minutes on large repos; watch logs for page and rate-limit updates.", + }; + } return { title: "No preview yet", - description: "Select a Hugging Face dataset and click Load to see a sample.", + description: + "Select a Hugging Face dataset and click Load to see a sample.", }; } @@ -175,24 +564,32 @@ async function fileToBase64Payload(file: File): Promise { }); } -export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactElement { +export function SeedDialog({ + config, + onUpdate, + open, +}: SeedDialogProps): ReactElement { const [inspectError, setInspectError] = useState(null); const [isInspecting, setIsInspecting] = useState(false); const advancedOpen = config.advancedOpen === true; const [previewRows, setPreviewRows] = useState[]>([]); - const [expandedPreviewRows, setExpandedPreviewRows] = useState>({}); + const [expandedPreviewRows, setExpandedPreviewRows] = useState< + Record + >({}); const [localFile, setLocalFile] = useState(null); - const [unstructuredFiles, setUnstructuredFiles] = useState(() => { - if (config.unstructured_file_ids?.length) { - return config.unstructured_file_ids.map((id, i) => ({ - id, - name: config.unstructured_file_names?.[i] ?? "Unknown", - size: config.unstructured_file_sizes?.[i] ?? 0, - status: "ok" as const, - })); - } - return []; - }); + const [unstructuredFiles, setUnstructuredFiles] = useState( + () => { + if (config.unstructured_file_ids?.length) { + return config.unstructured_file_ids.map((id, i) => ({ + id, + name: config.unstructured_file_names?.[i] ?? "Unknown", + size: config.unstructured_file_sizes?.[i] ?? 0, + status: "ok" as const, + })); + } + return []; + }, + ); const mode = config.seed_source_type ?? "hf"; const previewEmpty = getPreviewEmptyStateCopy(mode); @@ -241,7 +638,14 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl status: "ok" as const, })), ); - }, [open, mode, unstructuredFiles.length, config.unstructured_file_ids, config.unstructured_file_names, config.unstructured_file_sizes]); + }, [ + open, + mode, + unstructuredFiles.length, + config.unstructured_file_ids, + config.unstructured_file_names, + config.unstructured_file_sizes, + ]); const handleUnstructuredFilesChange = useCallback( (updater: FileEntry[] | ((prev: FileEntry[]) => FileEntry[])) => { @@ -291,142 +695,135 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl const { chunkSize, chunkOverlap } = resolveChunking(config); const fileKey = okFiles.map((f) => `${f.id}|${f.name}`).join(","); return `unstructured:${fileKey}|${chunkSize}|${chunkOverlap}`; - }, [ - config, - localFile, - mode, - unstructuredFiles, - ]); + }, [config, localFile, mode, unstructuredFiles]); - const loadSeedMetadata = useCallback(async (opts?: { silent?: boolean }): Promise => { - const loadKey = getCurrentLoadKey(); - if (!opts?.silent) { - setInspectError(null); - } - setIsInspecting(true); - try { - if (mode === "hf") { - const datasetName = config.hf_repo_id.trim(); - if (!datasetName) { - throw new Error("Dataset repo is required."); - } - const response = await inspectSeedDataset({ - dataset_name: datasetName, - hf_token: config.hf_token?.trim() || undefined, - split: config.hf_split?.trim() || undefined, - subset: config.hf_subset?.trim() || undefined, - preview_size: 10, - }); - onUpdate({ - hf_path: response.resolved_path, - seed_columns: response.columns, - seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => - response.columns.includes(name), - ), - seed_preview_rows: response.preview_rows ?? [], - hf_split: response.split ?? "", - hf_subset: response.subset ?? "", - local_file_name: "", - unstructured_file_ids: [], - unstructured_file_names: [], - unstructured_file_sizes: [], - }); - setPreviewRows(response.preview_rows ?? []); - setLastLoadedKey(loadKey); - return true; - } - - if (mode === "local") { - if (!localFile) { - throw new Error("Select a local CSV/JSON/JSONL file first."); - } - if (localFile.size > MAX_UPLOAD_BYTES) { - throw new Error("File too large (max 50MB)."); - } - const payload = await fileToBase64Payload(localFile); - const response = await inspectSeedUpload({ - filename: localFile.name, - content_base64: payload, - preview_size: 10, - }); - onUpdate({ - hf_path: response.resolved_path, - seed_columns: response.columns, - seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => - response.columns.includes(name), - ), - seed_preview_rows: response.preview_rows ?? [], - hf_repo_id: "", - hf_subset: "", - hf_split: "", - local_file_name: localFile.name, - unstructured_file_ids: [], - unstructured_file_names: [], - unstructured_file_sizes: [], - }); - setPreviewRows(response.preview_rows ?? []); - setLastLoadedKey(loadKey); - return true; - } - - if (mode === "unstructured") { - const fileIds = unstructuredFiles - .filter((f) => f.status === "ok") - .map((f) => f.id); - const fileNames = unstructuredFiles - .filter((f) => f.status === "ok") - .map((f) => f.name); - - if (fileIds.length === 0) { - setInspectError("No files uploaded"); - return false; - } - - const { chunkSize, chunkOverlap } = resolveChunking(config); - const response = await inspectSeedUpload({ - block_id: config.id, - file_ids: fileIds, - file_names: fileNames, - preview_size: 10, - seed_source_type: "unstructured", - unstructured_chunk_size: chunkSize, - unstructured_chunk_overlap: chunkOverlap, - }); - - onUpdate({ - hf_path: response.resolved_path, - resolved_paths: response.resolved_paths ?? [], - seed_columns: response.columns, - seed_preview_rows: response.preview_rows ?? [], - unstructured_file_ids: fileIds, - unstructured_file_names: fileNames, - unstructured_file_sizes: unstructuredFiles - .filter((f) => f.status === "ok") - .map((f) => f.size), - }); - setPreviewRows(response.preview_rows ?? []); - setLastLoadedKey(loadKey); - return true; - } - - return false; - } catch (error) { + const loadSeedMetadata = useCallback( + async (opts?: { silent?: boolean }): Promise => { + const loadKey = getCurrentLoadKey(); if (!opts?.silent) { - setInspectError(getErrorMessage(error, "Failed to load seed metadata.")); + setInspectError(null); } - setPreviewRows([]); - return false; - } finally { - setIsInspecting(false); - } - }, [ - config, - getCurrentLoadKey, - localFile, - mode, - onUpdate, - unstructuredFiles, - ]); + setIsInspecting(true); + try { + if (mode === "hf") { + const datasetName = config.hf_repo_id.trim(); + if (!datasetName) { + throw new Error("Dataset repo is required."); + } + const response = await inspectSeedDataset({ + dataset_name: datasetName, + hf_token: config.hf_token?.trim() || undefined, + split: config.hf_split?.trim() || undefined, + subset: config.hf_subset?.trim() || undefined, + preview_size: 10, + }); + onUpdate({ + hf_path: response.resolved_path, + seed_columns: response.columns, + seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => + response.columns.includes(name), + ), + seed_preview_rows: response.preview_rows ?? [], + hf_split: response.split ?? "", + hf_subset: response.subset ?? "", + local_file_name: "", + unstructured_file_ids: [], + unstructured_file_names: [], + unstructured_file_sizes: [], + }); + setPreviewRows(response.preview_rows ?? []); + setLastLoadedKey(loadKey); + return true; + } + + if (mode === "local") { + if (!localFile) { + throw new Error("Select a local CSV/JSON/JSONL file first."); + } + if (localFile.size > MAX_UPLOAD_BYTES) { + throw new Error("File too large (max 50MB)."); + } + const payload = await fileToBase64Payload(localFile); + const response = await inspectSeedUpload({ + filename: localFile.name, + content_base64: payload, + preview_size: 10, + }); + onUpdate({ + hf_path: response.resolved_path, + seed_columns: response.columns, + seed_drop_columns: (config.seed_drop_columns ?? []).filter((name) => + response.columns.includes(name), + ), + seed_preview_rows: response.preview_rows ?? [], + hf_repo_id: "", + hf_subset: "", + hf_split: "", + local_file_name: localFile.name, + unstructured_file_ids: [], + unstructured_file_names: [], + unstructured_file_sizes: [], + }); + setPreviewRows(response.preview_rows ?? []); + setLastLoadedKey(loadKey); + return true; + } + + if (mode === "unstructured") { + const fileIds = unstructuredFiles + .filter((f) => f.status === "ok") + .map((f) => f.id); + const fileNames = unstructuredFiles + .filter((f) => f.status === "ok") + .map((f) => f.name); + + if (fileIds.length === 0) { + setInspectError("No files uploaded"); + return false; + } + + const { chunkSize, chunkOverlap } = resolveChunking(config); + const response = await inspectSeedUpload({ + block_id: config.id, + file_ids: fileIds, + file_names: fileNames, + preview_size: 10, + seed_source_type: "unstructured", + unstructured_chunk_size: chunkSize, + unstructured_chunk_overlap: chunkOverlap, + }); + + onUpdate({ + hf_path: response.resolved_path, + resolved_paths: response.resolved_paths ?? [], + seed_columns: response.columns, + seed_preview_rows: response.preview_rows ?? [], + unstructured_file_ids: fileIds, + unstructured_file_names: fileNames, + unstructured_file_sizes: unstructuredFiles + .filter((f) => f.status === "ok") + .map((f) => f.size), + }); + setPreviewRows(response.preview_rows ?? []); + setLastLoadedKey(loadKey); + return true; + } + + return false; + } catch (error) { + if (!opts?.silent) { + setInspectError( + getErrorMessage(error, "Failed to load seed metadata."), + ); + } + setPreviewRows([]); + return false; + } finally { + setIsInspecting(false); + } + }, + [config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles], + ); useEffect(() => { const wasOpen = wasOpenRef.current; @@ -463,7 +860,8 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl return []; }, [config.seed_columns, previewRows]); const selectedSeedDropColumns = useMemo( - () => (config.seed_drop_columns ?? []).filter((name) => name.trim().length > 0), + () => + (config.seed_drop_columns ?? []).filter((name) => name.trim().length > 0), [config.seed_drop_columns], ); const selectedSeedDropSet = useMemo( @@ -540,10 +938,11 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl className="nodrag" placeholder="hf_..." value={config.hf_token ?? ""} - onChange={(event) => onUpdate({ hf_token: event.target.value })} + onChange={(event) => + onUpdate({ hf_token: event.target.value }) + } />
- )} @@ -600,9 +999,15 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl /> )} - {inspectError &&

{inspectError}

} + {mode === "github_repo" && ( + + )} - {mode !== "unstructured" && ( + {inspectError && ( +

{inspectError}

+ )} + + {mode !== "unstructured" && mode !== "github_repo" && (
{ const isChecked = value === true; const next = isChecked - ? Array.from(new Set([...selectedSeedDropColumns, columnName])) - : selectedSeedDropColumns.filter((name) => name !== columnName); + ? Array.from( + new Set([ + ...selectedSeedDropColumns, + columnName, + ]), + ) + : selectedSeedDropColumns.filter( + (name) => name !== columnName, + ); onUpdate({ seed_drop_columns: next }); }} /> @@ -660,7 +1072,9 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl onUpdate({ selection_start: event.target.value })} + onChange={(event) => + onUpdate({ selection_start: event.target.value }) + } />
- + onUpdate({ selection_end: event.target.value })} + onChange={(event) => + onUpdate({ selection_end: event.target.value }) + } />
@@ -772,17 +1203,24 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl className="nodrag" inputMode="numeric" value={config.selection_index ?? ""} - onChange={(event) => onUpdate({ selection_index: event.target.value })} + onChange={(event) => + onUpdate({ selection_index: event.target.value }) + } />
- + - onUpdate({ selection_num_partitions: event.target.value }) + onUpdate({ + selection_num_partitions: event.target.value, + }) } />
@@ -830,7 +1268,8 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl { @@ -850,7 +1289,9 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl className="max-w-[260px] whitespace-pre-wrap break-words text-xs" > {(() => { - const imagePreview = resolveImagePreview(row[col]); + const imagePreview = resolveImagePreview( + row[col], + ); if (imagePreview?.kind === "ready") { return ( ; + rows: number; + setRows: (rows: number) => void; + updateConfig: (id: string, patch: Partial) => void; + onRun: () => void; + runLoading: boolean; + runErrors: string[]; + onSwitchToAdvanced: () => void; +}; + +export function GithubCrawlerEasyView({ + configs, + rows, + setRows, + updateConfig, + onRun, + runLoading, + runErrors, + onSwitchToAdvanced, +}: GithubCrawlerEasyViewProps): ReactElement { + const seedConfig = useMemo( + () => + Object.values(configs).find((c): c is SeedConfig => c.kind === "seed") ?? + null, + [configs], + ); + const modelConfig = useMemo( + () => + Object.values(configs).find( + (c): c is ModelConfig => c.kind === "model_config", + ) ?? null, + [configs], + ); + + // Local buffer for the Rows input so the user can hold transient invalid + // state (empty while backspacing, partial digits, etc.) without the parent + // snapping them back to 1 on every keystroke. The canonical ``rows`` value + // only advances when the buffer parses to a valid positive integer; on + // blur we clamp back to a sane default if the user left it empty. + const [rowsText, setRowsText] = useState(String(rows)); + useEffect(() => { + setRowsText(String(rows)); + }, [rows]); + + const handleSeedUpdate = (patch: Partial): void => { + if (!seedConfig) return; + updateConfig(seedConfig.id, patch); + }; + + const handleModelChange = (value: string): void => { + if (!modelConfig) return; + updateConfig(modelConfig.id, { model: value }); + }; + + if (!seedConfig) { + return ( +
+

+ This recipe has no seed node. Switch to{" "} + {" "} + to configure it. +

+
+ ); + } + + return ( +
+
+
+ +
+
+

GitHub Crawler

+

+ Crawl real GitHub issues and PRs and turn each thread into a{" "} + {"{User, Assistant}"} training pair. + Defaults use the server's GH_TOKEN env var and the + bundled local model. +

+
+
+ + + +
+

+ Run settings +

+
+
+ + { + // Allow empty / partial strings while the user is editing. + // type="text" avoids the browser's number spinner and the + // related backspace quirks; we still parse + clamp below. + const raw = event.target.value.replace(/[^0-9]/g, ""); + setRowsText(raw); + const next = Number.parseInt(raw, 10); + if (Number.isFinite(next) && next > 0 && next <= 10000) { + setRows(next); + } + }} + onBlur={() => { + const next = Number.parseInt(rowsText, 10); + if (!Number.isFinite(next) || next < 1) { + setRows(1); + setRowsText("1"); + } else if (next > 10000) { + setRows(10000); + setRowsText("10000"); + } else { + setRowsText(String(next)); + } + }} + /> +
+
+ + handleModelChange(event.target.value)} + placeholder="unsloth/gemma-4-E2B-it-GGUF" + disabled={!modelConfig} + /> +
+
+
+ + {runErrors.length > 0 && ( +
+

Cannot run:

+
    + {runErrors.slice(0, 4).map((err) => ( +
  • {err}
    • + ))} +
+
+ )} + +
+ + +
+
+ ); +} diff --git a/studio/frontend/src/features/recipe-studio/execution-types.ts b/studio/frontend/src/features/recipe-studio/execution-types.ts index c1976ab3f9..92dd88c4f5 100644 --- a/studio/frontend/src/features/recipe-studio/execution-types.ts +++ b/studio/frontend/src/features/recipe-studio/execution-types.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -export type RecipeStudioView = "editor" | "executions"; +export type RecipeStudioView = "easy" | "editor" | "executions"; export type RecipeExecutionKind = "preview" | "full"; @@ -30,6 +30,28 @@ export type RecipeExecutionBatch = { total?: number | null; }; +export type RecipeSourceProgress = { + source?: string | null; + status?: string | null; + repo?: string | null; + resource?: string | null; + page?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + page_items?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + fetched_items?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + estimated_total?: number | null; + percent?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + rate_remaining?: number | null; + // biome-ignore lint/style/useNamingConvention: backend schema + retry_after_sec?: number | null; + message?: string | null; + // biome-ignore lint/style/useNamingConvention: backend schema + updated_at?: number | null; +}; + export type RecipeExecutionAnalysis = { num_records?: number; target_num_records?: number; @@ -65,6 +87,8 @@ export type RecipeExecutionRecord = { column_progress: RecipeExecutionProgress | null; batch: RecipeExecutionBatch | null; // biome-ignore lint/style/useNamingConvention: backend schema + source_progress: RecipeSourceProgress | null; + // biome-ignore lint/style/useNamingConvention: backend schema model_usage: Record | null; // biome-ignore lint/style/useNamingConvention: backend schema lastEventId: number | null; diff --git a/studio/frontend/src/features/recipe-studio/executions/runtime.ts b/studio/frontend/src/features/recipe-studio/executions/runtime.ts index bdc3e653fc..e37a1c9929 100644 --- a/studio/frontend/src/features/recipe-studio/executions/runtime.ts +++ b/studio/frontend/src/features/recipe-studio/executions/runtime.ts @@ -6,6 +6,7 @@ import type { RecipeExecutionBatch, RecipeExecutionKind, RecipeExecutionRecord, + RecipeSourceProgress, } from "../execution-types"; import { DATASET_PAGE_SIZE, @@ -72,6 +73,14 @@ export function toExecutionLogLine(event: JobEvent): string | null { return null; } +function normalizeSourceProgress(input: unknown): RecipeSourceProgress | null { + const raw = normalizeObject(input); + if (!raw) { + return null; + } + return raw as RecipeSourceProgress; +} + export function applyExecutionStatusSnapshot( execution: RecipeExecutionRecord, status: JobStatusResponse, @@ -100,6 +109,7 @@ export function applyExecutionStatusSnapshot( (normalizeObject(status.column_progress) as RecipeExecutionRecord["column_progress"]) ?? null, batch, + source_progress: normalizeSourceProgress(status.source_progress), model_usage: normalizeObject(status.model_usage), artifact_path: status.artifact_path ?? execution.artifact_path, error: status.error ?? null, @@ -137,6 +147,7 @@ export function createBaseExecutionRecord(input: { progress: null, column_progress: null, batch: null, + source_progress: null, model_usage: null, lastEventId: null, artifact_path: null, diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts index 24d02e0f6c..1040025a5d 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts @@ -3,7 +3,12 @@ import { useCallback, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; +import { toast } from "sonner"; import { toastError } from "@/shared/toast"; +import { + getInferenceStatus, + loadModel, +} from "@/features/chat/api/chat-api"; import { cancelRecipeJob, createRecipeJob, @@ -39,10 +44,84 @@ import { } from "../stores/recipe-executions"; import type { RecipePayload, RecipePayloadResult } from "../utils/payload/types"; +/** + * Auto-load the local model before running a recipe that uses it. + * + * Looks at payload.recipe.model_providers for any provider with is_local=true, + * finds the bound model_configs and asks the backend to load whichever model + * the first local-bound model_config points at. Skips when the inference + * server already has that exact model active. This removes the "open /chat + * first" prerequisite that users kept tripping on. + */ +async function ensureLocalModelLoaded( + payload: RecipePayload, +): Promise { + const providers = Array.isArray(payload.recipe.model_providers) + ? (payload.recipe.model_providers as Array>) + : []; + const localProviderNames = new Set(); + for (const p of providers) { + if (p.is_local === true && typeof p.name === "string") { + localProviderNames.add(p.name); + } + } + if (localProviderNames.size === 0) { + return null; + } + const modelConfigs = Array.isArray(payload.recipe.model_configs) + ? (payload.recipe.model_configs as Array>) + : []; + const boundConfig = modelConfigs.find( + (c) => typeof c.provider === "string" && localProviderNames.has(c.provider), + ); + const target = + typeof boundConfig?.model === "string" ? boundConfig.model.trim() : ""; + if (!target) { + return null; + } + + try { + const status = await getInferenceStatus(); + if ( + status.active_model && + status.active_model.toLowerCase() === target.toLowerCase() + ) { + return null; + } + } catch { + // Fall through to load attempt; the backend will re-error if needed. + } + + const toastId = toast.loading(`Loading ${target}…`, { + description: "Starting the local inference server for this recipe.", + }); + try { + const isGguf = /gguf/i.test(target); + await loadModel({ + model_path: target, + hf_token: null, + max_seq_length: isGguf ? 0 : 4096, + load_in_4bit: true, + is_lora: false, + gguf_variant: null, + trust_remote_code: false, + chat_template_override: null, + cache_type_kv: null, + speculative_type: null, + }); + toast.success(`Loaded ${target}`, { id: toastId, duration: 2000 }); + return null; + } catch (error) { + toast.dismiss(toastId); + return error instanceof Error ? error.message : String(error); + } +} + type UseRecipeExecutionsParams = { recipeId: string; currentSignature: string; payloadResult: RecipePayloadResult; + initialRunRows?: number | null; onExecutionStart?: () => void; onPreviewSuccess?: () => void; }; @@ -101,6 +180,7 @@ export function useRecipeExecutions({ recipeId, currentSignature, payloadResult, + initialRunRows, onExecutionStart, onPreviewSuccess, }: UseRecipeExecutionsParams): UseRecipeExecutionsResult { @@ -181,6 +261,19 @@ export function useRecipeExecutions({ resetForRecipe(); + // Seed previewRows from the recipe's original run.rows (read from the + // loaded JSON, not the rebuilt payload (the builder hardcodes 5). + // Templates ship their own suggested preview size (e.g. GitHub Support + // Bot: 10); we honor it so users don't see a surprise 5. + if ( + typeof initialRunRows === "number" && + Number.isFinite(initialRunRows) && + initialRunRows > 0 && + initialRunRows !== 5 + ) { + setPreviewRows(Math.floor(initialRunRows)); + } + async function hydrate(): Promise { try { const records = await loadSortedRecipeExecutions(recipeId); @@ -216,10 +309,12 @@ export function useRecipeExecutions({ cancelled = true; }; }, [ + initialRunRows, onPreviewSuccess, recipeId, resetForRecipe, setExecutions, + setPreviewRows, setRunErrors, upsertAndPersist, ]); @@ -340,6 +435,20 @@ export function useRecipeExecutions({ return false; } + // Flip to the Runs pane BEFORE we run ensureLocalModelLoaded + validate. + // Validation re-crawls the seed (multiple seconds for the github_repo + // reader) and the user otherwise stares at a "Running..." button with + // nothing else changing. runExecution() later no-ops this callback if + // the view has already been flipped, so we fire it once here. + onExecutionStart?.(); + + const localLoadError = await ensureLocalModelLoaded(payload); + if (localLoadError) { + setRunErrors([localLoadError]); + toastError("Local model failed to load", localLoadError); + return false; + } + const normalizedRows = sanitizeExecutionRows(rows, kind); const executionPayload = buildExecutionPayload({ payload, @@ -374,7 +483,13 @@ export function useRecipeExecutions({ runName, }); }, - [readExecutablePayload, runExecution, runSettings, setRunErrors], + [ + onExecutionStart, + readExecutablePayload, + runExecution, + runSettings, + setRunErrors, + ], ); const runPreview = useCallback(async (): Promise => { diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts index f6f7390cf7..79c75c1885 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-studio-actions.ts @@ -112,6 +112,10 @@ export function useRecipeStudioActions({ recipeId, currentSignature: persistence.currentSignature, payloadResult, + initialRunRows: + typeof initialPayload?.run?.rows === "number" + ? initialPayload.run.rows + : null, onExecutionStart, onPreviewSuccess, }); diff --git a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx index 586d5d8ad8..5dbc004ffd 100644 --- a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx +++ b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx @@ -47,6 +47,7 @@ import { ConfigDialog } from "./dialogs/config-dialog"; import { ImportDialog } from "./dialogs/import-dialog"; import { RunDialog } from "./dialogs/preview-dialog"; import { ProcessorsDialog } from "./dialogs/processors-dialog"; +import { GithubCrawlerEasyView } from "./easy/github-crawler-easy-view"; import type { RecipeExecutionRecord, RecipeStudioView, @@ -201,7 +202,40 @@ export function RecipeStudioPage({ null, ); const flowContainerRef = useRef(null); - const [activeView, setActiveView] = useState("editor"); + const supportsEasyMode = + initialPayload?.ui?.seed_source_type === "github_repo" || + (initialPayload?.recipe?.seed_config as { source?: { seed_type?: string } } | undefined) + ?.source?.seed_type === "github_repo"; + const viewModeStorageKey = `recipe-studio:view-mode:${recipeId}`; + const [activeView, setActiveViewState] = useState(() => { + if (typeof window !== "undefined") { + const stored = window.localStorage.getItem(viewModeStorageKey); + if (stored === "easy" && supportsEasyMode) return "easy"; + if (stored === "editor" || stored === "executions") return stored; + } + return supportsEasyMode ? "easy" : "editor"; + }); + const setActiveView = useCallback( + (next: RecipeStudioView | ((prev: RecipeStudioView) => RecipeStudioView)) => { + setActiveViewState((prev) => { + const resolved = typeof next === "function" ? next(prev) : next; + if (typeof window !== "undefined") { + window.localStorage.setItem(viewModeStorageKey, resolved); + } + return resolved; + }); + }, + [viewModeStorageKey], + ); + // Easy mode has no canvas overlay/progress island, so once a run starts the + // user sees the Run button stuck on "Running..." with nothing else changing. + // Flip to the Runs pane so they land where progress is actually rendered. + // Advanced (editor) keeps its island and stays put. + const handleExecutionStart = useCallback(() => { + setActiveView((currentView) => + currentView === "easy" ? "executions" : currentView, + ); + }, [setActiveView]); const [processorsOpen, setProcessorsOpen] = useState(false); const [interactive, setInteractive] = useState(true); const [runtimeIslandMinimized, setRuntimeIslandMinimized] = useState(false); @@ -326,6 +360,8 @@ export function RecipeStudioPage({ validateResult, cancelExecution, loadExecutionDatasetPage, + runPreview, + runFull, copyRecipe, importRecipe, } = useRecipeStudioActions({ @@ -338,6 +374,7 @@ export function RecipeStudioPage({ resetRecipe, loadRecipe, getCurrentPayloadFromStore, + onExecutionStart: handleExecutionStart, }); const { activeExecution, @@ -359,6 +396,19 @@ export function RecipeStudioPage({ const runBusy = previewLoading || fullLoading || executionLocked; const islandExecution = activeExecution ?? recentCompletedExecution; + // Easy mode runs a full run (artifact persisted, tracked in Runs pane) + // using runFull. runFull requires a non-empty fullRunName but the Easy form + // has no run-name input, so seed a default here as soon as Easy is active. + // User can still rename it from Advanced/Runs dialogs before clicking Run. + useEffect(() => { + if (!supportsEasyMode) return; + if (activeView !== "easy") return; + if (fullRunName.trim()) return; + const stamp = new Date().toISOString().replace(/[:.]/g, "-").slice(0, 16); + const base = workflowName.trim() || "Easy run"; + setFullRunName(`${base} ${stamp}`); + }, [supportsEasyMode, activeView, fullRunName, workflowName, setFullRunName]); + const toggleInteractive = useCallback(() => { if (executionLocked) { return; @@ -739,6 +789,7 @@ export function RecipeStudioPage({ savedAtLabel={savedAtLabel} workflowName={workflowName} warnings={getGraphWarnings(configs, edges)} + supportsEasyMode={supportsEasyMode} onWorkflowNameChange={setWorkflowName} onViewChange={setActiveView} onSaveRecipe={() => { @@ -749,7 +800,25 @@ export function RecipeStudioPage({ className="h-[75vh] w-full rounded-t-none" ref={flowContainerRef} > - {activeView === "editor" ? ( + {activeView === "easy" ? ( + { + // Easy mode is a full run (artifact persisted, tracked in + // the Runs pane) capped at the user's row count. runFull + // requires a non-empty fullRunName; the effect below + // populates one on mount so the closure in runFull is + // already up to date by the time the user clicks Run. + void runFull(); + }} + runLoading={fullLoading || executionLocked} + runErrors={runErrors} + onSwitchToAdvanced={() => setActiveView("editor")} + /> + ) : activeView === "editor" ? ( editorContent ) : ( ((set, get) => ({ nextSourceType = "local"; } else if (type === "seed_unstructured") { nextSourceType = "unstructured"; + } else if (type === "seed_github") { + nextSourceType = "github_repo"; } - const nextConfig = { + const nextConfig: typeof existing = { ...existing, seed_source_type: nextSourceType, hf_repo_id: "", @@ -415,6 +417,12 @@ export const useRecipeStudioStore = create((set, get) => ({ seed_preview_rows: [], unstructured_chunk_size: "1200", unstructured_chunk_overlap: "200", + github_repo_slug: "", + github_token: "", + github_limit: "100", + github_item_types: ["issues", "pulls"], + github_include_comments: true, + github_max_comments_per_item: "30", }; return { configs: { diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index cd1127a4e2..9c720a06d3 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -38,7 +38,10 @@ export type LayoutDirection = "LR" | "TB"; export type SeedSamplingStrategy = "ordered" | "shuffle"; export type SeedSelectionType = "none" | "index_range" | "partition_block"; -export type SeedSourceType = "hf" | "local" | "unstructured"; +export type SeedSourceType = "hf" | "local" | "unstructured" | "github_repo"; + +export type GithubItemType = "issues" | "pulls" | "commits"; +export type GithubStateFilter = "all" | "open" | "closed"; export const INFRA_NODE_KINDS = new Set([ "model_provider", "model_config", @@ -338,6 +341,20 @@ export type SeedConfig = { unstructured_file_ids?: string[]; unstructured_file_names?: string[]; unstructured_file_sizes?: number[]; + // biome-ignore lint/style/useNamingConvention: api schema + github_repo_slug?: string; + // biome-ignore lint/style/useNamingConvention: api schema + github_token?: string; + // biome-ignore lint/style/useNamingConvention: api schema + github_limit?: string; + // biome-ignore lint/style/useNamingConvention: api schema + github_item_types?: GithubItemType[]; + // biome-ignore lint/style/useNamingConvention: api schema + github_state?: GithubStateFilter; + // biome-ignore lint/style/useNamingConvention: api schema + github_include_comments?: boolean; + // biome-ignore lint/style/useNamingConvention: api schema + github_max_comments_per_item?: string; resolved_paths?: string[]; // ui-only seed_preview_rows?: Record[]; diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts index 0d92129232..21eadb3195 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts @@ -80,6 +80,12 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { let resolved_paths: string[] = []; let unstructured_chunk_size = "1200"; let unstructured_chunk_overlap = "200"; + let github_repo_slug = ""; + let github_token = ""; + let github_limit = "100"; + let github_item_types: ("issues" | "pulls" | "commits")[] = ["issues", "pulls"]; + let github_include_comments = true; + let github_max_comments_per_item = "30"; const sourceRaw = seedConfigRaw.source; if (isRecord(sourceRaw)) { const seedType = readString(sourceRaw.seed_type); @@ -107,6 +113,26 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { unstructuredFileNames = []; unstructured_chunk_size = readNumberString(sourceRaw.chunk_size) || "1200"; unstructured_chunk_overlap = readNumberString(sourceRaw.chunk_overlap) || "200"; + } else if (seedType === "github_repo") { + seed_source_type = "github_repo"; + const rawRepos = Array.isArray(sourceRaw.repos) ? sourceRaw.repos : []; + const repos = rawRepos.filter((r): r is string => typeof r === "string"); + github_repo_slug = repos.join("\n"); + github_token = readString(sourceRaw.token) ?? ""; + github_limit = readNumberString(sourceRaw.limit) || "100"; + const rawItems = Array.isArray(sourceRaw.item_types) ? sourceRaw.item_types : []; + const validItems = rawItems.filter( + (t): t is "issues" | "pulls" | "commits" => + t === "issues" || t === "pulls" || t === "commits", + ); + if (validItems.length > 0) { + github_item_types = validItems; + } + if (typeof sourceRaw.include_comments === "boolean") { + github_include_comments = sourceRaw.include_comments; + } + github_max_comments_per_item = + readNumberString(sourceRaw.max_comments_per_item) || "30"; } } @@ -147,6 +173,12 @@ function parseSeedSettings(seedConfigRaw: unknown): Partial { resolved_paths, unstructured_chunk_size, unstructured_chunk_overlap, + github_repo_slug, + github_token, + github_limit, + github_item_types, + github_include_comments, + github_max_comments_per_item, sampling_strategy, selection_type, selection_start, diff --git a/studio/frontend/src/features/recipe-studio/utils/node-data.ts b/studio/frontend/src/features/recipe-studio/utils/node-data.ts index 2fc5205db8..96042b5919 100644 --- a/studio/frontend/src/features/recipe-studio/utils/node-data.ts +++ b/studio/frontend/src/features/recipe-studio/utils/node-data.ts @@ -70,7 +70,9 @@ export function nodeDataFromConfig( ? "Hugging Face dataset" : seedSourceType === "local" ? "CSV or JSON file" - : "Document file"; + : seedSourceType === "github_repo" + ? "GitHub repositories" + : "Document file"; return { title: "Source data", kind: "seed", diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts index f73a2b46da..bb48b43857 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-seed.ts @@ -6,6 +6,7 @@ import type { NodeConfig, SeedConfig } from "../../types"; const DEFAULT_CHUNK_SIZE = 1200; const DEFAULT_CHUNK_OVERLAP = 200; const MAX_CHUNK_SIZE = 20000; +const GITHUB_ITEM_TYPES = new Set(["issues", "pulls", "commits"]); function parseIntStrict(value: string | undefined): number | null { const trimmed = value?.trim(); @@ -56,33 +57,86 @@ export function buildSeedConfig( selectionStrategy = { index, num_partitions: numPartitions }; } - const source = - seedSourceType === "hf" - ? { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "hf", - path, - token, - endpoint, - } - : seedSourceType === "unstructured" - ? (() => { - const { chunkSize, chunkOverlap } = resolveChunking(config); - return { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "unstructured", - paths: config.resolved_paths?.length ? config.resolved_paths : [config.hf_path], - // biome-ignore lint/style/useNamingConvention: api schema - chunk_size: chunkSize, - // biome-ignore lint/style/useNamingConvention: api schema - chunk_overlap: chunkOverlap, - }; - })() - : { - // biome-ignore lint/style/useNamingConvention: api schema - seed_type: "local", - path, - }; + let source: Record; + if (seedSourceType === "hf") { + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "hf", + path, + token, + endpoint, + }; + } else if (seedSourceType === "unstructured") { + const { chunkSize, chunkOverlap } = resolveChunking(config); + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "unstructured", + paths: config.resolved_paths?.length ? config.resolved_paths : [config.hf_path], + // biome-ignore lint/style/useNamingConvention: api schema + chunk_size: chunkSize, + // biome-ignore lint/style/useNamingConvention: api schema + chunk_overlap: chunkOverlap, + }; + } else if (seedSourceType === "github_repo") { + const repos = (config.github_repo_slug ?? "") + .split(/[\n,]/) + .map((r) => r.trim()) + .filter(Boolean); + if (repos.length === 0) { + errors.push(`Seed ${config.name}: at least one repo is required.`); + return undefined; + } + const invalidRepo = repos.find((repo) => { + const parts = repo.split("/"); + return parts.length !== 2 || parts.some((part) => !part); + }); + if (invalidRepo) { + errors.push( + `Seed ${config.name}: GitHub repositories must use owner/name format.`, + ); + return undefined; + } + const itemTypes = config.github_item_types?.length + ? config.github_item_types + : ["issues", "pulls"]; + if (itemTypes.some((itemType) => !GITHUB_ITEM_TYPES.has(itemType))) { + errors.push(`Seed ${config.name}: GitHub item types invalid.`); + return undefined; + } + const limitNum = parseIntStrict(config.github_limit ?? "100"); + if (limitNum === null || limitNum < 1 || limitNum > 5000) { + errors.push( + `Seed ${config.name}: GitHub items per repo must be an integer from 1 to 5000.`, + ); + return undefined; + } + const maxCommentsNum = parseIntStrict(config.github_max_comments_per_item ?? "30"); + if (maxCommentsNum === null || maxCommentsNum < 0 || maxCommentsNum > 200) { + errors.push( + `Seed ${config.name}: GitHub max comments per item must be an integer from 0 to 200.`, + ); + return undefined; + } + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "github_repo", + repos, + token: (config.github_token ?? "").trim(), + // biome-ignore lint/style/useNamingConvention: api schema + item_types: itemTypes, + limit: limitNum, + // biome-ignore lint/style/useNamingConvention: api schema + include_comments: config.github_include_comments ?? true, + // biome-ignore lint/style/useNamingConvention: api schema + max_comments_per_item: maxCommentsNum, + }; + } else { + source = { + // biome-ignore lint/style/useNamingConvention: api schema + seed_type: "local", + path, + }; + } return { source, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts index b8b1aba93b..902ea796b4 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts @@ -57,7 +57,7 @@ export type RecipePayload = { // ui-only: graph orientation layout_direction?: "LR" | "TB"; // ui-only, used to preserve seed block mode across imports/refresh - seed_source_type?: "hf" | "local" | "unstructured"; + seed_source_type?: "hf" | "local" | "unstructured" | "github_repo"; // ui-only, persisted aux node positions by llm name + aux key aux_nodes?: Array<{ llm: string; diff --git a/studio/frontend/src/features/recipe-studio/utils/validation.ts b/studio/frontend/src/features/recipe-studio/utils/validation.ts index 5f48fd44a5..d5200eeed0 100644 --- a/studio/frontend/src/features/recipe-studio/utils/validation.ts +++ b/studio/frontend/src/features/recipe-studio/utils/validation.ts @@ -8,6 +8,7 @@ import { isOxcCodeShape } from "./validators/oxc-code-shape"; import { isOxcValidationMode } from "./validators/oxc-mode"; const TRACE_MODES = new Set(["none", "last_message", "all_messages"]); +const GITHUB_ITEM_TYPES = new Set(["issues", "pulls", "commits"]); // biome-ignore lint/complexity/noExcessiveCognitiveComplexity: validation rules export function getConfigErrors(config: NodeConfig | null): string[] { @@ -268,15 +269,49 @@ export function getConfigErrors(config: NodeConfig | null): string[] { } if (config.kind === "seed") { const seedSourceType = config.seed_source_type ?? "hf"; - if (seedSourceType === "hf" && !config.hf_repo_id.trim()) { + if (seedSourceType === "github_repo") { + const repos = (config.github_repo_slug ?? "") + .split(/[\n,]/) + .map((repo) => repo.trim()) + .filter(Boolean); + if (repos.length === 0) { + errors.push("Add at least one GitHub repository."); + } + if ( + repos.some((repo) => { + const parts = repo.split("/"); + return parts.length !== 2 || parts.some((part) => !part); + }) + ) { + errors.push("GitHub repositories must use owner/name format."); + } + const itemTypes = config.github_item_types?.length + ? config.github_item_types + : ["issues", "pulls"]; + if (itemTypes.length === 0) { + errors.push("Choose at least one GitHub item type."); + } else if (itemTypes.some((itemType) => !GITHUB_ITEM_TYPES.has(itemType))) { + errors.push("GitHub item types must be issues, pulls, or commits."); + } + const limit = parseIntNumber(config.github_limit ?? "100"); + if (limit === null || limit < 1 || limit > 5000) { + errors.push("Items per repo must be an integer from 1 to 5000."); + } + const maxComments = parseIntNumber(config.github_max_comments_per_item ?? "30"); + if (maxComments === null || maxComments < 0 || maxComments > 200) { + errors.push("Max comments per item must be an integer from 0 to 200."); + } + } else if (seedSourceType === "hf" && !config.hf_repo_id.trim()) { errors.push("Choose a Hugging Face dataset."); } - const hasPath = - seedSourceType === "unstructured" - ? (config.resolved_paths?.length ?? 0) > 0 - : Boolean(config.hf_path.trim()); - if (!hasPath) { - errors.push("Load the source-data preview first."); + if (seedSourceType !== "github_repo") { + const hasPath = + seedSourceType === "unstructured" + ? (config.resolved_paths?.length ?? 0) > 0 + : Boolean(config.hf_path.trim()); + if (!hasPath) { + errors.push("Load the source-data preview first."); + } } if ( seedSourceType === "hf" && @@ -304,7 +339,7 @@ export function getConfigErrors(config: NodeConfig | null): string[] { ) { errors.push("Chunk overlap must be less than chunk size."); } - } else { + } else if (seedSourceType !== "github_repo") { const selectedDropColumns = (config.seed_drop_columns ?? []) .map((value) => value.trim()) .filter(Boolean); diff --git a/studio/frontend/src/lib/api-base.ts b/studio/frontend/src/lib/api-base.ts index 664b9e58bd..c8be442653 100644 --- a/studio/frontend/src/lib/api-base.ts +++ b/studio/frontend/src/lib/api-base.ts @@ -8,6 +8,12 @@ if (isTauri && !isViteDev) { apiBase = 'http://127.0.0.1:8888' } +const initialApiBase = apiBase + +export function resetApiBase() { + apiBase = initialApiBase +} + export function setApiBase(port: number) { apiBase = `http://127.0.0.1:${port}` } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index b56b737cfa..5fe4415970 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -418,6 +418,9 @@ CONSTRAINTS = SINGLE_ENV / "constraints.txt" LOCAL_DD_UNSTRUCTURED_PLUGIN = ( SCRIPT_DIR / "backend" / "plugins" / "data-designer-unstructured-seed" ) +LOCAL_DD_GITHUB_PLUGIN = ( + SCRIPT_DIR / "backend" / "plugins" / "data-designer-github-repo-seed" +) # -- Unicode-safe printing --------------------------------------------- # On Windows the default console encoding can be a legacy code page @@ -1135,22 +1138,28 @@ def install_python_stack() -> int: req = SINGLE_ENV / "data-designer.txt", ) - # 11. Local Data Designer seed plugin - if not LOCAL_DD_UNSTRUCTURED_PLUGIN.is_dir(): - _safe_print( - _red( - f"❌ Missing local plugin directory: {LOCAL_DD_UNSTRUCTURED_PLUGIN}", - ), - ) - return 1 + # 11. Local Data Designer seed plugins + local_dd_plugins = [ + ("unstructured", LOCAL_DD_UNSTRUCTURED_PLUGIN), + ("github", LOCAL_DD_GITHUB_PLUGIN), + ] + for _plugin_name, plugin_dir in local_dd_plugins: + if not plugin_dir.is_dir(): + _safe_print( + _red( + f"❌ Missing local plugin directory: {plugin_dir}", + ), + ) + return 1 _progress("local plugin") - pip_install( - "Installing local data-designer unstructured plugin", - "--no-cache-dir", - "--no-deps", - str(LOCAL_DD_UNSTRUCTURED_PLUGIN), - constrain = False, - ) + for plugin_name, plugin_dir in local_dd_plugins: + pip_install( + f"Installing local data-designer {plugin_name} plugin", + "--no-cache-dir", + "--no-deps", + str(plugin_dir), + constrain = False, + ) # 12. Patch metadata for single-env compatibility _progress("finalizing")