lmstudiofixed/tools/lmstudio/wrapper_generator.py

#!/usr/bin/env python3
import argparse
import json
from datetime import datetime
from pathlib import Path


DEFAULT_OFFICIAL_MODEL_YAML = Path(
    r"G:\LMStudio\.lmstudio\hub\models\qwen\qwen3.6-35b-a3b\model.yaml"
)
DEFAULT_OFFICIAL_MANIFEST = Path(
    r"G:\LMStudio\.lmstudio\hub\models\qwen\qwen3.6-35b-a3b\manifest.json"
)
DEFAULT_OUTPUT_ROOT = Path("artifacts/lmstudio-model-wrappers")
DEFAULT_TARGET_MODEL_KEY = "qwen3.6-35b-a3b-uncensored-hauhaucs-aggressive"
DEFAULT_WRAPPER_OWNER = "local"
DEFAULT_WRAPPER_NAME = "qwen3.6-35b-a3b-hauhaucs-aggressive"


def build_wrapper_model_yaml(
    official_model_yaml_text: str,
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
) -> str:
    lines = official_model_yaml_text.splitlines()
    output = []
    replaced_model = False
    inserted_base = False
    skipped_original_base = False

    for line in lines:
        stripped = line.strip()

        if not replaced_model and stripped.startswith("model: "):
            output.append(f"model: {wrapper_owner}/{wrapper_name}")
            replaced_model = True
            continue

        if not inserted_base and stripped.startswith("base:"):
            output.append(f'base: "{target_model_key}"')
            inserted_base = True
            skipped_original_base = True
            continue

        if skipped_original_base:
            if line.startswith("metadataOverrides:"):
                skipped_original_base = False
                output.append(line)
            continue

        output.append(line)

    if not replaced_model:
        raise ValueError("Official model.yaml is missing a model declaration")
    if not inserted_base:
        raise ValueError("Official model.yaml is missing a base declaration")

    return "\n".join(output).rstrip() + "\n"


def build_wrapper_manifest(
    official_manifest: dict,
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
) -> dict:
    manifest = {
        "type": "model",
        "owner": wrapper_owner,
        "name": wrapper_name,
        "dependencies": [
            {
                "type": "model",
                "purpose": "baseModel",
                "modelKeys": [target_model_key],
            }
        ],
        "revision": int(official_manifest.get("revision", 1)),
    }

    source_dep = (official_manifest.get("dependencies") or [{}])[0]
    if source_dep.get("sources"):
        manifest["dependencies"][0]["sources"] = source_dep["sources"]
    return manifest


def build_readme(
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
) -> str:
    return (
        f"# {wrapper_owner}/{wrapper_name}\n\n"
        "This wrapper exposes LM Studio reasoning metadata for a local HauhauCS GGUF.\n\n"
        "## Source Model\n\n"
        f"- Local model key: `{target_model_key}`\n\n"
        "## Intent\n\n"
        "- Preserve the original HauhauCS weights unchanged\n"
        "- Reuse the official Qwen3.6 reasoning-capable model wrapper behavior\n"
        "- Keep the wrapper portable and separate from LM Studio `.internal` state\n\n"
        "## Next Steps\n\n"
        "- Place or link this wrapper into an LM Studio-recognized local model package location\n"
        "- Reload LM Studio model metadata\n"
        "- Validate the wrapper via UI reasoning toggle and `/api/v1/chat`\n"
    )


def build_wrapper_files(
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
    official_model_yaml_text: str,
    official_manifest: dict,
) -> dict:
    model_yaml = build_wrapper_model_yaml(
        official_model_yaml_text=official_model_yaml_text,
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
    )
    manifest = build_wrapper_manifest(
        official_manifest=official_manifest,
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
    )
    readme = build_readme(
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
    )
    return {
        "model.yaml": model_yaml,
        "manifest.json": json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
        "README.md": readme,
    }


def generate_wrapper(
    output_dir: Path,
    target_model_key: str,
    wrapper_owner: str,
    wrapper_name: str,
    official_model_yaml_text: str,
    official_manifest: dict,
    dry_run: bool,
) -> dict:
    files = build_wrapper_files(
        target_model_key=target_model_key,
        wrapper_owner=wrapper_owner,
        wrapper_name=wrapper_name,
        official_model_yaml_text=official_model_yaml_text,
        official_manifest=official_manifest,
    )

    summary = {
        "generated_at": datetime.now().isoformat(),
        "dry_run": dry_run,
        "output_dir": str(output_dir),
        "wrapper_model": f"{wrapper_owner}/{wrapper_name}",
        "target_model_key": target_model_key,
        "files": sorted(files.keys()),
    }

    if dry_run:
        return summary

    output_dir.mkdir(parents=True, exist_ok=True)
    for name, content in files.items():
        (output_dir / name).write_text(content, encoding="utf-8")
    (output_dir / "summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    return summary


def load_inputs(model_yaml_path: Path, manifest_path: Path) -> tuple[str, dict]:
    return (
        model_yaml_path.read_text(encoding="utf-8"),
        json.loads(manifest_path.read_text(encoding="utf-8")),
    )


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--target-model-key", default=DEFAULT_TARGET_MODEL_KEY)
    parser.add_argument("--wrapper-owner", default=DEFAULT_WRAPPER_OWNER)
    parser.add_argument("--wrapper-name", default=DEFAULT_WRAPPER_NAME)
    parser.add_argument("--official-model-yaml", default=str(DEFAULT_OFFICIAL_MODEL_YAML))
    parser.add_argument("--official-manifest", default=str(DEFAULT_OFFICIAL_MANIFEST))
    parser.add_argument("--output-dir", default="")
    parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT))
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    official_model_yaml_text, official_manifest = load_inputs(
        Path(args.official_model_yaml),
        Path(args.official_manifest),
    )

    if args.output_dir:
        output_dir = Path(args.output_dir)
    else:
        output_dir = Path(args.output_root) / args.wrapper_owner / args.wrapper_name

    summary = generate_wrapper(
        output_dir=output_dir,
        target_model_key=args.target_model_key,
        wrapper_owner=args.wrapper_owner,
        wrapper_name=args.wrapper_name,
        official_model_yaml_text=official_model_yaml_text,
        official_manifest=official_manifest,
        dry_run=args.dry_run,
    )
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())