diff --git a/CHANGELOG.md b/CHANGELOG.md index 3526159..bad0cff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,55 @@ Format loosely follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Unreleased] +## [2.6.0] - 2026-09-24 + +A simpler interactive menu with a fixed screen layout, help at every stage and +export at any stage, plus a fix for runs that looked like an endless error +loop when the API key was missing or invalid. One behaviour change: Ctrl+C at +a stage checkpoint now stops and saves the run. + +### Added +- **Export at any stage** (`exporters/stage_data.py`). Discovered URLs (txt, + csv, json), scraped pages (one Markdown file per page with front matter, + plain text, json, jsonl), chunks (jsonl, json, txt) and all samples (jsonl, + json). "Export what I have so far" at a checkpoint offers whatever exists, + plus the training set when there are samples. `dataforge export ` uses + the same flow and no longer refuses a session without samples; Browse my + data has an export entry too. +- **Help**: "What happens next?" at each stage checkpoint explains every + stage and the next one. A Help page in the main menu; `help` and `h` work in + the URL review. +- "Scrape pages (no AI)" in the main menu runs `dataforge scrape`. + +### Changed +- **Main menu worded as tasks**: Scrape pages (no AI), Build an AI training + dataset, Continue a paused project, Browse my data, Export data, Settings + and tools, Help, Clear the screen, Exit. Rarely used items (provider and + key, system info, update, uninstall, project list) moved to a submenu. + Every choice shows a one-line description. +- **Screen layout**: wizard steps, pipeline stages and the URL review start on + a fresh screen, in one order (title, content, messages, then key hints + directly above the prompt). The URL review redraws in place and the URL + list is no longer printed twice. `DATAFORGE_NO_CLEAR=1` keeps all output. +- The interactive menu shows WARNING and above on the terminal unless + `DATAFORGE_LOG_LEVEL` is set; the log file still records INFO. +- **Input**: paths and URLs accept surrounding quotes (Windows "Copy as + path"), PowerShell drag-and-drop (`& 'C:\...'`), curly quotes, ``, + and URL lists separated by commas, spaces or lines. +- **Behaviour change**: Ctrl+C at a stage checkpoint stops and saves the run; + it previously continued to the next stage. + +### Fixed +- **A missing or invalid API key no longer loops.** Fatal credential and + connection errors now reach the generator and streaming agent instead of + being swallowed per chunk; auth failures are not retried; LiteLLM debug + banners are suppressed; `generate_batch` stops dispatching and cancels + pending work on the first fatal error. With 40 chunks and a bad key, the run + stops in about 3 seconds with one message, where it previously made 3 calls + per chunk. +- "Export what I have so far" after collection or processing wrote nothing, + because it exported only samples. + ## [2.5.0] - 2026-09-23 Quick scrape (#69, first part): `dataforge scrape` and the MCP tool diff --git a/CITATION.cff b/CITATION.cff index 82fab26..3f67f97 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -12,8 +12,8 @@ authors: orcid: 'https://orcid.org/0009-0000-4888-1941' affiliation: William Jessup University repository-code: 'https://github.com/ianktoo/data-forge' -version: 2.5.0 -date-released: '2026-09-23' +version: 2.6.0 +date-released: '2026-09-24' abstract: >- DataForge is an open-source pipeline for building synthetic instruction-tuning datasets from websites. It honors HTTP diff --git a/README.md b/README.md index 4d9d40b..1c8e500 100644 --- a/README.md +++ b/README.md @@ -6,13 +6,14 @@ [![Python](https://img.shields.io/badge/python-3.11%20%7C%203.12%20%7C%203.13%20%7C%203.14-blue)](https://www.python.org/) [![License](https://img.shields.io/badge/license-MIT-green)](LICENSE) -**Current release: 2.5.0.** New: `dataforge scrape` saves pages and their -tables (CSV, JSON) with no AI, API key or cost, also available to AI agents as -the MCP tool `scrape_page`. Discovery is faster and politer: a crawl without a -sitemap follows site navigation, visits the most useful pages first and -downloads each page once. Tested on Linux, Windows and macOS with Python 3.11 -to 3.14, with standalone binaries for Windows, Linux and Apple Silicon Macs. -See the [changelog](CHANGELOG.md#250---2026-09-23). +**Current release: 2.6.0.** An easier interactive menu: tasks in plain +words, one screen per step, help at every stage, and "export what I have so +far" at any point (discovered URLs, pages, chunks or samples). A missing or +invalid API key now stops the run with one clear message instead of retrying +every chunk. `dataforge scrape` saves pages and their tables with no AI, API +key or cost. Tested on Linux, Windows and macOS with Python 3.11 to 3.14, +with standalone binaries for Windows, Linux and Apple Silicon Macs. See the +[changelog](CHANGELOG.md#260---2026-09-24). Fine-tuning needs data, and good domain data is trapped in documentation sites, knowledge bases and public archives. Getting it out usually means diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md index 379341e..9720201 100644 --- a/docs/CONFIGURATION.md +++ b/docs/CONFIGURATION.md @@ -23,7 +23,8 @@ Run `dataforge config` to set your provider and API key interactively. | `DATAFORGE_MAX_CRAWL_DEPTH` | `3` | Max link depth for the fallback crawl (the seed is depth 0) | | `DATAFORGE_CHUNK_SIZE` | `512` | Tokens per chunk | | `DATAFORGE_CHUNK_OVERLAP` | `64` | Token overlap between chunks | -| `DATAFORGE_LOG_LEVEL` | `INFO` | `DEBUG` / `INFO` / `WARNING` / `ERROR` | +| `DATAFORGE_LOG_LEVEL` | `INFO` | `DEBUG` / `INFO` / `WARNING` / `ERROR`. Terminal level; the interactive menu shows `WARNING` and above unless this is set. The log file always records `INFO`. | +| `DATAFORGE_NO_CLEAR` | unset | Set to `1` to keep all output instead of clearing the screen between steps | | `DATAFORGE_OUTPUT_DIR` | `./output` | Session output directory (logs also stored here in `logs/`) | | `DATAFORGE_DB_PATH` | `./dataforge.db` | SQLite database path | | `DATAFORGE_AUTOSAVE` | `true` | Checkpoint progress to the session DB after every stage | diff --git a/docs/TECHNICAL.tex b/docs/TECHNICAL.tex index 84dad35..e8ecf3b 100644 --- a/docs/TECHNICAL.tex +++ b/docs/TECHNICAL.tex @@ -357,6 +357,11 @@ \subsection{Export} version~2.4.2 the Unsloth file had no lineage at all (project issue \#19), so a consumer training from it lost the ability to cite a sample's source document. +Earlier stages are exportable too, without an LLM: the discovered URL +list (text, CSV, JSON), the scraped pages (one Markdown or plain-text file +per page, or JSON), and the chunks with their source URL (JSON, JSONL, +text), so a session that stops before generation still yields a usable +corpus. \subsubsection{Dataset introspection} \label{sec:introspection} diff --git a/pyproject.toml b/pyproject.toml index d29513f..bb6e25b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "llm-web-crawler" -version = "2.5.0" +version = "2.6.0" description = "LLM data collection and synthetic fine-tuning dataset pipeline" readme = "README.md" requires-python = ">=3.11" diff --git a/src/dataforge/cli/app.py b/src/dataforge/cli/app.py index 36ef75d..eaed43c 100644 --- a/src/dataforge/cli/app.py +++ b/src/dataforge/cli/app.py @@ -115,12 +115,18 @@ def _apply_project_file(s, cwd: Path) -> None: pass # Malformed file: ignore and fall back to defaults -def _bootstrap() -> None: +def _bootstrap(interactive: bool = False) -> None: from dataforge.cli.preflight import check_env_file check_env_file() s = get_settings() _apply_project_file(s, Path.cwd()) - setup_logging(s.logs_dir(), s.log_level) + level = s.log_level + # The interactive screens already show progress; INFO log lines on top of + # them are what made the terminal scroll away. They still go to the log + # file. An explicit DATAFORGE_LOG_LEVEL wins. + if interactive and not os.getenv("DATAFORGE_LOG_LEVEL"): + level = "WARNING" + setup_logging(s.logs_dir(), level) init_db(s.db_path) @@ -159,7 +165,7 @@ def main( # Rich respects the NO_COLOR env var; set it so all consoles pick it up os.environ["NO_COLOR"] = "1" if ctx.invoked_subcommand is None: - _bootstrap() + _bootstrap(interactive=True) if not quiet: ui.banner() asyncio.run(_interactive_pipeline()) @@ -170,7 +176,7 @@ def main( @app.command() def pipeline() -> None: """Start a new interactive pipeline.""" - _bootstrap() + _bootstrap(interactive=True) if not _QUIET: ui.banner() asyncio.run(_interactive_pipeline()) @@ -224,20 +230,7 @@ def scrape( "files": written, }, indent=2, ensure_ascii=False)) else: - for p in pages: - if p.ok: - ui.success(f"{p.url} [dim]{p.title[:60]}[/] {p.word_count} words, " - f"{len(p.tables)} table{'s' if len(p.tables) != 1 else ''}") - for w in p.warnings: - ui.warn(f" {w}") - else: - ui.error(f"{p.url} {p.status}") - for p in ok: - for k, t in enumerate(p.tables[:3], 1): - _print_table_preview(t, f"{p.title or p.url}: table {k}") - n_tables = sum(len(p.tables) for p in ok) - ui.info(f"{len(ok)}/{len(pages)} page(s), {n_tables} table(s) saved to " - f"[bold]{out_dir.resolve()}[/]") + _print_scrape_results(pages, out_dir) if not ok: raise typer.Exit(code=1) @@ -554,47 +547,8 @@ async def _export_session(session_id: str, approved_only: bool) -> None: ui.error(f"Session '{session_id}' not found") raise typer.Exit(1) - # Check what's available - with open_session(s.db_path) as db: - sample_count = len(db.exec( - select(SyntheticSample).where(SyntheticSample.session_id == session_id) - ).all()) - - if sample_count == 0: - ui.warn("No synthetic samples yet. Run at least through the generation stage.") - raise typer.Exit(1) - - ui.info(f"{sample_count} samples available") - targets = await prompts.ask_export_targets( - hf_configured=bool(s.huggingface_token), - kg_configured=bool(s.kaggle_username and s.kaggle_key), - ) - - export_kw: dict = {"targets": targets, "approved_only": approved_only, - "stage_snapshot": session.stage} - - if "huggingface" in targets: - export_kw["hf_repo_id"] = await prompts.ask_hf_repo() - export_kw["hf_private"] = await prompts.ask_hf_private() - - if "kaggle" in targets: - export_kw["kaggle_slug"] = await prompts.ask_kaggle_slug(s.kaggle_username) - export_kw["kaggle_title"] = session.name - - from dataforge.agents import PipelineContext - from dataforge.agents.exporter import ExporterAgent - ctx = PipelineContext( - session_id=session.id, - session_name=session.name, - goal=session.goal, - format=DataFormat(session.format), - seed_urls=session.seed_url_list(), - settings=s, - ) - agent = ExporterAgent(ctx, **export_kw) - ctx = await agent.run() - ui.export_summary(ctx.export_records) - ui.success("Export complete") + # Any stage has something worth exporting (URLs, pages, chunks, samples). + await _export_flow(_context_for(session), session.stage, approved_only=approved_only) # ── view command ────────────────────────────────────────────────────────────── @@ -1464,7 +1418,7 @@ def _show_pipeline_plan() -> None: PipelineStage.generation, PipelineStage.quality, PipelineStage.export, ] - ui.section("Pipeline Plan") + ui.screen("Pipeline steps and project status") # Find current project sessions for context pf = find_project_file(Path.cwd()) @@ -1564,6 +1518,7 @@ async def _explore_menu(limit: int = 5) -> None: f" stage=[cyan]{session.stage}[/] limit=[bold]{limit}[/]" ) stage_choices = [questionary.Choice(label, value=key) for key, label in _STAGE_OPTIONS] + stage_choices.append(questionary.Choice("Export this project's data", value="__export__")) stage_choices.append(questionary.Choice(f"Change sample limit (current: {limit})", value="__limit__")) stage_choices.append(questionary.Choice("← Back to main menu", value="__back__")) @@ -1574,6 +1529,10 @@ async def _explore_menu(limit: int = 5) -> None: if not picked or picked == "__back__": break + if picked == "__export__": + await _export_flow(_context_for(session), session.stage) + continue + if picked == "__limit__": raw = await questionary.text( f"Enter number of records to show per page (current: {limit}):", @@ -1742,6 +1701,20 @@ async def _interactive_pipeline() -> None: action = await _main_menu() if action == "new": break # fall through to wizard + if action == "back": + continue + if action == "clear": + ui.screen("Main menu") + continue + if action == "help": + _show_help() + continue + if action == "scrape": + await _quick_scrape_flow() + continue + if action == "export": + await _export_menu() + continue if action == "resume": await _pick_and_resume() continue @@ -1752,6 +1725,7 @@ async def _interactive_pipeline() -> None: _show_pipeline_plan() continue if action == "sessions": + ui.screen("All projects") with open_session(s.db_path) as db: all_s = db.exec(select(PipelineSession)).all() rows = [{"id": x.id, "name": x.name, "stage": x.stage, @@ -1763,6 +1737,7 @@ async def _interactive_pipeline() -> None: await _configure() continue if action == "info": + ui.screen("System info") _show_info() continue if action == "update": @@ -1845,15 +1820,16 @@ async def _run_wizard(state: dict) -> str: step = STEPS[step_idx] if step == "urls": - ui.section("Input") + ui.screen("New dataset", "step 1 of 4: which site?") result = await _step_urls(state) elif step == "output": - ui.section("Output Location") + ui.screen("New dataset", "step 2 of 4: where to save") result = await _step_output(state) elif step == "config": - ui.section("Configuration") + ui.screen("New dataset", "step 3 of 4: what kind of data") result = await _step_config(state) elif step == "review": + ui.screen("New dataset", "step 4 of 4: check and start") result = await _step_review(state) else: result = "next" @@ -1903,6 +1879,14 @@ async def _run_wizard(state: dict) -> str: _STAGE_TOTAL = 6 +_NEXT_STAGE = { + PipelineStage.discovery: PipelineStage.collection, + PipelineStage.collection: PipelineStage.processing, + PipelineStage.processing: PipelineStage.generation, + PipelineStage.generation: PipelineStage.quality, + PipelineStage.quality: PipelineStage.export, +} + _LANG_RE = re.compile( r"/([a-z]{2}(?:-[a-z]{2})?)/|[?&]lang(?:uage)?=([a-z]{2})|[?&]locale=([a-z]{2})", re.IGNORECASE, @@ -1979,8 +1963,11 @@ async def _run_orchestrator(ctx: PipelineContext, start_from: str | None = None) async def stage_hook(stage: str, context: PipelineContext) -> bool: name, step = _stage_map.get(stage, (stage, 0)) + nxt = _NEXT_STAGE.get(stage) + next_name = _stage_map[nxt][0] if nxt in _stage_map else "" - # Show summary + contextual tip + # Results first, then the tip, then the menu: the helper text always + # sits directly above the prompt instead of scrolling off the top. _print_stage_summary(stage, context) ui.tip(stage) @@ -1988,14 +1975,19 @@ async def stage_hook(stage: str, context: PipelineContext) -> bool: if stage in (PipelineStage.collection, PipelineStage.processing, PipelineStage.generation, PipelineStage.quality): while True: - action = await prompts.ask_stage_action(name) + action = await prompts.ask_stage_action(name, next_name) if action == "adjust": await _adjust_settings(context) continue # re-show the same menu so the user can continue/export/pause next + if action == "explain": + ui.pipeline_overview_panel(current_stage=stage, next_stage=nxt) + if nxt in _STAGE_PRE_DESCRIPTIONS: + ui.stage_description(next_name, _stage_map[nxt][1], _STAGE_TOTAL, + _STAGE_PRE_DESCRIPTIONS[nxt]) + continue if action == "export": - await _quick_export(context, stage) - cont = await prompts.ask_confirm("Continue pipeline after export?") - return cont + await _export_flow(context, stage) + continue # back to the checkpoint: continue, export more, or stop if action == "pause": ui.info(f"Session saved. Resume with: [bold]dataforge resume {context.session_id[:8]}[/]") return False @@ -2014,19 +2006,19 @@ async def post_discovery_hook(stage: str, context: PipelineContext) -> bool: ui.warn("Discovery returned 0 URLs. Check that the site has a reachable sitemap or provide a direct sitemap URL.") return False - ui.url_table(context.discovered_urls) - - # Locale hint + # The review screen lists the URLs itself; printing them here as well + # only pushed the review's key hints off screen. + notice = "" lang_groups = _detect_language_groups(context.discovered_urls) if len(lang_groups) >= 2: - ui.language_groups_panel(lang_groups, total) - ui.warn( - "Multiple language variants detected — use the filter " - "(e.g. [bold]/en/[/]) in the review step to keep one locale." + top = ", ".join(f"{k} ({v})" for k, v in sorted(lang_groups.items(), key=lambda x: -x[1])[:5]) + notice = ( + f"Several languages found: {top}. " + "To keep one, filter it, e.g. [bold]f /en/[/], then [bold]none[/] / [bold]all[/] as needed." ) from .url_review import run_url_review - selected = await run_url_review(context.discovered_urls) + selected = await run_url_review(context.discovered_urls, notice=notice) if not selected: ui.warn("No URLs selected — returning to menu.") @@ -2049,6 +2041,9 @@ async def combined_hook(stage: str, context: PipelineContext) -> bool: async def pre_stage_hook(stage: str, context: PipelineContext) -> bool: name, step = _stage_map.get(stage, (stage, 0)) detail = _STAGE_PRE_DESCRIPTIONS.get(stage, "") + # Each stage starts on a fresh screen so progress output from the + # previous one does not pile up. + ui.screen(f"Step {step} of {_STAGE_TOTAL}: {name}", context.session_name) ui.stage_description(name, step, _STAGE_TOTAL, detail) if stage == PipelineStage.generation: from dataforge.generators.templates import build_prompt @@ -2165,25 +2160,88 @@ def _print_stage_summary(stage: str, ctx: PipelineContext) -> None: ui.quality_distribution_panel(scores, ctx.quality_threshold) -async def _quick_export(ctx: PipelineContext, stage: str) -> None: +async def _export_flow(ctx: PipelineContext, stage: str = "", approved_only: bool = True) -> None: + """Export whatever the session has so far: training data and/or plain files. + + Works at every stage. Before generation there are no samples, but the + discovered URLs, scraped pages and chunks are still worth having as + Markdown, text or JSON. + """ + from datetime import datetime + + from dataforge.exporters.stage_data import FORMAT_LABELS, KINDS, available_data + s = ctx.settings - targets = await prompts.ask_export_targets( - hf_configured=bool(s.huggingface_token), - kg_configured=bool(s.kaggle_username), - ) - export_kw: dict = {"targets": targets, "stage_snapshot": stage} - if "huggingface" in targets: - export_kw["hf_repo_id"] = await prompts.ask_hf_repo() - export_kw["hf_private"] = await prompts.ask_hf_private() - if "kaggle" in targets: - export_kw["kaggle_slug"] = await prompts.ask_kaggle_slug(s.kaggle_username) - export_kw["kaggle_title"] = ctx.session_name + have = available_data(s.db_path, ctx.session_id) + if not have: + ui.warn("Nothing to export yet: this project has no URLs, pages or samples.") + return + + with open_session(s.db_path) as db: + n_approved = len(db.exec( + select(SyntheticSample.id) + .where(SyntheticSample.session_id == ctx.session_id) + .where(SyntheticSample.approved == True) # noqa: E712 + ).all()) + n_training = n_approved if approved_only else have.get("samples", 0) + + options: list[tuple[str, str, str]] = [] + if n_training: + options.append(( + f"Training dataset ({n_training} samples)", "training", + "JSONL, Parquet, CSV and Unsloth files ready for fine-tuning. " + "Can also upload to HuggingFace or Kaggle.", + )) + descriptions = { + "urls": "The list of links found on the site.", + "pages": "The text of each page, cleaned. Good for reading, search or RAG.", + "chunks": "Pages split into smaller passages, with their source URL.", + "samples": "Every generated sample, approved or not, with its quality score.", + } + # Latest stage first: that is usually what people want right now. + for kind in reversed(list(KINDS)): + if kind in have: + options.append((f"{KINDS[kind].label} ({have[kind]})", kind, descriptions[kind])) + + picked = await prompts.ask_export_what(options) + if not picked: + ui.info("Nothing exported.") + return - from dataforge.agents.exporter import ExporterAgent - agent = ExporterAgent(ctx, **export_kw) - ctx_out = await agent.run() - ui.info("Export complete!") - ui.export_summary(ctx_out.export_records) + export_dir = ctx.session_dir() / "exports" / datetime.now().strftime("%Y%m%d_%H%M%S") + written: list[tuple[str, Path]] = [] + for kind in picked: + if kind == "training": + continue + fmts = await prompts.ask_export_formats(KINDS[kind].label, KINDS[kind].formats, FORMAT_LABELS) + if not fmts: + continue + from dataforge.exporters.stage_data import export_stage_data + paths = export_stage_data(s.db_path, ctx.session_id, kind, fmts, export_dir) + written += [(f"{KINDS[kind].label} ({fmt})", p) for fmt, p in paths.items()] + + if "training" in picked: + targets = await prompts.ask_export_targets( + hf_configured=bool(s.huggingface_token), + kg_configured=bool(s.kaggle_username and s.kaggle_key), + ) or ["local"] + export_kw: dict = {"targets": targets, "approved_only": approved_only, + "stage_snapshot": stage or ctx.current_stage} + if "huggingface" in targets: + export_kw["hf_repo_id"] = await prompts.ask_hf_repo() + export_kw["hf_private"] = await prompts.ask_hf_private() + if "kaggle" in targets: + export_kw["kaggle_slug"] = await prompts.ask_kaggle_slug(s.kaggle_username) + export_kw["kaggle_title"] = ctx.session_name + from dataforge.agents.exporter import ExporterAgent + ctx_out = await ExporterAgent(ctx, **export_kw).run() + ui.export_summary(ctx_out.export_records) + + if written: + ui.info(f"Files saved in [bold]{export_dir}[/]") + for label, path in written: + suffix = "/" if path.is_dir() else "" + ui.success(f"{label} [dim]{path.name}{suffix}[/]") async def _ask_export_config(s) -> dict: @@ -2243,6 +2301,168 @@ async def _collect_urls() -> list[str] | None: return [] +async def _pick_session(question: str) -> PipelineSession | None: + """Choose any session, newest first. None on back/cancel or when there are none.""" + import questionary + s = get_settings() + with open_session(s.db_path) as db: + rows = sorted(db.exec(select(PipelineSession)).all(), + key=lambda x: x.created_at, reverse=True) + if not rows: + ui.info("No projects yet. Start one from the main menu.") + return None + choices = [ + questionary.Choice( + f"{x.name} [{x.id[:8]}]", value=x.id, + description=f"Reached: {x.stage} · status: {x.status} · " + f"started {x.created_at.strftime('%Y-%m-%d %H:%M')}", + ) + for x in rows + ] + choices.append(questionary.Choice("← Back", value="__back__")) + sid = await questionary.select(question, choices=choices, **prompts._q()).ask_async() + if not sid or sid == "__back__": + return None + return next(x for x in rows if x.id == sid) + + +def _context_for(session: PipelineSession) -> PipelineContext: + from dataforge.agents import PipelineContext + return PipelineContext( + session_id=session.id, + session_name=session.name, + goal=session.goal, + format=DataFormat(session.format), + seed_urls=session.seed_url_list(), + settings=get_settings(), + ) + + +async def _export_menu() -> None: + session = await _pick_session("Export data from which project?") + if session: + await _export_flow(_context_for(session), session.stage) + + +async def _quick_scrape_flow() -> None: + """Menu version of `dataforge scrape`: pages and tables, no AI, no API key.""" + from datetime import datetime + + import questionary + + from dataforge import scrape as qs + + ui.screen("Scrape pages", "no AI, no API key needed") + ui.info("Fetches the exact pages you give it (it does not crawl the whole site) " + "and saves their text and tables.") + urls = await _collect_urls() + if not urls: + return + fmts = await questionary.checkbox( + "Save as:", + choices=[ + questionary.Choice("Markdown, one .md file per page", value="md", checked=True), + questionary.Choice("JSON Lines (pages.jsonl, tables.jsonl)", value="jsonl", checked=True), + questionary.Choice("Tables as CSV and JSON", value="csv", checked=True), + ], + instruction="(Space to tick, Enter to confirm)", + validate=lambda v: bool(v) or "Tick at least one format", + **prompts._q(), + ).ask_async() + if not fmts: + return + default_out = str(Path("scrape") / datetime.now().strftime("%Y%m%d-%H%M%S")) + out = await prompts.ask_output_dir(default_out) + if out is None: + return + s = get_settings() + with ui.console.status(f"Fetching {len(urls)} page(s)…"): + pages = await qs.scrape_urls(urls, rate_limit=s.rate_limit, tables="csv" in fmts) + out_dir = Path(out).expanduser() + qs.write_outputs(pages, out_dir, tuple(fmts)) + _print_scrape_results(pages, out_dir) + + +def _print_scrape_results(pages: list, out_dir: Path) -> None: + ok = [p for p in pages if p.ok] + for p in pages: + if p.ok: + ui.success(f"{p.url} [dim]{p.title[:60]}[/] {p.word_count} words, " + f"{len(p.tables)} table{'s' if len(p.tables) != 1 else ''}") + for w in p.warnings: + ui.warn(f" {w}") + else: + ui.error(f"{p.url} {p.status}") + for p in ok: + for k, t in enumerate(p.tables[:3], 1): + _print_table_preview(t, f"{p.title or p.url}: table {k}") + n_tables = sum(len(p.tables) for p in ok) + ui.info(f"{len(ok)}/{len(pages)} page(s), {n_tables} table(s) saved to " + f"[bold]{out_dir.resolve()}[/]") + + +_MENU_HELP = """\ +[bold]What each option does[/] + + [bold cyan]Scrape pages (no AI)[/] + Give it one or more page URLs; get their text as Markdown/JSON and any + tables as CSV. No API key, costs nothing. + + [bold cyan]Build an AI training dataset[/] + The full pipeline. Finds every page on a site, lets you pick which to + keep, scrapes them, splits the text into chunks, has an LLM write + training examples, filters them for quality, and exports. + You can stop after any step and export what you have so far + (pages as Markdown or text, chunks as JSON, and so on). + Needs an LLM: an API key, or a local model through Ollama. + + [bold cyan]Continue a paused project[/] + Pick up a project you stopped, from the step where it stopped. + + [bold cyan]Browse my data[/] + Look through the URLs, pages, chunks and samples of any project. + + [bold cyan]Export data[/] + Save any project's data as Markdown, text, JSON, CSV or a training set. + + [bold cyan]Settings and tools[/] + Set your LLM provider and API key, see system info, update or uninstall. + +[bold]Moving around[/] + Arrow keys to move, Enter to choose, Space to tick in lists with boxes. + Ctrl+C goes back or stops safely; your progress is always saved. + Set DATAFORGE_NO_CLEAR=1 to keep all output instead of clearing the screen.""" + + +def _show_help() -> None: + ui.screen("Help") + ui.console.print(_MENU_HELP) + ui.console.print("") + ui.pipeline_overview_panel() + + +async def _settings_menu() -> str: + """Less common actions, grouped so the main menu stays short.""" + import questionary + return await questionary.select( + "Settings and tools", + choices=[ + questionary.Choice("LLM provider and API key", value="config", + description="Choose OpenAI, Anthropic, Gemini, Groq, Ollama… and save your key."), + questionary.Choice("List all projects", value="sessions", + description="Every project in this folder's database, with its status."), + questionary.Choice("Pipeline steps and project status", value="plan", + description="What each step does and where your latest project is."), + questionary.Choice("System info", value="info", + description="Version, paths, provider and machine resources."), + questionary.Choice("Update DataForge", value="update"), + questionary.Choice("Uninstall DataForge", value="uninstall"), + questionary.Choice("← Back", value="back"), + ], + **prompts._q(), + ).ask_async() or "back" + + async def _main_menu() -> str: import questionary s = get_settings() @@ -2255,35 +2475,57 @@ async def _main_menu() -> str: except Exception: paused_count, total_sessions = 0, 0 - resume_label = "Resume a paused session" + (f" ({paused_count} waiting)" if paused_count else "") - + # Worded as tasks, each with a one-line explanation that questionary shows + # in the same spot under the list as the user moves through it. choices = [ - questionary.Choice("New pipeline", value="new"), - questionary.Choice(resume_label, value="resume"), + questionary.Choice( + "Scrape pages (no AI)", value="scrape", + description="Save the text and tables of pages you choose. No API key needed.", + ), + questionary.Choice( + "Build an AI training dataset", value="new", + description="Crawl a site, pick pages, and turn them into training examples. " + "You can stop and export after any step.", + ), ] - if total_sessions > 0: + if paused_count: + choices.append(questionary.Choice( + f"Continue a paused project ({paused_count} waiting)", value="resume", + description="Pick up where you stopped. Nothing is lost.", + )) + if total_sessions: choices += [ - questionary.Choice("Explore sessions", value="explore"), - questionary.Choice("List all sessions", value="sessions"), + questionary.Choice( + "Browse my data", value="explore", + description="Look through the URLs, pages, chunks and samples of a project.", + ), + questionary.Choice( + "Export data", value="export", + description="Save a project's data as Markdown, text, JSON, CSV or a training set.", + ), ] choices += [ - questionary.Choice("Plan / project status", value="plan"), - questionary.Choice("Configure LLM provider", value="config"), - questionary.Choice("System info", value="info"), - questionary.Choice("Update DataForge", value="update"), - questionary.Choice("Uninstall DataForge", value="uninstall"), - questionary.Choice("Exit", value="exit"), + questionary.Choice( + "Settings and tools", value="settings", + description="LLM provider and API key, system info, update, uninstall.", + ), + questionary.Choice("Help", value="help", description="What every option does, and how the pipeline works."), + questionary.Choice("Clear the screen", value="clear", description="Tidy up the terminal and show this menu again."), + questionary.Choice("Exit", value="exit"), ] ui.console.print("") result = await questionary.select( "What would you like to do?", choices=choices, + instruction="(arrows to move, Enter to choose)", **prompts._q(), ).ask_async() if result is None: raise typer.Exit() + if result == "settings": + return await _settings_menu() return result diff --git a/src/dataforge/cli/prompts.py b/src/dataforge/cli/prompts.py index 6bde7bd..c33cac6 100644 --- a/src/dataforge/cli/prompts.py +++ b/src/dataforge/cli/prompts.py @@ -123,12 +123,13 @@ async def ask_input_method() -> str | None: async def ask_single_url() -> str | None: url = await questionary.text( "Enter URL:", - validate=lambda v: _valid_url(v) or "Enter a valid http(s) URL", + instruction="(paste it; quotes are fine)", + validate=lambda v: _valid_url(clean_input(v)) or "Enter a valid http(s) URL, e.g. https://example.com", **_q(), ).ask_async() if url is None: return None - return url.strip() + return clean_input(url) async def ask_multiple_urls() -> list[str] | None: @@ -139,18 +140,19 @@ async def ask_multiple_urls() -> list[str] | None: ).ask_async() if raw is None: return None - return [u.strip() for u in raw.splitlines() if u.strip() and _valid_url(u.strip())] + return [u for u in _split_urls(raw) if _valid_url(u)] async def ask_file_path() -> Path | None: path = await questionary.path( "Path to URL file:", - validate=lambda v: Path(v).exists() or "File not found", + instruction="(drag the file here or paste its path; quotes are fine)", + validate=lambda v: _path_exists(v) or f"File not found: {clean_input(v)}", **_q(), ).ask_async() if path is None: return None - return Path(path) + return Path(clean_input(path)).expanduser() async def ask_goal() -> str | None: @@ -223,7 +225,7 @@ async def ask_output_dir(default: str = "./output") -> str | None: ).ask_async() if answer is None: return None - return answer.strip() or default + return clean_input(answer) or default async def ask_session_name() -> str | None: @@ -274,17 +276,36 @@ async def ask_skip_known(domain: str) -> bool: # ── Stage checkpoints ───────────────────────────────────────────────────────── -async def ask_stage_action(stage: str) -> str: +async def ask_stage_action(stage: str, next_stage: str = "") -> str: + """Checkpoint menu between stages. Each choice explains itself below the list.""" + nxt = f": {next_stage}" if next_stage else "" return await questionary.select( - f"Stage '{stage}' complete. What next?", + f"{stage} finished. What next?", choices=[ - questionary.Choice("Continue to next stage", value="continue"), - questionary.Choice("Adjust settings (model / output dir)", value="adjust"), - questionary.Choice("Export available data now", value="export"), - questionary.Choice("Save and exit (resume later)", value="pause"), + questionary.Choice( + f"Continue{nxt}", value="continue", + description=f"Run the next step{nxt}.", + ), + questionary.Choice( + "Export what I have so far", value="export", + description="Save pages, chunks or samples as Markdown, text, JSON or CSV. " + "You can keep going afterwards.", + ), + questionary.Choice( + "What happens next?", value="explain", + description="Explain every step and what the next one will do.", + ), + questionary.Choice( + "Change model or output folder", value="adjust", + description="Takes effect from the next step.", + ), + questionary.Choice( + "Stop here (resume later)", value="pause", + description="Everything is saved. Pick 'Continue a paused project' from the menu later.", + ), ], **_q(), - ).ask_async() + ).ask_async() or "pause" async def ask_adjust_settings_target() -> str | None: @@ -303,6 +324,34 @@ async def ask_adjust_settings_target() -> str | None: # ── Export ──────────────────────────────────────────────────────────────────── +async def ask_export_what(options: list[tuple[str, str, str]]) -> list[str]: + """Pick what to export. ``options`` is [(label, value, description)].""" + answer = await questionary.checkbox( + "What do you want to export?", + choices=[ + questionary.Choice(label, value=value, description=desc, checked=(i == 0)) + for i, (label, value, desc) in enumerate(options) + ], + instruction="(Space to tick, Enter to confirm)", + **_q(), + ).ask_async() + return answer or [] + + +async def ask_export_formats(label: str, formats: tuple[str, ...], names: dict[str, str]) -> list[str]: + answer = await questionary.checkbox( + f"{label}: which formats?", + choices=[ + questionary.Choice(names.get(f, f), value=f, checked=(i == 0)) + for i, f in enumerate(formats) + ], + instruction="(Space to tick, Enter to confirm)", + validate=lambda v: bool(v) or "Tick at least one format", + **_q(), + ).ask_async() + return answer or [] + + async def ask_export_targets(hf_configured: bool, kg_configured: bool) -> list[str]: choices = [questionary.Choice("Local files (JSONL / Parquet / CSV)", value="local")] if hf_configured: @@ -414,4 +463,39 @@ def _valid_url(v: str) -> bool: def read_url_file(path: Path) -> list[str]: lines = path.read_text(encoding="utf-8").splitlines() - return [line.strip() for line in lines if line.strip() and not line.startswith("#") and _valid_url(line.strip())] + return [ + u for line in lines + if not line.strip().startswith("#") + for u in _split_urls(line) + if _valid_url(u) + ] + + +# Straight and curly quotes people paste around paths and URLs. +_QUOTE_PAIRS = {'"': '"', "'": "'", "`": "`", "“": "”", "‘": "’", "<": ">"} + + +def clean_input(value: str) -> str: + r"""Normalise a pasted path or URL. + + Handles what terminals and file managers add around a path: surrounding + quotes (Windows "Copy as path"), PowerShell's drag-and-drop form + ``& 'C:\dir\file.txt'``, and stray whitespace. + """ + v = (value or "").strip() + if v.startswith("& "): + v = v[2:].strip() + while len(v) >= 2 and _QUOTE_PAIRS.get(v[0]) == v[-1]: + v = v[1:-1].strip() + return v + + +def _split_urls(text: str) -> list[str]: + """URLs from free text: one per line, or separated by commas or spaces.""" + parts = text.replace(",", " ").split() + return [c for c in (clean_input(p) for p in parts) if c] + + +def _path_exists(value: str) -> bool: + v = clean_input(value) + return bool(v) and Path(v).expanduser().exists() diff --git a/src/dataforge/cli/ui.py b/src/dataforge/cli/ui.py index c33d278..d0eb623 100644 --- a/src/dataforge/cli/ui.py +++ b/src/dataforge/cli/ui.py @@ -67,6 +67,36 @@ def section(title: str) -> None: console.rule(f"[bold]{title}[/]", style="cyan") +def _can_clear() -> bool: + import os + # DATAFORGE_NO_CLEAR=1 keeps the full scrollback (useful when debugging). + return console.is_terminal and not os.getenv("DATAFORGE_NO_CLEAR") + + +def screen(title: str, subtitle: str = "") -> None: + """Start a fresh screen: clear the terminal, then a one-line title bar. + + Every interactive screen follows the same layout so users always know + where to look: title at the top, content in the middle, then status + messages, then key hints directly above the prompt (see ``hints``). + """ + if _can_clear(): + console.clear() + else: + console.print("") + head = f"[bold cyan]DataForge[/] [bold]{title}[/]" + if subtitle: + head += f" [dim]{subtitle}[/]" + console.print(head) + console.rule(style="cyan") + + +def hints(text: str) -> None: + """Key hints for the prompt below. Always the last thing before a prompt.""" + console.rule(style="dim") + console.print(f"[dim]{text}[/]") + + def url_table(urls: list[str], selected: list[bool] | None = None, max_rows: int = 30) -> None: t = Table(box=box.SIMPLE_HEAD, show_footer=False) t.add_column("#", style="dim", width=5) diff --git a/src/dataforge/cli/url_review.py b/src/dataforge/cli/url_review.py index 3e75bce..f3b9cd3 100644 --- a/src/dataforge/cli/url_review.py +++ b/src/dataforge/cli/url_review.py @@ -18,13 +18,12 @@ _PAGE_SIZE = 30 _HELP_TEXT = ( - "[dim] n[/dim]=next [dim]p[/dim]=prev " - "[dim]f [/dim]=filter " - "[dim]x <#>[/dim]=deselect [dim]+ <#>[/dim]=select " - "[dim]x [/dim]=range " - "[dim]all[/dim]/[dim]none[/dim] " - "[dim]i <#>[/dim]=inspect " - "[dim]done[/dim]=proceed [dim]back[/dim]=menu [dim]?[/dim]=help" + "Type a command, then Enter:\n" + " [bold]done[/] scrape the ticked pages [bold]f blog[/] show only URLs containing 'blog'\n" + " [bold]x 3[/] untick row 3 ([bold]x 3-8[/] for a range) [bold]+ 3[/] tick it again " + "[bold]all[/] / [bold]none[/]\n" + " [bold]n[/] / [bold]p[/] next / previous page [bold]i 3[/] details of row 3 " + "[bold]back[/] main menu [bold]?[/] full help" ) _HELP_FULL = """\ @@ -79,6 +78,7 @@ def __init__(self, urls: list[str]) -> None: self._selected: set[str] = set(urls) # selected URLs (mutable) self._page: int = 0 # 0-indexed current page self._filter: str = "" # active filter pattern + self._detail: str = "" # extra block for the next render (help, inspect) # ── Display ─────────────────────────────────────────────────────────────── @@ -86,18 +86,16 @@ def _page_slice(self) -> list[str]: start = self._page * _PAGE_SIZE return self._view[start : start + _PAGE_SIZE] - def _render(self) -> None: + def _render(self, message: str = "", notice: str = "") -> None: + """Draw the whole review screen: table, then messages, then key hints.""" page_urls = self._page_slice() total_pages = _page_count(len(self._view)) n_selected = len(self._selected & set(self._view)) - # Header line - filter_hint = f" [dim]filter: {self._filter}[/]" if self._filter else "" - ui.console.print( - f"\n[bold cyan]URL Review[/]" - f" Page [bold]{self._page + 1}[/] / {total_pages}" - f" [green]{n_selected}[/] / {len(self._view)} selected" - f"{filter_hint}" + filter_hint = f" · filter: {self._filter}" if self._filter else "" + ui.screen( + "Choose pages to scrape", + f"page {self._page + 1}/{total_pages} · {n_selected}/{len(self._view)} ticked{filter_hint}", ) table = Table(box=box.SIMPLE, show_header=True, header_style="bold") @@ -112,7 +110,14 @@ def _render(self) -> None: table.add_row(str(row_num), check, _label(url)) ui.console.print(table) - ui.info(_HELP_TEXT) + if notice: + ui.warn(notice) + if self._detail: + ui.console.print(self._detail) + self._detail = "" + if message: + ui.info(message) + ui.hints(_HELP_TEXT) # ── Command handlers ────────────────────────────────────────────────────── @@ -175,8 +180,9 @@ def _do_inspect(self, spec: str) -> str: return f"Row {n} not on current page." url = page_urls[local] parsed = urlparse(url) - ui.console.print( - f"\n[bold]URL #{n}[/]\n" + # Shown on the next redraw, so it survives the screen clear. + self._detail = ( + f"[bold]URL #{n}[/]\n" f" Full URL: [cyan]{url}[/]\n" f" Scheme: {parsed.scheme}\n" f" Host: {parsed.netloc}\n" @@ -248,8 +254,8 @@ def handle(self, raw: str) -> tuple[str, bool]: if cmd in ("q", "quit", "back", "b"): return "", None # type: ignore[return-value] - if cmd == "?": - ui.console.print(_HELP_FULL) + if cmd in ("?", "h", "help"): + self._detail = _HELP_FULL return "", False return f"Unknown command '{raw.strip()}' — type ? for help.", False @@ -260,9 +266,13 @@ def selected_urls(self) -> list[str]: return sorted(self._selected, key=lambda u: order.get(u, 0)) -async def run_url_review(urls: list[str]) -> list[str]: +async def run_url_review(urls: list[str], notice: str = "") -> list[str]: """Interactive URL review: paginated browser with filter, select, and inspect. + The screen is redrawn after every command instead of appended to, so the + table, messages and key hints stay in the same place. ``notice`` (e.g. a + mixed-language warning) is shown on every redraw. + Returns the user-approved subset in original discovery order. """ if not urls: @@ -272,8 +282,9 @@ async def run_url_review(urls: list[str]) -> list[str]: session: PromptSession[str] = PromptSession(style=_PT_STYLE) + msg = "" while True: - reviewer._render() + reviewer._render(msg, notice) try: raw = await session.prompt_async(" review> ") @@ -288,32 +299,25 @@ async def run_url_review(urls: list[str]) -> list[str]: ui.info("Returning to menu.") return [] - if msg: - ui.info(msg) - if done: selected = reviewer.selected_urls() total = len(urls) picked = len(selected) if picked == 0: - ui.warn("No URLs selected — returning to review.") + msg = "No URLs ticked yet. Tick some with + or all, then type done." continue - ui.info( - f"[bold]{picked}[/] / {total} URL{'s' if picked != 1 else ''} selected." - ) - _QSTYLE = QStyle([ ("qmark", "fg:cyan bold"), ("question", "bold"), ("answer", "fg:cyan bold"), ("pointer", "fg:cyan bold"), ]) confirmed = await questionary.confirm( - f"Proceed with {picked} URL{'s' if picked != 1 else ''}?", + f"Scrape {picked} of {total} page{'s' if picked != 1 else ''}?", default=True, style=_QSTYLE, ).ask_async() if confirmed: return selected - # User said no — keep browsing + msg = "Selection kept. Keep editing, then type done." diff --git a/src/dataforge/exporters/stage_data.py b/src/dataforge/exporters/stage_data.py new file mode 100644 index 0000000..ac536b3 --- /dev/null +++ b/src/dataforge/exporters/stage_data.py @@ -0,0 +1,235 @@ +"""Export what a session has produced so far, in general-purpose formats. + +The training export (``ExporterAgent``) only knows about generated samples. +Each earlier stage also leaves useful data behind: the discovered URL list, +the scraped pages as Markdown, and the cleaned chunks. This module writes any +of those out as plain files anyone can open (Markdown, text, JSON, CSV), with +no LLM involved, so a session is useful even if it never reaches generation. +""" +from __future__ import annotations + +import csv +import json +import re +from dataclasses import dataclass +from pathlib import Path + +from sqlmodel import select + +from dataforge.storage import ( + DiscoveredURL, + ProcessedChunk, + ScrapedPage, + SyntheticSample, + open_session, +) + +from .local import write_jsonl + + +@dataclass(frozen=True) +class DataKind: + key: str + label: str + formats: tuple[str, ...] + + +# Ordered by pipeline stage. The first format of each kind is the default. +KINDS: dict[str, DataKind] = { + "urls": DataKind("urls", "Discovered URLs", ("txt", "csv", "json")), + "pages": DataKind("pages", "Scraped pages", ("md", "txt", "json", "jsonl")), + "chunks": DataKind("chunks", "Text chunks", ("jsonl", "json", "txt")), + "samples": DataKind("samples", "Generated samples (all)", ("jsonl", "json")), +} + +FORMAT_LABELS = { + "md": "Markdown (one .md file per page)", + "txt": "Plain text", + "json": "JSON (one file)", + "jsonl": "JSON Lines (one record per line)", + "csv": "CSV spreadsheet", +} + + +def available_data(db_path: Path, session_id: str) -> dict[str, int]: + """Row counts per data kind for a session; kinds with no rows are omitted.""" + with open_session(db_path) as db: + counts = { + "urls": len(db.exec(select(DiscoveredURL.id).where(DiscoveredURL.session_id == session_id)).all()), + "pages": len(db.exec(select(ScrapedPage.id).where(ScrapedPage.session_id == session_id)).all()), + "chunks": len(db.exec(select(ProcessedChunk.id).where(ProcessedChunk.session_id == session_id)).all()), + "samples": len(db.exec(select(SyntheticSample.id).where(SyntheticSample.session_id == session_id)).all()), + } + return {k: n for k, n in counts.items() if n} + + +def export_stage_data( + db_path: Path, + session_id: str, + kind: str, + formats: list[str], + out_dir: Path, +) -> dict[str, Path]: + """Write one kind of session data in each requested format. + + Returns ``{format: path}``; for per-page Markdown/text the path is a folder. + """ + if kind not in KINDS: + raise ValueError(f"Unknown data kind: {kind}") + bad = [f for f in formats if f not in KINDS[kind].formats] + if bad: + raise ValueError(f"{kind} cannot be exported as: {', '.join(bad)}") + + records = _load(db_path, session_id, kind) + out_dir.mkdir(parents=True, exist_ok=True, mode=0o750) + writer = _WRITERS[kind] + return {fmt: writer(records, fmt, out_dir) for fmt in formats} + + +# -- Loading ------------------------------------------------------------------- + +def _load(db_path: Path, session_id: str, kind: str) -> list[dict]: + with open_session(db_path) as db: + if kind == "urls": + rows = db.exec(select(DiscoveredURL).where(DiscoveredURL.session_id == session_id)).all() + return [ + {"url": r.url, "source": r.source, "selected": r.selected, "scraped": r.scraped} + for r in rows + ] + if kind == "pages": + rows = db.exec(select(ScrapedPage).where(ScrapedPage.session_id == session_id)).all() + return [ + { + "id": r.id, "url": r.url, "title": r.title, "author": r.author, + "published_date": r.published_date, "word_count": r.word_count, + "markdown": _read(r.raw_path), + } + for r in rows + ] + if kind == "chunks": + rows = db.exec(select(ProcessedChunk).where(ProcessedChunk.session_id == session_id)).all() + return [ + { + "id": r.id, "page_id": r.page_id, "chunk_index": r.chunk_index, + "source_url": r.parsed_meta().get("source_url", ""), + "title": r.parsed_meta().get("title", ""), + "token_count": r.token_count, "content": r.content, + } + for r in rows + ] + rows = db.exec(select(SyntheticSample).where(SyntheticSample.session_id == session_id)).all() + return [ + { + "id": r.id, "chunk_id": r.chunk_id, "format": r.format, + "messages": r.messages(), "quality_score": r.quality_score, + "approved": r.approved, + } + for r in rows + ] + + +def _read(path: str) -> str: + if not path: + return "" + try: + return Path(path).read_text(encoding="utf-8") + except OSError: + return "" + + +# -- Writers ------------------------------------------------------------------- + +def _write_json(records: list[dict], path: Path) -> Path: + path.write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding="utf-8") + return path + + +def _write_urls(records: list[dict], fmt: str, out: Path) -> Path: + if fmt == "txt": + path = out / "urls.txt" + path.write_text("".join(r["url"] + "\n" for r in records), encoding="utf-8") + return path + if fmt == "csv": + path = out / "urls.csv" + with path.open("w", newline="", encoding="utf-8") as f: + w = csv.DictWriter(f, fieldnames=["url", "source", "selected", "scraped"]) + w.writeheader() + w.writerows(records) + return path + return _write_json(records, out / "urls.json") + + +def _write_pages(records: list[dict], fmt: str, out: Path) -> Path: + if fmt in ("md", "txt"): + folder = out / ("pages_md" if fmt == "md" else "pages_txt") + folder.mkdir(parents=True, exist_ok=True) + used: set[str] = set() + for i, r in enumerate(records, 1): + name = _unique(_slug(r["title"] or r["url"]) or f"page-{i}", used) + if fmt == "md": + header = ( + f"---\ntitle: {json.dumps(r['title'], ensure_ascii=False)}\nsource: {r['url']}\n" + f"words: {r['word_count']}\n---\n\n" + ) + (folder / f"{name}.md").write_text(header + r["markdown"], encoding="utf-8") + else: + body = f"{r['title']}\n{r['url']}\n\n{_plain(r['markdown'])}\n" + (folder / f"{name}.txt").write_text(body, encoding="utf-8") + return folder + if fmt == "jsonl": + path = out / "pages.jsonl" + write_jsonl(records, path) + return path + return _write_json(records, out / "pages.json") + + +def _write_chunks(records: list[dict], fmt: str, out: Path) -> Path: + if fmt == "txt": + path = out / "chunks.txt" + path.write_text( + "\n\n".join(f"### {r['source_url']} (chunk {r['chunk_index']})\n{r['content']}" for r in records), + encoding="utf-8", + ) + return path + if fmt == "jsonl": + path = out / "chunks.jsonl" + write_jsonl(records, path) + return path + return _write_json(records, out / "chunks.json") + + +def _write_samples(records: list[dict], fmt: str, out: Path) -> Path: + if fmt == "jsonl": + path = out / "samples_all.jsonl" + write_jsonl(records, path) + return path + return _write_json(records, out / "samples_all.json") + + +_WRITERS = {"urls": _write_urls, "pages": _write_pages, "chunks": _write_chunks, "samples": _write_samples} + + +# -- Helpers ------------------------------------------------------------------- + +_SLUG_RE = re.compile(r"[^a-z0-9]+") +_MD_LINK_RE = re.compile(r"!?\[([^\]]*)\]\([^)]*\)") +_MD_MARK_RE = re.compile(r"^\s{0,3}(#{1,6}\s+|>\s?|[-*+]\s+)|[*_`]{1,3}", re.M) + + +def _slug(text: str, max_len: int = 60) -> str: + text = re.sub(r"^https?://", "", text.lower()) + return _SLUG_RE.sub("-", text).strip("-")[:max_len].strip("-") + + +def _unique(name: str, used: set[str]) -> str: + candidate, n = name, 2 + while candidate in used: + candidate = f"{name}-{n}" + n += 1 + used.add(candidate) + return candidate + + +def _plain(markdown: str) -> str: + """Markdown to readable plain text: keep link text, drop markup.""" + return _MD_MARK_RE.sub("", _MD_LINK_RE.sub(r"\1", markdown)).strip() diff --git a/src/dataforge/generators/llm.py b/src/dataforge/generators/llm.py index 4c5f385..fa53c7c 100644 --- a/src/dataforge/generators/llm.py +++ b/src/dataforge/generators/llm.py @@ -7,13 +7,17 @@ from typing import Any import litellm -from tenacity import retry, stop_after_attempt, wait_exponential +from tenacity import retry, retry_if_not_exception_type, stop_after_attempt, wait_exponential from dataforge.config import endpoint_kwargs, get_settings, litellm_model, model_supports_thinking from dataforge.utils import get_logger +from dataforge.utils.errors import MissingCredentialError log = get_logger("llm") litellm.set_verbose = False +# Otherwise LiteLLM prints a "Give Feedback / Get Help" banner on every +# failed call, which floods the terminal when a key is missing or invalid. +litellm.suppress_debug_info = True class BudgetTracker: @@ -122,6 +126,9 @@ async def complete( @retry( stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=20), + # A bad or missing key fails the same way every time; retrying only + # adds backoff delay before the run can stop. + retry=retry_if_not_exception_type(MissingCredentialError), reraise=True, ) async def _complete_impl( @@ -165,8 +172,7 @@ async def _complete_impl( self.usage.errors += 1 msg = str(exc).lower() # Surface actionable errors without a raw traceback - if "auth" in msg or "api key" in msg or "401" in msg or "403" in msg: - from dataforge.utils.errors import MissingCredentialError + if "auth" in msg or "api key" in msg or "api_key" in msg or "401" in msg or "403" in msg: _KEY_MAP = { "openai": "OPENAI_API_KEY", "anthropic": "ANTHROPIC_API_KEY", "google": "GEMINI_API_KEY", @@ -295,8 +301,7 @@ async def complete_stream( except Exception as exc: self.usage.errors += 1 msg = str(exc).lower() - if "auth" in msg or "api key" in msg or "401" in msg or "403" in msg: - from dataforge.utils.errors import MissingCredentialError + if "auth" in msg or "api key" in msg or "api_key" in msg or "401" in msg or "403" in msg: _KEY_MAP = { "openai": "OPENAI_API_KEY", "anthropic": "ANTHROPIC_API_KEY", "google": "GEMINI_API_KEY", diff --git a/src/dataforge/generators/synthetic.py b/src/dataforge/generators/synthetic.py index 1f6320c..ddd5a03 100644 --- a/src/dataforge/generators/synthetic.py +++ b/src/dataforge/generators/synthetic.py @@ -8,7 +8,7 @@ from dataforge.processors.formatter import DataRecord from dataforge.utils import get_logger -from dataforge.utils.errors import BudgetExceededError +from dataforge.utils.errors import BudgetExceededError, LLMConnectionError, MissingCredentialError from .contracts import GenerationParseError, to_message_dict, to_message_list from .llm import LLMClient @@ -60,6 +60,10 @@ async def generate_from_chunk( except BudgetExceededError as exc: log.debug(f"Skipping chunk {record.chunk_id}: {exc}") return [] + except (MissingCredentialError, LLMConnectionError): + # Fatal for every chunk, not just this one: let the caller stop the run + # instead of repeating the same failure once per remaining chunk. + raise except Exception as exc: log.warning(f"Generation failed for chunk {record.chunk_id}: {exc}") return [] @@ -115,11 +119,15 @@ def on_thinking(chunk: str) -> None: messages, on_thinking=on_thinking, ) + except (MissingCredentialError, LLMConnectionError): + raise except Exception as exc: log.warning(f"Thinking-stream generation failed for chunk {record.chunk_id}: {exc}") # Fall back to regular completion try: resp = await client.complete(messages) + except (MissingCredentialError, LLMConnectionError): + raise except Exception: return [] @@ -138,20 +146,36 @@ async def generate_batch( concurrency: int = 3, ) -> AsyncIterator[GeneratedSample]: sem = asyncio.Semaphore(concurrency) + stopped = False async def _worker(rec: DataRecord): + nonlocal stopped async with sem: - return await generate_from_chunk( - client, rec, - format=format, goal=goal, - n_per_chunk=n_per_chunk, custom_system=custom_system, - ) + # A fatal error in another worker means this call would fail too. + if stopped: + return [] + try: + return await generate_from_chunk( + client, rec, + format=format, goal=goal, + n_per_chunk=n_per_chunk, custom_system=custom_system, + ) + except (MissingCredentialError, LLMConnectionError): + stopped = True + raise tasks = [asyncio.create_task(_worker(r)) for r in records] - for coro in asyncio.as_completed(tasks): - samples = await coro - for s in samples: - yield s + try: + for coro in asyncio.as_completed(tasks): + samples = await coro + for s in samples: + yield s + finally: + # On a fatal error (or the consumer stopping early) the queued chunks + # must not keep calling the LLM in the background. + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) # ── Parsing helpers ──────────────────────────────────────────────────────────── diff --git a/tests/test_generation_fatal_errors.py b/tests/test_generation_fatal_errors.py new file mode 100644 index 0000000..52a9441 --- /dev/null +++ b/tests/test_generation_fatal_errors.py @@ -0,0 +1,51 @@ +"""A missing or invalid API key must stop generation once, not once per chunk.""" +from __future__ import annotations + +import litellm +import pytest + +from dataforge.generators.llm import LLMClient +from dataforge.generators.synthetic import generate_batch, generate_from_chunk +from dataforge.processors.formatter import DataRecord +from dataforge.utils.errors import MissingCredentialError + + +def _record(i: int) -> DataRecord: + return DataRecord( + chunk_id=i, source_url="https://x", title="t", + content="Floods are dangerous.", token_count=5, + ) + + +@pytest.fixture +def auth_failure(monkeypatch): + calls = {"n": 0} + + async def _fail(**kwargs): + calls["n"] += 1 + raise Exception("AuthenticationError: Incorrect API key provided") + + monkeypatch.setattr(litellm, "acompletion", _fail) + return calls + + +async def test_generate_from_chunk_propagates_missing_credential(auth_failure): + with pytest.raises(MissingCredentialError): + await generate_from_chunk(LLMClient(), _record(1), format="qa", goal="g", n_per_chunk=1) + + +async def test_auth_failure_is_not_retried(auth_failure): + with pytest.raises(MissingCredentialError): + await LLMClient().complete([{"role": "user", "content": "hi"}]) + assert auth_failure["n"] == 1 + + +async def test_generate_batch_stops_on_first_auth_failure(auth_failure): + records = [_record(i) for i in range(50)] + with pytest.raises(MissingCredentialError): + async for _ in generate_batch( + LLMClient(), records, format="qa", goal="g", n_per_chunk=1, concurrency=3 + ): + pass + # Only the calls already in flight may run; the rest are cancelled. + assert auth_failure["n"] <= 3 diff --git a/tests/test_input_cleaning.py b/tests/test_input_cleaning.py new file mode 100644 index 0000000..ce9d6f6 --- /dev/null +++ b/tests/test_input_cleaning.py @@ -0,0 +1,43 @@ +"""Pasted paths and URLs are accepted with the quotes terminals add around them.""" +from __future__ import annotations + +import pytest + +from dataforge.cli.prompts import _path_exists, _split_urls, clean_input, read_url_file + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ('"C:\\my files\\urls.txt"', "C:\\my files\\urls.txt"), # Windows "Copy as path" + ("& 'C:\\my files\\urls.txt'", "C:\\my files\\urls.txt"), # PowerShell drag-and-drop + ("'/home/me/urls.txt'", "/home/me/urls.txt"), + (" https://example.com ", "https://example.com"), + ("", "https://example.com"), + ("\u201chttps://example.com\u201d", "https://example.com"), # curly quotes + ("'\"https://example.com\"'", "https://example.com"), # nested + ("'unbalanced", "'unbalanced"), # left alone + ("", ""), + ], +) +def test_clean_input(raw, expected): + assert clean_input(raw) == expected + + +def test_split_urls_accepts_quotes_commas_and_lines(): + raw = '"https://a.com", https://b.com\n \'https://c.com\' https://d.com' + assert _split_urls(raw) == ["https://a.com", "https://b.com", "https://c.com", "https://d.com"] + + +def test_path_exists_with_quotes(tmp_path): + f = tmp_path / "my urls.txt" + f.write_text("x", encoding="utf-8") + assert _path_exists(f'"{f}"') + assert _path_exists(f"& '{f}'") + assert not _path_exists('"' + str(tmp_path / "missing.txt") + '"') + + +def test_read_url_file_strips_quotes_and_comments(tmp_path): + f = tmp_path / "urls.txt" + f.write_text('# comment\n"https://a.com"\nhttps://b.com, https://c.com\nnot-a-url\n', encoding="utf-8") + assert read_url_file(f) == ["https://a.com", "https://b.com", "https://c.com"] diff --git a/tests/test_stage_data_export.py b/tests/test_stage_data_export.py new file mode 100644 index 0000000..ab500ae --- /dev/null +++ b/tests/test_stage_data_export.py @@ -0,0 +1,72 @@ +"""Every stage's data can be exported as plain files, with no samples required.""" +from __future__ import annotations + +import json + +import pytest + +from dataforge.exporters.stage_data import available_data, export_stage_data +from dataforge.storage import DiscoveredURL, ProcessedChunk, ScrapedPage, open_session + +SID = "sess-1" + + +@pytest.fixture +def seeded(tmp_settings, tmp_path): + raw = tmp_path / "page.md" + raw.write_text("# Floods\n\nFloods are **dangerous**. See [the guide](https://x/guide).", encoding="utf-8") + with open_session(tmp_settings.db_path) as db: + db.add(DiscoveredURL(session_id=SID, url="https://x/floods")) + db.add(DiscoveredURL(session_id=SID, url="https://x/fires", selected=False)) + db.add(ScrapedPage(session_id=SID, url_id=1, url="https://x/floods", + title="Floods", raw_path=str(raw), word_count=6)) + db.add(ScrapedPage(session_id=SID, url_id=2, url="https://x/floods-2", + title="Floods", raw_path=str(raw), word_count=6)) + db.add(ProcessedChunk(session_id=SID, page_id=1, content="Floods are dangerous.", + metadata_json=json.dumps({"source_url": "https://x/floods"}))) + db.commit() + return tmp_settings + + +def test_available_data_counts_only_what_exists(seeded): + assert available_data(seeded.db_path, SID) == {"urls": 2, "pages": 2, "chunks": 1} + assert available_data(seeded.db_path, "other") == {} + + +def test_pages_as_markdown_one_file_each_with_front_matter(seeded, tmp_path): + out = export_stage_data(seeded.db_path, SID, "pages", ["md"], tmp_path / "exp") + files = sorted(out["md"].glob("*.md")) + # Same title twice: names must not collide. + assert [f.name for f in files] == ["floods-2.md", "floods.md"] + text = files[1].read_text(encoding="utf-8") + assert text.startswith('---\ntitle: "Floods"\nsource: https://x/floods') + assert "Floods are **dangerous**" in text + + +def test_pages_as_plain_text_strip_markdown(seeded, tmp_path): + out = export_stage_data(seeded.db_path, SID, "pages", ["txt"], tmp_path / "exp") + body = (out["txt"] / "floods.txt").read_text(encoding="utf-8") + assert "Floods are dangerous. See the guide." in body + assert "**" not in body and "](" not in body + + +def test_pages_and_chunks_as_json(seeded, tmp_path): + out = export_stage_data(seeded.db_path, SID, "pages", ["json", "jsonl"], tmp_path / "exp") + pages = json.loads(out["json"].read_text(encoding="utf-8")) + assert {p["url"] for p in pages} == {"https://x/floods", "https://x/floods-2"} + assert len(out["jsonl"].read_text(encoding="utf-8").splitlines()) == 2 + + out = export_stage_data(seeded.db_path, SID, "chunks", ["json"], tmp_path / "exp") + chunks = json.loads(out["json"].read_text(encoding="utf-8")) + assert chunks[0]["source_url"] == "https://x/floods" + + +def test_urls_as_text(seeded, tmp_path): + out = export_stage_data(seeded.db_path, SID, "urls", ["txt", "csv"], tmp_path / "exp") + assert out["txt"].read_text(encoding="utf-8").split() == ["https://x/floods", "https://x/fires"] + assert out["csv"].read_text(encoding="utf-8").startswith("url,source,selected,scraped") + + +def test_rejects_format_that_does_not_fit_the_data(seeded, tmp_path): + with pytest.raises(ValueError): + export_stage_data(seeded.db_path, SID, "urls", ["md"], tmp_path / "exp") diff --git a/uv.lock b/uv.lock index 883e5e5..147529c 100644 --- a/uv.lock +++ b/uv.lock @@ -1559,7 +1559,7 @@ wheels = [ [[package]] name = "llm-web-crawler" -version = "2.5.0" +version = "2.6.0" source = { editable = "." } dependencies = [ { name = "aiofiles" },