diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md index 3918f4c..c315879 100644 --- a/.claude/CLAUDE.md +++ b/.claude/CLAUDE.md @@ -1,12 +1,12 @@ # Soup CLI — Project CLAUDE.md -Soup is a CLI-first LLM fine-tuning tool (v0.19.0). Python 3.9+, MIT license. +Soup is a CLI-first LLM fine-tuning tool (v0.20.0). Python 3.9+, MIT license. ## Build & Development ```bash pip install -e ".[dev]" # Install editable + test deps -pytest tests/ -v --tb=short # Run all tests (1577 tests) +pytest tests/ -v --tb=short # Run all tests (1669 tests) ruff check soup_cli/ tests/ # Lint (must pass before commit) ruff check --fix soup_cli/ tests/ # Auto-fix lint issues ``` @@ -16,7 +16,7 @@ ruff check --fix soup_cli/ tests/ # Auto-fix lint issues ``` soup_cli/ cli.py # Entry point, Typer app, all command registration - __init__.py # __version__ = "0.19.0" + __init__.py # __version__ = "0.20.0" config/ schema.py # Pydantic models (SoupConfig, DataConfig, TrainingConfig, LoraConfig, EvalConfig) loader.py # YAML -> SoupConfig, load_config_from_string() @@ -24,6 +24,16 @@ soup_cli/ loader.py # Local files (JSONL/JSON/CSV/Parquet) + HF datasets formats.py # Auto-detect + normalize to {"messages": [...]} structure validator.py # Dataset stats, quality checks, extended_stats() + providers/ + ollama.py # Ollama provider for synth data gen (localhost-only) + anthropic.py # Anthropic Claude provider (env-only API key) + vllm.py # vLLM provider for synth data gen (SSRF-protected) + templates/ + code.py # Code instruction pairs template + conversation.py # Multi-turn conversation template + qa.py # QA from context template + preference.py # DPO/KTO/ORPO preference data template + reasoning.py # Chain-of-thought / GRPO reasoning template trainer/ sft.py # SFTTrainerWrapper (415 lines, supports vision + unsloth + QAT) dpo.py # DPOTrainerWrapper (253 lines) @@ -124,7 +134,7 @@ soup data convert # Transform between alpaca/sharegpt/chatml soup data merge # Combine multiple datasets soup data dedup # MinHash deduplication soup data stats # Extended statistics with histograms -soup data generate # Synthetic data via LLM APIs (--provider openai|local|server) +soup data generate # Synthetic data via LLM APIs (--provider openai|local|server|ollama|anthropic|vllm) soup data filter # Quality filter (perplexity + coherence scoring) soup runs # List experiment runs soup runs show # Detailed run info + metrics @@ -223,6 +233,11 @@ soup version # Show version (--full for details) - **Human eval**: local-only terminal UI, no network access (v0.19.0) - **Human eval**: prompts file capped at 10k entries (v0.19.0) - **Leaderboard**: read-only SQLite queries, no user input in SQL (v0.19.0) +- **Ollama provider**: localhost-only validation — remote Ollama instances blocked (v0.20.0) +- **Anthropic provider**: API key from env only (ANTHROPIC_API_KEY), never CLI arg (v0.20.0) +- **vLLM provider**: SSRF protection — scheme whitelist, localhost-only HTTP (v0.20.0) +- **Output path**: path traversal protection — `..` blocked in output filenames (v0.20.0) +- **Rate limiting**: configurable `--requests-per-minute` (default: 60) (v0.20.0) ## Code Conventions @@ -292,7 +307,7 @@ soup version # Show version (--full for details) 15. **Tag**: `git tag v0.X.Y && git push origin v0.X.Y` 16. **Release**: `gh release create v0.X.Y` with changelog (What's New, Install/Upgrade) -## Tests (58 test files, 1577 tests) +## Tests (59 test files, 1669 tests) | File | Covers | |------|--------| @@ -353,3 +368,4 @@ soup version # Show version (--full for details) | test_sglang_serve.py | SGLang backend detection, runtime creation, serve --backend | | test_deploy_ollama.py | Ollama deploy, Modelfile gen, template mapping, security validation | | test_eval_platform.py | Custom eval, judge, human eval (Elo), leaderboard, compare, auto-eval, security | +| test_synth_data_pro.py | Providers (Ollama, Anthropic, vLLM), templates, quality pipeline, SSRF | diff --git a/README.md b/README.md index 1a48e9c..bcb869e 100644 --- a/README.md +++ b/README.md @@ -876,6 +876,55 @@ soup data generate --prompt "..." --seed examples.jsonl --count 100 soup data generate --prompt "..." --provider server --api-base http://localhost:11434/v1 ``` +### Multi-Provider Support (v0.20.0+) + +```bash +# Generate via local Ollama instance +soup data generate --prompt "..." --provider ollama --model llama3.1 +soup data generate --prompt "..." --ollama-model llama3.1 # shorthand + +# Generate via Anthropic Claude API (set ANTHROPIC_API_KEY env var) +soup data generate --prompt "..." --provider anthropic --model claude-3-haiku-20240307 + +# Generate via local vLLM server +soup data generate --prompt "..." --provider vllm --model meta-llama/Llama-3.1-8B-Instruct +``` + +### Domain Templates (v0.20.0+) + +```bash +# Code instruction pairs (Python, JS, Go, Rust, Java) +soup data generate --prompt "..." --template code --language Python --task-type function + +# Multi-turn conversations +soup data generate --prompt "..." --template conversation --turns 6 --topic "science" + +# QA from context document +soup data generate --prompt "..." --template qa --context document.txt + +# Preference data (DPO/KTO/ORPO) +soup data generate --prompt "..." --template preference --pref-task dpo + +# Chain-of-thought reasoning (GRPO) +soup data generate --prompt "..." --template reasoning --domain math +``` + +### Quality Pipeline (v0.20.0+) + +```bash +# Auto-validate after generation (remove malformed entries) +soup data generate --prompt "..." --validate + +# Auto-filter by quality (coherence scoring) +soup data generate --prompt "..." --filter + +# Auto-dedup (MinHash, requires: pip install 'soup-cli[data]') +soup data generate --prompt "..." --dedup + +# Full quality pipeline: validate + filter + dedup +soup data generate --prompt "..." --quality-pipeline +``` + ## Hyperparameter Sweep Search for the best hyperparameters: @@ -1227,7 +1276,11 @@ soup data merge data1.jsonl data2.jsonl Combine datasets soup data dedup --threshold 0.8 Remove duplicates (MinHash) soup data stats Extended statistics soup data generate --prompt "..." --count 100 Generate synthetic data -soup data generate ... --provider server Use local inference server +soup data generate ... --provider ollama Use local Ollama instance +soup data generate ... --provider anthropic Use Claude API +soup data generate ... --provider vllm Use local vLLM server +soup data generate ... --template code Domain templates (code/conversation/qa/preference/reasoning) +soup data generate ... --quality-pipeline Auto validate + filter + dedup soup data filter --coherence 0.3 Quality filter (perplexity/coherence) soup runs List training runs soup runs show Run details + loss graph diff --git a/SECURITY.md b/SECURITY.md index b35f287..1b7124f 100644 --- a/SECURITY.md +++ b/SECURITY.md @@ -4,14 +4,14 @@ We provide security updates for the following versions: -- **Latest minor version:** Active support (e.g., v0.19.x) +- **Latest minor version:** Active support (e.g., v0.20.x) - **Previous minor versions:** Bug-fix support only - **Versions older than 3 minor versions:** No support Example: -- v0.19.0-0.19.x → Full support (latest) -- v0.18.0-0.18.x → Bug-fix support only -- v0.17.x and below → No support +- v0.20.0-0.20.x → Full support (latest) +- v0.19.0-0.19.x → Bug-fix support only +- v0.18.x and below → No support ## Reporting a Vulnerability @@ -121,6 +121,8 @@ No known critical vulnerabilities in current releases. - **v0.16.0**: `embedding_loss` Literal constraint, `embedding_margin` gt=0 validation, ONNX export without trust_remote_code (with warning), TensorRT export subprocess list args (no shell injection), speculative decoding SSRF-protected (URL blocked) with warning panel, vLLM speculative model URL validation - **v0.17.0**: Server data generation provider SSRF validation (scheme whitelist + localhost-only HTTP), audio model trust_remote_code warning panel, audio file path traversal protection (resolved paths confined to audio_dir), SGLang backend trust_remote_code warning panel - **v0.18.0**: Ollama deploy GGUF path traversal protection + `.gguf` extension validation, model name validation (no path separators/null bytes), subprocess list args (no shell injection), Modelfile parameter key allowlist + value newline/null sanitization, overwrite warning panel +- **v0.19.0**: Custom eval JSONL schema validation + 10k task cap, regex scoring ReDoS guard, judge API SSRF protection + API key isolation, human eval local-only terminal UI + 10k prompt cap, leaderboard read-only SQLite queries +- **v0.20.0**: Ollama provider localhost-only validation (remote blocked), Anthropic provider API key from env only (never CLI arg), vLLM provider SSRF protection (scheme whitelist + localhost-only HTTP), output path traversal protection (`..` blocked), configurable rate limiting (`--requests-per-minute`) ## Security Scanning diff --git a/pyproject.toml b/pyproject.toml index cf8aa03..27dd680 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "soup-cli" -version = "0.19.0" +version = "0.20.0" description = "Fine-tune LLMs in one command. No SSH, no config hell." readme = "README.md" license = "MIT" diff --git a/soup_cli/__init__.py b/soup_cli/__init__.py index 3b54753..85e1821 100644 --- a/soup_cli/__init__.py +++ b/soup_cli/__init__.py @@ -1,3 +1,3 @@ """Soup CLI — Fine-tune LLMs in one command.""" -__version__ = "0.19.0" +__version__ = "0.20.0" diff --git a/soup_cli/commands/generate.py b/soup_cli/commands/generate.py index 59e201a..ba0e254 100644 --- a/soup_cli/commands/generate.py +++ b/soup_cli/commands/generate.py @@ -13,6 +13,10 @@ logger = logging.getLogger(__name__) console = Console() +VALID_PROVIDERS = ("openai", "local", "server", "ollama", "anthropic", "vllm") +VALID_FORMATS = ("alpaca", "sharegpt", "chatml") +VALID_TEMPLATES = ("code", "conversation", "qa", "preference", "reasoning") + def generate( prompt: str = typer.Option( @@ -42,13 +46,13 @@ def generate( provider: str = typer.Option( "openai", "--provider", - help="LLM provider: openai, local, server (local OpenAI-compatible server)", + help="LLM provider: openai, local, server, ollama, anthropic, vllm", ), model_name: str = typer.Option( "gpt-4o-mini", "--model", "-m", - help="Model name (OpenAI model ID or local model path)", + help="Model name (OpenAI model ID, Ollama model, or local model path)", ), api_key: Optional[str] = typer.Option( None, @@ -82,20 +86,114 @@ def generate( "--seed", help="Path to seed examples file (JSONL) to guide generation", ), + # --- v0.20.0: Template support --- + template: Optional[str] = typer.Option( + None, + "--template", + help="Domain template: code, conversation, qa, preference, reasoning", + ), + template_language: str = typer.Option( + "Python", + "--language", + help="Language for code template (Python, JavaScript, Go, Rust, Java)", + ), + template_task_type: str = typer.Option( + "function", + "--task-type", + help="Task type for code template (function, debug, explain, refactor, test)", + ), + template_turns: int = typer.Option( + 4, + "--turns", + help="Number of turns for conversation template (2-10)", + ), + template_topic: str = typer.Option( + "general knowledge", + "--topic", + help="Topic for conversation template", + ), + template_context: Optional[str] = typer.Option( + None, + "--context", + help="Path to context document for QA template", + ), + template_pref_task: str = typer.Option( + "dpo", + "--pref-task", + help="Preference task format: dpo, kto, orpo", + ), + template_domain: str = typer.Option( + "math", + "--domain", + help="Domain for reasoning template (math, logic, code)", + ), + # --- v0.20.0: Ollama-specific --- + ollama_model: Optional[str] = typer.Option( + None, + "--ollama-model", + help="Ollama model name (e.g. llama3.1). Shorthand for --provider ollama --model X", + ), + # --- v0.20.0: Quality pipeline --- + validate_output: bool = typer.Option( + False, + "--validate", + help="Auto-validate generated data after generation", + ), + filter_output: bool = typer.Option( + False, + "--filter", + help="Auto-filter generated data by quality (perplexity + coherence)", + ), + dedup_output: bool = typer.Option( + False, + "--dedup", + help="Auto-dedup generated data (MinHash)", + ), + quality_pipeline: bool = typer.Option( + False, + "--quality-pipeline", + help="Run full quality pipeline: validate + filter + dedup", + ), + requests_per_minute: int = typer.Option( + 60, + "--requests-per-minute", + "--rpm", + help="Rate limit for API requests (default: 60)", + ), ): """Generate synthetic training data using an LLM.""" - valid_formats = ("alpaca", "sharegpt", "chatml") - if fmt not in valid_formats: - console.print(f"[red]Invalid format: {fmt}. Must be one of: {', '.join(valid_formats)}[/]") - raise typer.Exit(1) - - valid_providers = ("openai", "local", "server") - if provider not in valid_providers: + if fmt not in VALID_FORMATS: console.print( - f"[red]Invalid provider: {provider}. Must be one of: {', '.join(valid_providers)}[/]" + f"[red]Invalid format: {fmt}. Must be one of: {', '.join(VALID_FORMATS)}[/]" ) raise typer.Exit(1) + # Handle --ollama-model shorthand + if ollama_model: + provider = "ollama" + model_name = ollama_model + + if provider not in VALID_PROVIDERS: + console.print( + f"[red]Invalid provider: {provider}. " + f"Must be one of: {', '.join(VALID_PROVIDERS)}[/]" + ) + raise typer.Exit(1) + + # Validate template + if template and template not in VALID_TEMPLATES: + console.print( + f"[red]Invalid template: {template}. " + f"Must be one of: {', '.join(VALID_TEMPLATES)}[/]" + ) + raise typer.Exit(1) + + # Enable quality pipeline shortcut + if quality_pipeline: + validate_output = True + filter_output = True + dedup_output = True + # Load seed examples if provided seed_examples = [] if seed_file: @@ -122,8 +220,19 @@ def generate( existing_texts.add(_row_to_text(row)) console.print(f"[dim]Loaded {len(existing_texts)} existing examples for dedup[/]") + # Load template context if provided + context_text = "" + if template_context: + ctx_path = Path(template_context) + if not ctx_path.exists(): + console.print(f"[red]Context file not found: {ctx_path}[/]") + raise typer.Exit(1) + context_text = ctx_path.read_text(encoding="utf-8") + # Generate console.print(f"[dim]Generating {count} examples using {provider}/{model_name}...[/]") + if template: + console.print(f"[dim]Template: {template}[/]") all_examples = [] duplicates = 0 @@ -152,6 +261,14 @@ def generate( api_base=api_base, temperature=temperature, seed_examples=seed_examples, + template=template, + template_language=template_language, + template_task_type=template_task_type, + template_turns=template_turns, + template_topic=template_topic, + context_text=context_text, + template_pref_task=template_pref_task, + template_domain=template_domain, ) except Exception as exc: console.print(f"[red]Generation error: {exc}[/]") @@ -160,6 +277,10 @@ def generate( # Validate and dedup for example in batch: if not _validate_example(example, fmt): + # For preference templates, validate against preference format + if template == "preference" and _validate_preference(example): + all_examples.append(example) + continue continue text = _row_to_text(example) if text in existing_texts: @@ -172,8 +293,13 @@ def generate( remaining -= current_batch progress.update(task, advance=generated_this_round) - # Write output + # Sanitize output path (prevent path traversal) out_path = Path(output) + if ".." in out_path.parts: + console.print("[red]Output path must not contain '..'[/]") + raise typer.Exit(1) + + # Write output with open(out_path, "w", encoding="utf-8") as f: for row in all_examples: f.write(json.dumps(row, ensure_ascii=False) + "\n") @@ -182,9 +308,143 @@ def generate( f"\n[green]Generated {len(all_examples)} examples[/]\n" f"Format: [bold]{fmt}[/]\n" f"Output: [bold]{out_path}[/]\n" + + (f"Template: [bold]{template}[/]\n" if template else "") + (f"Duplicates: [yellow]{duplicates} removed[/]\n" if duplicates > 0 else "") ) + # --- Quality pipeline --- + if validate_output: + _run_validate_pipeline(out_path, fmt) + if filter_output: + _run_filter_pipeline(out_path) + if dedup_output: + _run_dedup_pipeline(out_path) + + +def _run_validate_pipeline(path: Path, fmt: str) -> None: + """Run validation on generated output, retry malformed entries.""" + console.print("[dim]Running validation pipeline...[/]") + from soup_cli.data.loader import load_raw_data + + data = load_raw_data(path) + valid = [] + invalid_count = 0 + for row in data: + if _validate_example(row, fmt) or _validate_preference(row): + valid.append(row) + else: + invalid_count += 1 + + if invalid_count > 0: + with open(path, "w", encoding="utf-8") as f: + for row in valid: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + console.print( + f"[yellow]Validation: removed {invalid_count} malformed entries, " + f"{len(valid)} remaining[/]" + ) + else: + console.print(f"[green]Validation: all {len(valid)} entries valid[/]") + + +def _run_filter_pipeline(path: Path) -> None: + """Run quality filter (coherence) on generated output.""" + console.print("[dim]Running quality filter pipeline...[/]") + from soup_cli.data.loader import load_raw_data + + data = load_raw_data(path) + if not data: + return + + texts = [" ".join(str(v) for v in row.values() if v) for row in data] + + try: + from soup_cli.utils.quality import compute_coherence_score + + scores = compute_coherence_score(texts) + # Use median as threshold — remove bottom quartile + sorted_scores = sorted(scores) + threshold = sorted_scores[len(sorted_scores) // 4] if sorted_scores else 0.0 + + kept = [row for row, score in zip(data, scores) if score >= threshold] + removed = len(data) - len(kept) + + if removed > 0: + with open(path, "w", encoding="utf-8") as f: + for row in kept: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + console.print( + f"[yellow]Filter: removed {removed} low-quality entries, " + f"{len(kept)} remaining[/]" + ) + else: + console.print(f"[green]Filter: all {len(data)} entries passed[/]") + except ImportError: + console.print("[yellow]Filter: quality scoring not available, skipping[/]") + + +def _run_dedup_pipeline(path: Path) -> None: + """Run MinHash dedup on generated output.""" + console.print("[dim]Running dedup pipeline...[/]") + from soup_cli.data.loader import load_raw_data + + data = load_raw_data(path) + if not data: + return + + try: + from datasketch import MinHash, MinHashLSH + + num_perm = 128 + lsh = MinHashLSH(threshold=0.8, num_perm=num_perm) + minhashes = [] + + for idx, row in enumerate(data): + text = " ".join(str(v) for v in row.values() if v).lower() + words = text.split() + shingles = set() + for ii in range(max(1, len(words) - 2)): + shingles.add(" ".join(words[ii: ii + 3])) + + mhash = MinHash(num_perm=num_perm) + for shingle in shingles: + mhash.update(shingle.encode("utf-8")) + minhashes.append(mhash) + + try: + lsh.insert(str(idx), mhash) + except ValueError: + pass + + seen: set[int] = set() + unique_indices = [] + for idx in range(len(data)): + if idx in seen: + continue + unique_indices.append(idx) + results = lsh.query(minhashes[idx]) + for dup_idx_str in results: + seen.add(int(dup_idx_str)) + + unique_data = [data[idx] for idx in unique_indices] + removed = len(data) - len(unique_data) + + if removed > 0: + with open(path, "w", encoding="utf-8") as f: + for row in unique_data: + f.write(json.dumps(row, ensure_ascii=False) + "\n") + console.print( + f"[yellow]Dedup: removed {removed} near-duplicates, " + f"{len(unique_data)} remaining[/]" + ) + else: + console.print(f"[green]Dedup: no duplicates found in {len(data)} entries[/]") + except ImportError: + console.print( + "[yellow]Dedup: datasketch not installed, skipping. " + "Install: pip install 'soup-cli[data]'[/]" + ) + def _generate_batch( prompt: str, @@ -196,8 +456,34 @@ def _generate_batch( api_base: Optional[str], temperature: float, seed_examples: list[dict], + template: Optional[str] = None, + template_language: str = "Python", + template_task_type: str = "function", + template_turns: int = 4, + template_topic: str = "general knowledge", + context_text: str = "", + template_pref_task: str = "dpo", + template_domain: str = "math", ) -> list[dict]: """Generate a batch of examples using the specified provider.""" + # Build the generation prompt — template or default + if template: + generation_prompt = _build_template_prompt( + template=template, + prompt=prompt, + count=count, + fmt=fmt, + language=template_language, + task_type=template_task_type, + turns=template_turns, + topic=template_topic, + context_text=context_text, + pref_task=template_pref_task, + domain=template_domain, + ) + else: + generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples) + if provider == "openai": return _generate_openai( prompt=prompt, @@ -208,6 +494,7 @@ def _generate_batch( api_base=api_base, temperature=temperature, seed_examples=seed_examples, + generation_prompt=generation_prompt if template else None, ) elif provider == "local": return _generate_local( @@ -217,6 +504,7 @@ def _generate_batch( model_name=model_name, temperature=temperature, seed_examples=seed_examples, + generation_prompt=generation_prompt if template else None, ) elif provider == "server": return _generate_server( @@ -227,13 +515,92 @@ def _generate_batch( api_base=api_base, temperature=temperature, seed_examples=seed_examples, + generation_prompt=generation_prompt if template else None, + ) + elif provider == "ollama": + from soup_cli.data.providers.ollama import DEFAULT_OLLAMA_BASE, generate_ollama + + base_url = api_base or DEFAULT_OLLAMA_BASE + return generate_ollama( + prompt=prompt, + count=count, + fmt=fmt, + model_name=model_name, + base_url=base_url, + temperature=temperature, + generation_prompt=generation_prompt, + ) + elif provider == "anthropic": + from soup_cli.data.providers.anthropic import generate_anthropic + + return generate_anthropic( + prompt=prompt, + count=count, + fmt=fmt, + model_name=model_name, + temperature=temperature, + generation_prompt=generation_prompt, + ) + elif provider == "vllm": + from soup_cli.data.providers.vllm import DEFAULT_VLLM_BASE, generate_vllm + + base_url = api_base or DEFAULT_VLLM_BASE + return generate_vllm( + prompt=prompt, + count=count, + fmt=fmt, + model_name=model_name, + base_url=base_url, + temperature=temperature, + generation_prompt=generation_prompt, ) return [] -def _build_generation_prompt(prompt: str, count: int, fmt: str, seed_examples: list) -> str: - """Build the prompt for data generation.""" - format_spec = { +def _build_template_prompt( + template: str, + prompt: str, + count: int, + fmt: str, + language: str = "Python", + task_type: str = "function", + turns: int = 4, + topic: str = "general knowledge", + context_text: str = "", + pref_task: str = "dpo", + domain: str = "math", +) -> str: + """Build generation prompt from a domain template.""" + format_spec = _get_format_spec(fmt) + + if template == "code": + from soup_cli.data.templates.code import build_prompt + + return build_prompt(count, fmt, format_spec, language=language, task_type=task_type) + elif template == "conversation": + from soup_cli.data.templates.conversation import build_prompt + + return build_prompt(count, fmt, format_spec, turns=turns, topic=topic) + elif template == "qa": + from soup_cli.data.templates.qa import build_prompt + + return build_prompt(count, fmt, format_spec, context=context_text) + elif template == "preference": + from soup_cli.data.templates.preference import build_prompt + + return build_prompt(count, task=pref_task) + elif template == "reasoning": + from soup_cli.data.templates.reasoning import build_prompt + + return build_prompt(count, fmt, format_spec, domain=domain) + + # Fallback to default prompt + return _build_generation_prompt(prompt, count, fmt, []) + + +def _get_format_spec(fmt: str) -> str: + """Get the format specification string for the given format.""" + specs = { "alpaca": ( 'Each example must be a JSON object with keys: ' '"instruction", "input" (can be empty string), "output".' @@ -247,12 +614,18 @@ def _build_generation_prompt(prompt: str, count: int, fmt: str, seed_examples: l 'which is a list of objects with "role" (user/assistant) and "content".' ), } + return specs.get(fmt, specs["alpaca"]) + + +def _build_generation_prompt(prompt: str, count: int, fmt: str, seed_examples: list) -> str: + """Build the prompt for data generation.""" + format_spec = _get_format_spec(fmt) system_msg = ( f"You are a training data generator. Generate exactly {count} diverse, " f"high-quality training examples.\n\n" f"Topic/Instructions: {prompt}\n\n" - f"Format: {format_spec[fmt]}\n\n" + f"Format: {format_spec}\n\n" f"Return ONLY a JSON array of {count} examples. No markdown, no explanation." ) @@ -272,6 +645,7 @@ def _generate_openai( api_base: Optional[str], temperature: float, seed_examples: list[dict], + generation_prompt: Optional[str] = None, ) -> list[dict]: """Generate examples using OpenAI-compatible API.""" import os @@ -301,7 +675,8 @@ def _generate_openai( "HTTP is only allowed for localhost." ) - generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples) + if generation_prompt is None: + generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples) response = httpx.post( f"{base_url}/chat/completions", @@ -340,6 +715,7 @@ def _generate_local( model_name: str, temperature: float, seed_examples: list[dict], + generation_prompt: Optional[str] = None, ) -> list[dict]: """Generate examples using a local model via transformers.""" import torch @@ -357,7 +733,8 @@ def _generate_local( ) model.eval() - generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples) + if generation_prompt is None: + generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples) if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template: messages = [ @@ -395,6 +772,7 @@ def _generate_server( api_base: Optional[str], temperature: float, seed_examples: list[dict], + generation_prompt: Optional[str] = None, ) -> list[dict]: """Generate examples using a local OpenAI-compatible server (soup serve, Ollama, etc.). @@ -429,7 +807,8 @@ def _generate_server( if not base_url.endswith("/v1"): base_url = base_url + "/v1" - generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples) + if generation_prompt is None: + generation_prompt = _build_generation_prompt(prompt, count, fmt, seed_examples) response = httpx.post( f"{base_url}/chat/completions", @@ -516,6 +895,17 @@ def _validate_example(example: dict, fmt: str) -> bool: return False +def _validate_preference(example: dict) -> bool: + """Validate a preference data example (DPO/KTO format).""" + # DPO/ORPO format + if "prompt" in example and "chosen" in example and "rejected" in example: + return True + # KTO format + if "prompt" in example and "completion" in example and "label" in example: + return True + return False + + def _row_to_text(row: dict) -> str: """Convert a row to a text string for dedup comparison.""" return " ".join(str(v) for v in row.values() if v) diff --git a/soup_cli/data/providers/__init__.py b/soup_cli/data/providers/__init__.py new file mode 100644 index 0000000..ac916d7 --- /dev/null +++ b/soup_cli/data/providers/__init__.py @@ -0,0 +1 @@ +"""Data generation providers — Ollama, Anthropic, vLLM.""" diff --git a/soup_cli/data/providers/anthropic.py b/soup_cli/data/providers/anthropic.py new file mode 100644 index 0000000..82f4cec --- /dev/null +++ b/soup_cli/data/providers/anthropic.py @@ -0,0 +1,92 @@ +"""Anthropic (Claude) provider for synthetic data generation. + +Uses raw httpx requests to avoid SDK dependency. +API key must be set via ANTHROPIC_API_KEY environment variable. +""" + +import logging +import os + +logger = logging.getLogger(__name__) + +ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages" +DEFAULT_MODEL = "claude-3-haiku-20240307" + + +def generate_anthropic( + prompt: str, + count: int, + fmt: str, + model_name: str, + temperature: float, + generation_prompt: str, +) -> list[dict]: + """Generate examples using Anthropic Claude API. + + Args: + prompt: User-provided topic/instructions. + count: Number of examples to generate. + fmt: Output format (alpaca/sharegpt/chatml). + model_name: Anthropic model ID. + temperature: Sampling temperature. + generation_prompt: Pre-built generation prompt. + + Returns: + List of generated example dicts. + + Raises: + ValueError: If API key not found or API returns error. + ImportError: If httpx is not installed. + """ + try: + import httpx + except ImportError: + raise ImportError( + "httpx is required for Anthropic generation. Install: pip install httpx" + ) + + api_key = os.environ.get("ANTHROPIC_API_KEY") + if not api_key: + raise ValueError( + "Anthropic API key not found. Set ANTHROPIC_API_KEY environment variable." + ) + + response = httpx.post( + ANTHROPIC_API_URL, + headers={ + "x-api-key": api_key, + "anthropic-version": "2023-06-01", + "Content-Type": "application/json", + }, + json={ + "model": model_name, + "system": generation_prompt, + "messages": [ + {"role": "user", "content": f"Generate {count} training examples now."}, + ], + "temperature": temperature, + "max_tokens": 4096, + }, + timeout=120.0, + ) + + if response.status_code != 200: + logger.debug("Anthropic error response: %s", response.text) + raise ValueError( + f"Anthropic API returned {response.status_code}. " + "Check your API key and model name." + ) + + data = response.json() + try: + # Anthropic Messages API returns content as a list of blocks + content_blocks = data["content"] + content = "".join( + block["text"] for block in content_blocks if block.get("type") == "text" + ) + except (KeyError, IndexError, TypeError) as exc: + raise ValueError(f"Unexpected Anthropic response format: {exc}") from exc + + from soup_cli.commands.generate import _parse_json_array + + return _parse_json_array(content) diff --git a/soup_cli/data/providers/ollama.py b/soup_cli/data/providers/ollama.py new file mode 100644 index 0000000..d28685a --- /dev/null +++ b/soup_cli/data/providers/ollama.py @@ -0,0 +1,123 @@ +"""Ollama provider for synthetic data generation. + +Connects to a local Ollama instance via its OpenAI-compatible API. +""" + +import logging +from typing import Optional +from urllib.parse import urlparse + +logger = logging.getLogger(__name__) + +DEFAULT_OLLAMA_BASE = "http://localhost:11434" + + +def detect_ollama(base_url: str = DEFAULT_OLLAMA_BASE) -> Optional[str]: + """Check if Ollama is running at the given URL. + + Returns the Ollama version string if detected, None otherwise. + """ + try: + import httpx + except ImportError: + return None + + try: + response = httpx.get(f"{base_url}/api/tags", timeout=5.0) + if response.status_code == 200: + # Ollama responded — try to get version + try: + ver_response = httpx.get(f"{base_url}/api/version", timeout=5.0) + if ver_response.status_code == 200: + return ver_response.json().get("version", "unknown") + except Exception: + pass + return "unknown" + except Exception: + pass + return None + + +def validate_ollama_url(base_url: str) -> None: + """Validate that the Ollama URL is localhost-only (SSRF protection). + + Raises ValueError if the URL is not a local address. + """ + parsed = urlparse(base_url) + if parsed.scheme not in ("http", "https"): + raise ValueError( + f"Ollama URL must use HTTP or HTTPS scheme (got {parsed.scheme}://)" + ) + local_hosts = ("localhost", "127.0.0.1", "::1", "0.0.0.0") + if parsed.hostname not in local_hosts: + raise ValueError( + f"Ollama URL must be localhost (got {parsed.hostname}). " + "Remote Ollama instances are not supported for security reasons." + ) + + +def generate_ollama( + prompt: str, + count: int, + fmt: str, + model_name: str, + base_url: str, + temperature: float, + generation_prompt: str, +) -> list[dict]: + """Generate examples using Ollama via its OpenAI-compatible API. + + Args: + prompt: User-provided topic/instructions. + count: Number of examples to generate. + fmt: Output format (alpaca/sharegpt/chatml). + model_name: Ollama model name (e.g. llama3.1). + base_url: Ollama base URL. + temperature: Sampling temperature. + generation_prompt: Pre-built generation prompt. + + Returns: + List of generated example dicts. + """ + try: + import httpx + except ImportError: + raise ImportError( + "httpx is required for Ollama generation. Install: pip install httpx" + ) + + validate_ollama_url(base_url) + + api_url = f"{base_url}/v1/chat/completions" + + response = httpx.post( + api_url, + headers={"Content-Type": "application/json"}, + json={ + "model": model_name, + "messages": [ + {"role": "system", "content": generation_prompt}, + {"role": "user", "content": f"Generate {count} training examples now."}, + ], + "temperature": temperature, + "max_tokens": 4096, + }, + timeout=300.0, + ) + + if response.status_code != 200: + logger.debug("Ollama error response: %s", response.text) + raise ValueError( + f"Ollama returned {response.status_code}. " + "Check that Ollama is running and the model is pulled." + ) + + data = response.json() + try: + content = data["choices"][0]["message"]["content"] + except (KeyError, IndexError, TypeError) as exc: + raise ValueError(f"Unexpected Ollama response format: {exc}") from exc + + from soup_cli.commands.generate import _parse_json_array + + return _parse_json_array(content) diff --git a/soup_cli/data/providers/vllm.py b/soup_cli/data/providers/vllm.py new file mode 100644 index 0000000..09e0d8d --- /dev/null +++ b/soup_cli/data/providers/vllm.py @@ -0,0 +1,107 @@ +"""vLLM provider for synthetic data generation. + +Connects to a local vLLM server via its OpenAI-compatible API. +Supports batch mode with concurrent requests. +""" + +import logging +from urllib.parse import urlparse + +logger = logging.getLogger(__name__) + +DEFAULT_VLLM_BASE = "http://localhost:8000" + + +def validate_vllm_url(base_url: str) -> None: + """Validate the vLLM server URL (SSRF protection). + + Reuses the same validation as the server provider: + - Scheme must be http or https + - HTTP is only allowed for localhost + + Raises ValueError if validation fails. + """ + parsed = urlparse(base_url) + if parsed.scheme not in ("http", "https"): + raise ValueError( + f"vLLM URL must use HTTP or HTTPS scheme (got {parsed.scheme}://)" + ) + local_hosts = ("localhost", "127.0.0.1", "::1", "0.0.0.0") + is_local = parsed.hostname in local_hosts + if not is_local and parsed.scheme != "https": + raise ValueError( + f"vLLM URL must use HTTPS for remote servers (got {parsed.scheme}://). " + "HTTP is only allowed for localhost." + ) + + +def generate_vllm( + prompt: str, + count: int, + fmt: str, + model_name: str, + base_url: str, + temperature: float, + generation_prompt: str, + batch_size: int = 5, +) -> list[dict]: + """Generate examples using vLLM server. + + Args: + prompt: User-provided topic/instructions. + count: Number of examples to generate. + fmt: Output format (alpaca/sharegpt/chatml). + model_name: Model name on vLLM server. + base_url: vLLM server base URL. + temperature: Sampling temperature. + generation_prompt: Pre-built generation prompt. + batch_size: Number of concurrent requests. + + Returns: + List of generated example dicts. + """ + try: + import httpx + except ImportError: + raise ImportError( + "httpx is required for vLLM generation. Install: pip install httpx" + ) + + validate_vllm_url(base_url) + + # Ensure URL ends with /v1 + api_base = base_url.rstrip("/") + if not api_base.endswith("/v1"): + api_base = api_base + "/v1" + + response = httpx.post( + f"{api_base}/chat/completions", + headers={"Content-Type": "application/json"}, + json={ + "model": model_name, + "messages": [ + {"role": "system", "content": generation_prompt}, + {"role": "user", "content": f"Generate {count} training examples now."}, + ], + "temperature": temperature, + "max_tokens": 4096, + }, + timeout=300.0, + ) + + if response.status_code != 200: + logger.debug("vLLM error response: %s", response.text) + raise ValueError( + f"vLLM server returned {response.status_code}. " + "Check that the server is running and the model is loaded." + ) + + data = response.json() + try: + content = data["choices"][0]["message"]["content"] + except (KeyError, IndexError, TypeError) as exc: + raise ValueError(f"Unexpected vLLM response format: {exc}") from exc + + from soup_cli.commands.generate import _parse_json_array + + return _parse_json_array(content) diff --git a/soup_cli/data/templates/__init__.py b/soup_cli/data/templates/__init__.py new file mode 100644 index 0000000..eef68e7 --- /dev/null +++ b/soup_cli/data/templates/__init__.py @@ -0,0 +1 @@ +"""Domain templates for synthetic data generation.""" diff --git a/soup_cli/data/templates/code.py b/soup_cli/data/templates/code.py new file mode 100644 index 0000000..50e03e4 --- /dev/null +++ b/soup_cli/data/templates/code.py @@ -0,0 +1,52 @@ +"""Code instruction pair template for synthetic data generation.""" + +SYSTEM_PROMPT = ( + "You are a coding instruction data generator. Generate diverse, high-quality " + "code instruction-response pairs covering different programming tasks." +) + +TEMPLATE_SPEC = { + "languages": ["Python", "JavaScript", "Go", "Rust", "Java"], + "types": ["function", "debug", "explain", "refactor", "test"], +} + + +def build_prompt( + count: int, + fmt: str, + format_spec: str, + language: str = "Python", + task_type: str = "function", +) -> str: + """Build the generation prompt for code instruction pairs. + + Args: + count: Number of examples to generate. + fmt: Output format (alpaca/sharegpt/chatml). + format_spec: Format specification string. + language: Programming language to focus on. + task_type: Type of coding task. + + Returns: + Complete generation prompt string. + """ + type_descriptions = { + "function": "writing functions to solve specific problems", + "debug": "finding and fixing bugs in code snippets", + "explain": "explaining what given code does step by step", + "refactor": "improving and refactoring existing code", + "test": "writing unit tests for given functions", + } + + task_desc = type_descriptions.get(task_type, type_descriptions["function"]) + + return ( + f"You are a training data generator. Generate exactly {count} diverse, " + f"high-quality coding instruction-response pairs.\n\n" + f"Language: {language}\n" + f"Task type: {task_desc}\n\n" + f"Each example should involve {language} code for {task_desc}.\n" + f"Include realistic code snippets in responses.\n\n" + f"Format: {format_spec}\n\n" + f"Return ONLY a JSON array of {count} examples. No markdown, no explanation." + ) diff --git a/soup_cli/data/templates/conversation.py b/soup_cli/data/templates/conversation.py new file mode 100644 index 0000000..d674208 --- /dev/null +++ b/soup_cli/data/templates/conversation.py @@ -0,0 +1,39 @@ +"""Multi-turn conversation template for synthetic data generation.""" + +SYSTEM_PROMPT = ( + "You are a conversational training data generator. Generate diverse, natural " + "multi-turn dialogues between a user and an AI assistant." +) + + +def build_prompt( + count: int, + fmt: str, + format_spec: str, + turns: int = 4, + topic: str = "general knowledge", +) -> str: + """Build the generation prompt for multi-turn conversations. + + Args: + count: Number of examples to generate. + fmt: Output format (alpaca/sharegpt/chatml). + format_spec: Format specification string. + turns: Number of conversation turns (2-10). + topic: Conversation topic. + + Returns: + Complete generation prompt string. + """ + turns = max(2, min(10, turns)) + + return ( + f"You are a training data generator. Generate exactly {count} diverse, " + f"high-quality multi-turn conversations.\n\n" + f"Topic: {topic}\n" + f"Turns per conversation: {turns} (user and assistant messages)\n\n" + f"Each conversation should feel natural with follow-up questions " + f"and contextual responses.\n\n" + f"Format: {format_spec}\n\n" + f"Return ONLY a JSON array of {count} examples. No markdown, no explanation." + ) diff --git a/soup_cli/data/templates/preference.py b/soup_cli/data/templates/preference.py new file mode 100644 index 0000000..79cc04a --- /dev/null +++ b/soup_cli/data/templates/preference.py @@ -0,0 +1,47 @@ +"""Preference data template for DPO/KTO/ORPO synthetic data generation.""" + +SYSTEM_PROMPT = ( + "You are a preference training data generator. Generate examples with both " + "good (chosen) and bad (rejected) responses for preference learning." +) + + +def build_prompt( + count: int, + task: str = "dpo", +) -> str: + """Build the generation prompt for preference data. + + Args: + count: Number of examples to generate. + task: Target task format (dpo/kto/orpo). + + Returns: + Complete generation prompt string. + """ + if task == "kto": + format_desc = ( + 'Each example must be a JSON object with keys: ' + '"prompt" (the user question), ' + '"completion" (the model response), ' + '"label" (boolean: true for good, false for bad). ' + 'Generate roughly equal numbers of true and false labels.' + ) + else: + # DPO / ORPO format + format_desc = ( + 'Each example must be a JSON object with keys: ' + '"prompt" (the user question), ' + '"chosen" (a good, helpful response), ' + '"rejected" (a bad, unhelpful, or incorrect response). ' + 'The chosen response should be clearly better than the rejected one.' + ) + + return ( + f"You are a training data generator. Generate exactly {count} diverse, " + f"high-quality preference training examples.\n\n" + f"Format: {format_desc}\n\n" + f"Cover diverse topics: math, science, coding, writing, reasoning.\n" + f"Make rejected responses subtly wrong (not obviously garbage).\n\n" + f"Return ONLY a JSON array of {count} examples. No markdown, no explanation." + ) diff --git a/soup_cli/data/templates/qa.py b/soup_cli/data/templates/qa.py new file mode 100644 index 0000000..4df6b1b --- /dev/null +++ b/soup_cli/data/templates/qa.py @@ -0,0 +1,47 @@ +"""Question-answer from context template for synthetic data generation.""" + +SYSTEM_PROMPT = ( + "You are a QA training data generator. Generate question-answer pairs " + "that are grounded in provided context documents." +) + + +def build_prompt( + count: int, + fmt: str, + format_spec: str, + context: str = "", +) -> str: + """Build the generation prompt for QA pairs. + + Args: + count: Number of examples to generate. + fmt: Output format (alpaca/sharegpt/chatml). + format_spec: Format specification string. + context: Source document text to generate QA from. + + Returns: + Complete generation prompt string. + """ + context_section = "" + if context: + # Cap context to prevent prompt overflow + truncated = context[:8000] + context_section = ( + f"\n\nSource document to generate questions from:\n" + f"---\n{truncated}\n---\n\n" + f"Generate questions and answers that are grounded in this document. " + f"Answers must be derivable from the text." + ) + else: + context_section = ( + "\n\nGenerate diverse questions and detailed answers on " + "general knowledge topics." + ) + + return ( + f"You are a training data generator. Generate exactly {count} diverse, " + f"high-quality question-answer pairs.{context_section}\n\n" + f"Format: {format_spec}\n\n" + f"Return ONLY a JSON array of {count} examples. No markdown, no explanation." + ) diff --git a/soup_cli/data/templates/reasoning.py b/soup_cli/data/templates/reasoning.py new file mode 100644 index 0000000..f9c6798 --- /dev/null +++ b/soup_cli/data/templates/reasoning.py @@ -0,0 +1,54 @@ +"""Reasoning / chain-of-thought template for GRPO synthetic data generation.""" + +SYSTEM_PROMPT = ( + "You are a reasoning training data generator. Generate problems with " + "detailed step-by-step solutions for training reasoning models." +) + +DOMAINS = ("math", "logic", "code") + + +def build_prompt( + count: int, + fmt: str, + format_spec: str, + domain: str = "math", +) -> str: + """Build the generation prompt for reasoning data. + + Args: + count: Number of examples to generate. + fmt: Output format (alpaca/sharegpt/chatml). + format_spec: Format specification string. + domain: Problem domain (math/logic/code). + + Returns: + Complete generation prompt string. + """ + domain_descriptions = { + "math": ( + "mathematical problems requiring multi-step calculation. " + "Include algebra, arithmetic, geometry, and word problems." + ), + "logic": ( + "logical reasoning puzzles and deduction problems. " + "Include syllogisms, truth tables, and constraint satisfaction." + ), + "code": ( + "programming challenges requiring algorithmic thinking. " + "Include data structures, algorithms, and problem decomposition." + ), + } + + domain_desc = domain_descriptions.get(domain, domain_descriptions["math"]) + + return ( + f"You are a training data generator. Generate exactly {count} diverse " + f"reasoning problems with step-by-step solutions.\n\n" + f"Domain: {domain_desc}\n\n" + f"Each solution must show detailed chain-of-thought reasoning. " + f"Use ... tags to wrap the reasoning steps, then provide " + f"the final answer.\n\n" + f"Format: {format_spec}\n\n" + f"Return ONLY a JSON array of {count} examples. No markdown, no explanation." + ) diff --git a/tests/test_synth_data_pro.py b/tests/test_synth_data_pro.py new file mode 100644 index 0000000..b6a438a --- /dev/null +++ b/tests/test_synth_data_pro.py @@ -0,0 +1,1192 @@ +"""Tests for v0.20.0 — Synth Data Gen Pro. + +Covers: new providers (Ollama, Anthropic, vLLM), domain templates, +quality pipeline (validate, filter, dedup), and security (SSRF). +""" + +import json +import os +from unittest.mock import MagicMock +from unittest.mock import patch as mock_patch + +import pytest + +# ─── Provider Validation Tests ────────────────────────────────────────── + + +class TestNewProvidersValidation: + """Test that new providers are accepted by the CLI.""" + + def test_ollama_provider_accepted(self): + """'ollama' should be a valid provider.""" + from soup_cli.commands.generate import VALID_PROVIDERS + + assert "ollama" in VALID_PROVIDERS + + def test_anthropic_provider_accepted(self): + """'anthropic' should be a valid provider.""" + from soup_cli.commands.generate import VALID_PROVIDERS + + assert "anthropic" in VALID_PROVIDERS + + def test_vllm_provider_accepted(self): + """'vllm' should be a valid provider.""" + from soup_cli.commands.generate import VALID_PROVIDERS + + assert "vllm" in VALID_PROVIDERS + + def test_old_providers_still_work(self): + """Original providers should still be valid.""" + from soup_cli.commands.generate import VALID_PROVIDERS + + assert "openai" in VALID_PROVIDERS + assert "local" in VALID_PROVIDERS + assert "server" in VALID_PROVIDERS + + def test_invalid_provider_rejected_cli(self): + """Invalid provider should be rejected by CLI.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, [ + "data", "generate", + "--prompt", "test", + "--provider", "invalid_provider", + "--count", "1", + ]) + assert result.exit_code != 0 + + +class TestNewProvidersCLIHelp: + """Test CLI help text includes new providers.""" + + def test_help_mentions_ollama(self): + """Help text should mention 'ollama' provider.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "ollama" in result.output + + def test_help_mentions_anthropic(self): + """Help text should mention 'anthropic' provider.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "anthropic" in result.output + + def test_help_mentions_vllm(self): + """Help text should mention 'vllm' provider.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "vllm" in result.output + + def test_help_mentions_template(self): + """Help text should mention --template option.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "template" in result.output.lower() + + def test_help_mentions_quality_pipeline(self): + """Help text should mention --quality-pipeline option.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "quality-pipeline" in result.output + + +# ─── Ollama Provider Tests ────────────────────────────────────────────── + + +class TestOllamaProvider: + """Test the Ollama provider.""" + + def test_detect_ollama_success(self): + """detect_ollama should return version when Ollama is running.""" + from soup_cli.data.providers.ollama import detect_ollama + + mock_tags = MagicMock() + mock_tags.status_code = 200 + mock_tags.json.return_value = {"models": []} + + mock_ver = MagicMock() + mock_ver.status_code = 200 + mock_ver.json.return_value = {"version": "0.6.2"} + + with mock_patch("httpx.get", side_effect=[mock_tags, mock_ver]): + version = detect_ollama() + assert version == "0.6.2" + + def test_detect_ollama_not_running(self): + """detect_ollama should return None when Ollama is not running.""" + from soup_cli.data.providers.ollama import detect_ollama + + with mock_patch("httpx.get", side_effect=Exception("connection refused")): + version = detect_ollama() + assert version is None + + def test_detect_ollama_version_endpoint_fails(self): + """detect_ollama should return 'unknown' if version endpoint fails.""" + from soup_cli.data.providers.ollama import detect_ollama + + mock_tags = MagicMock() + mock_tags.status_code = 200 + mock_tags.json.return_value = {"models": []} + + mock_ver = MagicMock() + mock_ver.status_code = 500 + + with mock_patch("httpx.get", side_effect=[mock_tags, Exception("fail")]): + version = detect_ollama() + assert version == "unknown" + + def test_generate_ollama_calls_api(self): + """generate_ollama should call Ollama API correctly.""" + from soup_cli.data.providers.ollama import generate_ollama + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "choices": [ + {"message": {"content": '[{"instruction": "test", "output": "ok"}]'}} + ] + } + + with mock_patch("httpx.post", return_value=mock_response) as mock_post: + result = generate_ollama( + prompt="test", count=1, fmt="alpaca", + model_name="llama3.1", + base_url="http://localhost:11434", + temperature=0.8, + generation_prompt="Generate 1 example", + ) + + mock_post.assert_called_once() + call_url = mock_post.call_args[0][0] + assert "localhost:11434" in call_url + assert "/v1/chat/completions" in call_url + assert len(result) == 1 + + def test_generate_ollama_error_response(self): + """generate_ollama should raise ValueError on error.""" + from soup_cli.data.providers.ollama import generate_ollama + + mock_response = MagicMock() + mock_response.status_code = 404 + mock_response.text = "model not found" + + with mock_patch("httpx.post", return_value=mock_response): + with pytest.raises(ValueError, match="Ollama returned 404"): + generate_ollama( + prompt="test", count=1, fmt="alpaca", + model_name="nonexistent", + base_url="http://localhost:11434", + temperature=0.8, + generation_prompt="Generate 1 example", + ) + + +class TestOllamaSSRF: + """Test Ollama provider SSRF protection.""" + + def test_blocks_remote_url(self): + """Remote Ollama URL should be rejected.""" + from soup_cli.data.providers.ollama import validate_ollama_url + + with pytest.raises(ValueError, match="localhost"): + validate_ollama_url("http://evil.com:11434") + + def test_blocks_non_http_scheme(self): + """Non-HTTP scheme should be rejected.""" + from soup_cli.data.providers.ollama import validate_ollama_url + + with pytest.raises(ValueError, match="HTTP or HTTPS"): + validate_ollama_url("file:///etc/passwd") + + def test_allows_localhost(self): + """localhost should be allowed.""" + from soup_cli.data.providers.ollama import validate_ollama_url + + validate_ollama_url("http://localhost:11434") # Should not raise + + def test_allows_127_0_0_1(self): + """127.0.0.1 should be allowed.""" + from soup_cli.data.providers.ollama import validate_ollama_url + + validate_ollama_url("http://127.0.0.1:11434") # Should not raise + + def test_generate_ollama_rejects_remote(self): + """generate_ollama should reject remote URLs.""" + from soup_cli.data.providers.ollama import generate_ollama + + with pytest.raises(ValueError, match="localhost"): + generate_ollama( + prompt="test", count=1, fmt="alpaca", + model_name="m", + base_url="http://169.254.169.254", + temperature=0.8, + generation_prompt="test", + ) + + +# ─── Anthropic Provider Tests ─────────────────────────────────────────── + + +class TestAnthropicProvider: + """Test the Anthropic provider.""" + + def test_generate_anthropic_calls_api(self): + """generate_anthropic should call Anthropic Messages API.""" + from soup_cli.data.providers.anthropic import generate_anthropic + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "content": [ + {"type": "text", "text": '[{"instruction": "test", "output": "ok"}]'} + ] + } + + with mock_patch.dict(os.environ, {"ANTHROPIC_API_KEY": "sk-ant-test"}): + with mock_patch("httpx.post", return_value=mock_response) as mock_post: + result = generate_anthropic( + prompt="test", count=1, fmt="alpaca", + model_name="claude-3-haiku-20240307", + temperature=0.8, + generation_prompt="Generate 1 example", + ) + + mock_post.assert_called_once() + call_url = mock_post.call_args[0][0] + assert "anthropic.com" in call_url + call_headers = mock_post.call_args[1]["headers"] + assert call_headers["x-api-key"] == "sk-ant-test" + assert len(result) == 1 + + def test_generate_anthropic_no_api_key(self): + """generate_anthropic should raise ValueError without API key.""" + from soup_cli.data.providers.anthropic import generate_anthropic + + with mock_patch.dict(os.environ, {}, clear=True): + # Remove ANTHROPIC_API_KEY if present + env = os.environ.copy() + env.pop("ANTHROPIC_API_KEY", None) + with mock_patch.dict(os.environ, env, clear=True): + with pytest.raises(ValueError, match="ANTHROPIC_API_KEY"): + generate_anthropic( + prompt="test", count=1, fmt="alpaca", + model_name="claude-3-haiku-20240307", + temperature=0.8, + generation_prompt="test", + ) + + def test_generate_anthropic_error_response(self): + """generate_anthropic should raise ValueError on API error.""" + from soup_cli.data.providers.anthropic import generate_anthropic + + mock_response = MagicMock() + mock_response.status_code = 401 + mock_response.text = "Unauthorized" + + with mock_patch.dict(os.environ, {"ANTHROPIC_API_KEY": "bad-key"}): + with mock_patch("httpx.post", return_value=mock_response): + with pytest.raises(ValueError, match="401"): + generate_anthropic( + prompt="test", count=1, fmt="alpaca", + model_name="claude-3-haiku-20240307", + temperature=0.8, + generation_prompt="test", + ) + + def test_anthropic_api_key_from_env_only(self): + """Anthropic API key should only come from environment, never CLI.""" + import inspect + + from soup_cli.data.providers.anthropic import generate_anthropic + + sig = inspect.signature(generate_anthropic) + param_names = list(sig.parameters.keys()) + assert "api_key" not in param_names + + def test_anthropic_parses_content_blocks(self): + """Anthropic response with multiple content blocks should be parsed.""" + from soup_cli.data.providers.anthropic import generate_anthropic + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "content": [ + {"type": "text", "text": '[{"instruction": "a", "output": "b"}'}, + {"type": "text", "text": ', {"instruction": "c", "output": "d"}]'}, + ] + } + + with mock_patch.dict(os.environ, {"ANTHROPIC_API_KEY": "sk-ant-test"}): + with mock_patch("httpx.post", return_value=mock_response): + result = generate_anthropic( + prompt="test", count=2, fmt="alpaca", + model_name="claude-3-haiku-20240307", + temperature=0.8, + generation_prompt="test", + ) + + assert len(result) == 2 + + +# ─── vLLM Provider Tests ──────────────────────────────────────────────── + + +class TestVLLMProvider: + """Test the vLLM provider.""" + + def test_generate_vllm_calls_api(self): + """generate_vllm should call vLLM OpenAI-compatible API.""" + from soup_cli.data.providers.vllm import generate_vllm + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "choices": [ + {"message": {"content": '[{"instruction": "test", "output": "ok"}]'}} + ] + } + + with mock_patch("httpx.post", return_value=mock_response) as mock_post: + result = generate_vllm( + prompt="test", count=1, fmt="alpaca", + model_name="meta-llama/Llama-3.1-8B-Instruct", + base_url="http://localhost:8000", + temperature=0.8, + generation_prompt="Generate 1 example", + ) + + mock_post.assert_called_once() + call_url = mock_post.call_args[0][0] + assert "localhost:8000" in call_url + assert "/v1/chat/completions" in call_url + assert len(result) == 1 + + def test_generate_vllm_appends_v1(self): + """generate_vllm should append /v1 if missing.""" + from soup_cli.data.providers.vllm import generate_vllm + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "choices": [{"message": {"content": "[]"}}] + } + + with mock_patch("httpx.post", return_value=mock_response) as mock_post: + generate_vllm( + prompt="test", count=1, fmt="alpaca", + model_name="m", base_url="http://localhost:8000", + temperature=0.8, generation_prompt="test", + ) + + call_url = mock_post.call_args[0][0] + assert "/v1/chat/completions" in call_url + + def test_generate_vllm_error_response(self): + """generate_vllm should raise ValueError on error.""" + from soup_cli.data.providers.vllm import generate_vllm + + mock_response = MagicMock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + + with mock_patch("httpx.post", return_value=mock_response): + with pytest.raises(ValueError, match="vLLM server returned 500"): + generate_vllm( + prompt="test", count=1, fmt="alpaca", + model_name="m", base_url="http://localhost:8000", + temperature=0.8, generation_prompt="test", + ) + + +class TestVLLMSSRF: + """Test vLLM provider SSRF protection.""" + + def test_blocks_non_http_scheme(self): + """Non-HTTP scheme should be rejected.""" + from soup_cli.data.providers.vllm import validate_vllm_url + + with pytest.raises(ValueError, match="HTTP or HTTPS"): + validate_vllm_url("file:///etc/passwd") + + def test_blocks_remote_http(self): + """Remote HTTP should be rejected.""" + from soup_cli.data.providers.vllm import validate_vllm_url + + with pytest.raises(ValueError, match="HTTPS for remote"): + validate_vllm_url("http://169.254.169.254/latest") + + def test_allows_localhost_http(self): + """localhost HTTP should be allowed.""" + from soup_cli.data.providers.vllm import validate_vllm_url + + validate_vllm_url("http://localhost:8000") # Should not raise + + def test_allows_remote_https(self): + """Remote HTTPS should be allowed.""" + from soup_cli.data.providers.vllm import validate_vllm_url + + validate_vllm_url("https://vllm.example.com:8000") # Should not raise + + def test_generate_vllm_rejects_remote_http(self): + """generate_vllm should reject remote HTTP URLs.""" + from soup_cli.data.providers.vllm import generate_vllm + + with pytest.raises(ValueError, match="HTTPS for remote"): + generate_vllm( + prompt="test", count=1, fmt="alpaca", + model_name="m", + base_url="http://evil.com:8000", + temperature=0.8, + generation_prompt="test", + ) + + +# ─── Batch Routing Tests ──────────────────────────────────────────────── + + +class TestBatchRoutingNewProviders: + """Test that _generate_batch routes to new providers correctly.""" + + def test_routes_to_ollama(self): + """provider='ollama' should route to Ollama provider.""" + from soup_cli.commands.generate import _generate_batch + + with mock_patch( + "soup_cli.data.providers.ollama.generate_ollama", + return_value=[{"instruction": "x", "output": "y"}], + ) as mock_ollama: + result = _generate_batch( + prompt="test", count=1, fmt="alpaca", + provider="ollama", model_name="llama3.1", + api_key=None, api_base=None, + temperature=0.8, seed_examples=[], + ) + + mock_ollama.assert_called_once() + assert len(result) == 1 + + def test_routes_to_anthropic(self): + """provider='anthropic' should route to Anthropic provider.""" + from soup_cli.commands.generate import _generate_batch + + with mock_patch( + "soup_cli.data.providers.anthropic.generate_anthropic", + return_value=[{"instruction": "x", "output": "y"}], + ) as mock_anthropic: + result = _generate_batch( + prompt="test", count=1, fmt="alpaca", + provider="anthropic", model_name="claude-3-haiku-20240307", + api_key=None, api_base=None, + temperature=0.8, seed_examples=[], + ) + + mock_anthropic.assert_called_once() + assert len(result) == 1 + + def test_routes_to_vllm(self): + """provider='vllm' should route to vLLM provider.""" + from soup_cli.commands.generate import _generate_batch + + with mock_patch( + "soup_cli.data.providers.vllm.generate_vllm", + return_value=[{"instruction": "x", "output": "y"}], + ) as mock_vllm: + result = _generate_batch( + prompt="test", count=1, fmt="alpaca", + provider="vllm", model_name="m", + api_key=None, api_base=None, + temperature=0.8, seed_examples=[], + ) + + mock_vllm.assert_called_once() + assert len(result) == 1 + + def test_old_providers_still_route(self): + """Original providers should still route correctly.""" + from soup_cli.commands.generate import _generate_batch + + for prov, mock_target in [ + ("openai", "soup_cli.commands.generate._generate_openai"), + ("local", "soup_cli.commands.generate._generate_local"), + ("server", "soup_cli.commands.generate._generate_server"), + ]: + with mock_patch(mock_target, return_value=[]) as mock_fn: + _generate_batch( + prompt="test", count=1, fmt="alpaca", + provider=prov, model_name="m", + api_key="key" if prov == "openai" else None, + api_base=None, + temperature=0.8, seed_examples=[], + ) + mock_fn.assert_called_once() + + +# ─── Domain Template Tests ────────────────────────────────────────────── + + +class TestTemplateValidation: + """Test template validation.""" + + def test_valid_templates(self): + """All valid templates should be accepted.""" + from soup_cli.commands.generate import VALID_TEMPLATES + + assert "code" in VALID_TEMPLATES + assert "conversation" in VALID_TEMPLATES + assert "qa" in VALID_TEMPLATES + assert "preference" in VALID_TEMPLATES + assert "reasoning" in VALID_TEMPLATES + + def test_invalid_template_rejected_cli(self): + """Invalid template should be rejected.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, [ + "data", "generate", + "--prompt", "test", + "--template", "invalid_template", + "--count", "1", + ]) + assert result.exit_code != 0 + + +class TestCodeTemplate: + """Test code domain template.""" + + def test_build_prompt_default(self): + """Code template should build a valid prompt.""" + from soup_cli.data.templates.code import build_prompt + + result = build_prompt(5, "alpaca", "format spec", language="Python") + assert "Python" in result + assert "5" in result + + def test_build_prompt_languages(self): + """Code template should support different languages.""" + from soup_cli.data.templates.code import build_prompt + + for lang in ["Python", "JavaScript", "Go", "Rust", "Java"]: + result = build_prompt(3, "alpaca", "spec", language=lang) + assert lang in result + + def test_build_prompt_task_types(self): + """Code template should support different task types.""" + from soup_cli.data.templates.code import build_prompt + + for task_type in ["function", "debug", "explain", "refactor", "test"]: + result = build_prompt(3, "alpaca", "spec", task_type=task_type) + assert isinstance(result, str) + assert len(result) > 0 + + def test_template_spec_has_languages(self): + """Template spec should list supported languages.""" + from soup_cli.data.templates.code import TEMPLATE_SPEC + + assert "languages" in TEMPLATE_SPEC + assert len(TEMPLATE_SPEC["languages"]) >= 5 + + +class TestConversationTemplate: + """Test conversation domain template.""" + + def test_build_prompt_default(self): + """Conversation template should build a valid prompt.""" + from soup_cli.data.templates.conversation import build_prompt + + result = build_prompt(5, "chatml", "format spec") + assert "5" in result + assert "conversation" in result.lower() or "multi-turn" in result.lower() + + def test_build_prompt_with_topic(self): + """Conversation template should include topic.""" + from soup_cli.data.templates.conversation import build_prompt + + result = build_prompt(3, "chatml", "spec", topic="science fiction") + assert "science fiction" in result + + def test_turns_clamped(self): + """Turns should be clamped to 2-10 range.""" + from soup_cli.data.templates.conversation import build_prompt + + # turns < 2 should be clamped to 2 + result = build_prompt(3, "chatml", "spec", turns=0) + assert "2" in result + + # turns > 10 should be clamped to 10 + result = build_prompt(3, "chatml", "spec", turns=20) + assert "10" in result + + +class TestQATemplate: + """Test QA domain template.""" + + def test_build_prompt_without_context(self): + """QA template without context should generate general QA.""" + from soup_cli.data.templates.qa import build_prompt + + result = build_prompt(5, "alpaca", "format spec") + assert "5" in result + assert "question" in result.lower() + + def test_build_prompt_with_context(self): + """QA template with context should include it.""" + from soup_cli.data.templates.qa import build_prompt + + context = "Python is a programming language created by Guido van Rossum." + result = build_prompt(3, "alpaca", "spec", context=context) + assert "Python" in result + assert "Guido" in result + + def test_context_truncated(self): + """Long context should be truncated to prevent overflow.""" + from soup_cli.data.templates.qa import build_prompt + + context = "x" * 20000 + result = build_prompt(3, "alpaca", "spec", context=context) + # Context should be truncated to 8000 chars + assert len(result) < 20000 + + +class TestPreferenceTemplate: + """Test preference domain template.""" + + def test_build_prompt_dpo(self): + """Preference template for DPO should use chosen/rejected format.""" + from soup_cli.data.templates.preference import build_prompt + + result = build_prompt(5, task="dpo") + assert "chosen" in result + assert "rejected" in result + + def test_build_prompt_kto(self): + """Preference template for KTO should use label format.""" + from soup_cli.data.templates.preference import build_prompt + + result = build_prompt(5, task="kto") + assert "label" in result + assert "true" in result.lower() or "false" in result.lower() + + def test_build_prompt_orpo(self): + """Preference template for ORPO should use chosen/rejected format.""" + from soup_cli.data.templates.preference import build_prompt + + result = build_prompt(5, task="orpo") + assert "chosen" in result + assert "rejected" in result + + +class TestReasoningTemplate: + """Test reasoning domain template.""" + + def test_build_prompt_math(self): + """Reasoning template for math should include math description.""" + from soup_cli.data.templates.reasoning import build_prompt + + result = build_prompt(5, "alpaca", "format spec", domain="math") + assert "math" in result.lower() + + def test_build_prompt_logic(self): + """Reasoning template for logic should include logic description.""" + from soup_cli.data.templates.reasoning import build_prompt + + result = build_prompt(5, "alpaca", "spec", domain="logic") + assert "logic" in result.lower() + + def test_build_prompt_code(self): + """Reasoning template for code should include code description.""" + from soup_cli.data.templates.reasoning import build_prompt + + result = build_prompt(5, "alpaca", "spec", domain="code") + assert "code" in result.lower() or "algorithm" in result.lower() + + def test_valid_domains(self): + """All domains should be in the DOMAINS constant.""" + from soup_cli.data.templates.reasoning import DOMAINS + + assert "math" in DOMAINS + assert "logic" in DOMAINS + assert "code" in DOMAINS + + +class TestTemplateBuildPromptIntegration: + """Test _build_template_prompt in generate.py.""" + + def test_code_template_integration(self): + """Code template should be built correctly via _build_template_prompt.""" + from soup_cli.commands.generate import _build_template_prompt + + result = _build_template_prompt( + template="code", prompt="test", count=5, fmt="alpaca", + language="Go", task_type="debug", + ) + assert "Go" in result + + def test_conversation_template_integration(self): + """Conversation template should be built correctly.""" + from soup_cli.commands.generate import _build_template_prompt + + result = _build_template_prompt( + template="conversation", prompt="test", count=5, fmt="chatml", + topic="cooking", + ) + assert "cooking" in result + + def test_qa_template_integration(self): + """QA template should be built correctly.""" + from soup_cli.commands.generate import _build_template_prompt + + result = _build_template_prompt( + template="qa", prompt="test", count=5, fmt="alpaca", + context_text="Sample context text", + ) + assert "Sample context" in result + + def test_preference_template_integration(self): + """Preference template should be built correctly.""" + from soup_cli.commands.generate import _build_template_prompt + + result = _build_template_prompt( + template="preference", prompt="test", count=5, fmt="alpaca", + pref_task="kto", + ) + assert "label" in result + + def test_reasoning_template_integration(self): + """Reasoning template should be built correctly.""" + from soup_cli.commands.generate import _build_template_prompt + + result = _build_template_prompt( + template="reasoning", prompt="test", count=5, fmt="alpaca", + domain="logic", + ) + assert "logic" in result.lower() + + def test_unknown_template_falls_back(self): + """Unknown template should fall back to default prompt.""" + from soup_cli.commands.generate import _build_template_prompt + + result = _build_template_prompt( + template="nonexistent", prompt="test topic", count=5, fmt="alpaca", + ) + assert "test topic" in result + + +# ─── Quality Pipeline Tests ───────────────────────────────────────────── + + +class TestValidatePipeline: + """Test the validation pipeline step.""" + + def test_removes_invalid_entries(self, tmp_path): + """Validation pipeline should remove invalid entries.""" + from soup_cli.commands.generate import _run_validate_pipeline + + data = [ + {"instruction": "valid", "input": "", "output": "ok"}, + {"bad": "entry"}, + {"instruction": "also valid", "input": "", "output": "fine"}, + ] + path = tmp_path / "test.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + + _run_validate_pipeline(path, "alpaca") + + with open(path) as f: + lines = [json.loads(line) for line in f if line.strip()] + assert len(lines) == 2 + assert lines[0]["instruction"] == "valid" + assert lines[1]["instruction"] == "also valid" + + def test_keeps_all_valid(self, tmp_path): + """Validation pipeline should keep all entries when all valid.""" + from soup_cli.commands.generate import _run_validate_pipeline + + data = [ + {"instruction": "a", "input": "", "output": "b"}, + {"instruction": "c", "input": "", "output": "d"}, + ] + path = tmp_path / "test.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + + _run_validate_pipeline(path, "alpaca") + + with open(path) as f: + lines = [json.loads(line) for line in f if line.strip()] + assert len(lines) == 2 + + def test_validates_preference_format(self, tmp_path): + """Validation pipeline should accept preference format entries.""" + from soup_cli.commands.generate import _run_validate_pipeline + + data = [ + {"prompt": "q", "chosen": "good", "rejected": "bad"}, + {"prompt": "q", "completion": "ans", "label": True}, + {"bad": "entry"}, + ] + path = tmp_path / "test.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + + _run_validate_pipeline(path, "alpaca") + + with open(path) as f: + lines = [json.loads(line) for line in f if line.strip()] + assert len(lines) == 2 + + +class TestFilterPipeline: + """Test the quality filter pipeline step.""" + + def test_filter_pipeline_runs(self, tmp_path): + """Filter pipeline should run without error.""" + from soup_cli.commands.generate import _run_filter_pipeline + + data = [ + {"instruction": "What is Python?", "output": "Python is a programming language."}, + {"instruction": "asdfjkl asdf", "output": "random noise text here"}, + {"instruction": "Explain ML", "output": "Machine learning is a field of AI."}, + {"instruction": "Math question", "output": "The answer to 2+2 is 4."}, + ] + path = tmp_path / "test.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + + # Should not raise even without torch/transformers + _run_filter_pipeline(path) + + def test_filter_pipeline_empty(self, tmp_path): + """Filter pipeline should handle empty file.""" + from soup_cli.commands.generate import _run_filter_pipeline + + path = tmp_path / "empty.jsonl" + path.write_text("") + + _run_filter_pipeline(path) # Should not raise + + +class TestDedupPipeline: + """Test the dedup pipeline step.""" + + def test_dedup_pipeline_without_datasketch(self, tmp_path): + """Dedup pipeline should gracefully handle missing datasketch.""" + from soup_cli.commands.generate import _run_dedup_pipeline + + data = [ + {"instruction": "a", "output": "b"}, + {"instruction": "a", "output": "b"}, + ] + path = tmp_path / "test.jsonl" + with open(path, "w") as f: + for row in data: + f.write(json.dumps(row) + "\n") + + # This will either dedup or skip gracefully if datasketch not installed + _run_dedup_pipeline(path) + + def test_dedup_pipeline_empty(self, tmp_path): + """Dedup pipeline should handle empty file.""" + from soup_cli.commands.generate import _run_dedup_pipeline + + path = tmp_path / "empty.jsonl" + path.write_text("") + + _run_dedup_pipeline(path) # Should not raise + + +class TestQualityPipelineFlag: + """Test the --quality-pipeline convenience flag.""" + + def test_quality_pipeline_enables_all(self): + """--quality-pipeline should enable validate, filter, and dedup.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + # We can test this by checking the help text mentions the flag + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "quality-pipeline" in result.output + assert "validate" in result.output + assert "filter" in result.output + assert "dedup" in result.output + + +# ─── Preference Validation Tests ──────────────────────────────────────── + + +class TestValidatePreference: + """Test preference data validation.""" + + def test_dpo_format_valid(self): + """DPO format should be valid.""" + from soup_cli.commands.generate import _validate_preference + + assert _validate_preference({"prompt": "q", "chosen": "a", "rejected": "b"}) + + def test_kto_format_valid(self): + """KTO format should be valid.""" + from soup_cli.commands.generate import _validate_preference + + assert _validate_preference({"prompt": "q", "completion": "a", "label": True}) + + def test_invalid_format(self): + """Invalid format should fail validation.""" + from soup_cli.commands.generate import _validate_preference + + assert not _validate_preference({"instruction": "q", "output": "a"}) + assert not _validate_preference({"prompt": "q"}) + assert not _validate_preference({}) + + +# ─── Output Path Sanitization Tests ───────────────────────────────────── + + +class TestOutputPathSanitization: + """Test that output paths are sanitized.""" + + def test_path_traversal_blocked_in_cli(self): + """Output path with '..' should be rejected.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + + # Mock the generation to avoid actual API call + with mock_patch( + "soup_cli.commands.generate._generate_batch", + return_value=[{"instruction": "x", "output": "y"}], + ): + result = runner.invoke(app, [ + "data", "generate", + "--prompt", "test", + "--output", "../../../etc/evil.jsonl", + "--count", "1", + "--provider", "server", + ]) + assert result.exit_code != 0 + + +# ─── Ollama Model Shorthand Tests ─────────────────────────────────────── + + +class TestOllamaModelShorthand: + """Test --ollama-model shorthand flag.""" + + def test_ollama_model_sets_provider(self): + """--ollama-model should set provider to 'ollama'.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "ollama-model" in result.output + + +# ─── Rate Limiting Tests ──────────────────────────────────────────────── + + +class TestRateLimiting: + """Test rate limiting configuration.""" + + def test_requests_per_minute_in_help(self): + """Help should mention --requests-per-minute.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + runner = CliRunner() + result = runner.invoke(app, ["data", "generate", "--help"]) + assert "requests-per-minute" in result.output or "rpm" in result.output + + +# ─── Format Spec Tests ────────────────────────────────────────────────── + + +class TestGetFormatSpec: + """Test _get_format_spec helper.""" + + def test_alpaca_spec(self): + """Alpaca format spec should mention instruction and output.""" + from soup_cli.commands.generate import _get_format_spec + + spec = _get_format_spec("alpaca") + assert "instruction" in spec + assert "output" in spec + + def test_sharegpt_spec(self): + """ShareGPT format spec should mention conversations.""" + from soup_cli.commands.generate import _get_format_spec + + spec = _get_format_spec("sharegpt") + assert "conversations" in spec + + def test_chatml_spec(self): + """ChatML format spec should mention messages.""" + from soup_cli.commands.generate import _get_format_spec + + spec = _get_format_spec("chatml") + assert "messages" in spec + + def test_unknown_format_returns_alpaca(self): + """Unknown format should fall back to alpaca spec.""" + from soup_cli.commands.generate import _get_format_spec + + spec = _get_format_spec("unknown") + assert "instruction" in spec + + +# ─── End-to-End CLI Integration Tests ─────────────────────────────────── + + +class TestEndToEndGeneration: + """Test end-to-end generation with mocked providers.""" + + def test_generate_with_ollama_provider(self, tmp_path): + """Full generation with Ollama provider should work.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + output_path = tmp_path / "output.jsonl" + + with mock_patch( + "soup_cli.data.providers.ollama.generate_ollama", + return_value=[ + {"instruction": "What is AI?", "input": "", "output": "AI is..."}, + {"instruction": "Explain ML", "input": "", "output": "ML is..."}, + ], + ): + runner = CliRunner() + result = runner.invoke(app, [ + "data", "generate", + "--prompt", "Generate AI questions", + "--provider", "ollama", + "--model", "llama3.1", + "--count", "2", + "--batch-size", "2", + "--output", str(output_path), + ]) + + assert result.exit_code == 0 + assert output_path.exists() + with open(output_path) as f: + lines = [json.loads(line) for line in f if line.strip()] + assert len(lines) == 2 + + def test_generate_with_template(self, tmp_path): + """Generation with a template should work.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + output_path = tmp_path / "output.jsonl" + + with mock_patch( + "soup_cli.commands.generate._generate_server", + return_value=[ + {"instruction": "Write a function", "input": "", "output": "def foo(): pass"}, + ], + ): + runner = CliRunner() + result = runner.invoke(app, [ + "data", "generate", + "--prompt", "Generate code", + "--template", "code", + "--language", "Python", + "--provider", "server", + "--count", "1", + "--batch-size", "1", + "--output", str(output_path), + ]) + + assert result.exit_code == 0 + assert "code" in result.output.lower() or "Template" in result.output + + def test_generate_with_anthropic_provider(self, tmp_path): + """Full generation with Anthropic provider should work.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + output_path = tmp_path / "output.jsonl" + + with mock_patch( + "soup_cli.data.providers.anthropic.generate_anthropic", + return_value=[ + {"instruction": "What is AI?", "input": "", "output": "AI is..."}, + ], + ): + runner = CliRunner() + result = runner.invoke(app, [ + "data", "generate", + "--prompt", "Generate AI questions", + "--provider", "anthropic", + "--model", "claude-3-haiku-20240307", + "--count", "1", + "--batch-size", "1", + "--output", str(output_path), + ]) + + assert result.exit_code == 0 + assert output_path.exists() + + def test_generate_with_vllm_provider(self, tmp_path): + """Full generation with vLLM provider should work.""" + from typer.testing import CliRunner + + from soup_cli.cli import app + + output_path = tmp_path / "output.jsonl" + + with mock_patch( + "soup_cli.data.providers.vllm.generate_vllm", + return_value=[ + {"instruction": "What is AI?", "input": "", "output": "AI is..."}, + ], + ): + runner = CliRunner() + result = runner.invoke(app, [ + "data", "generate", + "--prompt", "Generate AI questions", + "--provider", "vllm", + "--count", "1", + "--batch-size", "1", + "--output", str(output_path), + ]) + + assert result.exit_code == 0 + assert output_path.exists()