Spaces:

kneiff
/

kneifftools

Runtime error

File size: 24,530 Bytes

from __future__ import annotations

import json
from pathlib import Path

import pytest
from typer.testing import CliRunner

import kneiff.cli.llm as cli_llm
from kneiff.config import LmStudioConfig, PromptgenConfig
from kneiff.infer.lmstudio import (
    PromptGenerationModels,
    PromptGenerationRequest,
    PromptGenerationResult,
    PromptPassResult,
    PromptReviewIssue,
)
from kneiff.infer.lmstudio.schema import PromptFieldCatalog
from kneiff.project_scaffold import ensure_project_resources
from tests._cli_helpers import invoke_cli

pytestmark = pytest.mark.usefixtures("isolated_cli_project")


@pytest.fixture(autouse=True)
def _isolate_llm_cli_env(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
) -> None:
    assert isolated_cli_storage.is_dir()
    monkeypatch.delenv("KNF_PROMPTGEN_EXCHANGE_DIR", raising=False)
    monkeypatch.delenv("KNF_PROMPTGEN_SAVE_EXCHANGES", raising=False)


def _reviewed_prompt_result() -> PromptGenerationResult:
    draft = PromptPassResult(
        fields={"view": "rear_view"},
        positive_prompt="draft prompt",
        negative_prompt="draft negative",
    )
    final = PromptPassResult(
        fields={"view": "over_shoulder_view"},
        positive_prompt="final prompt",
        negative_prompt="final negative",
        assumptions=("kept the requested forest setting",),
        missing_input=("lighting was not specified",),
        review_issues=(
            PromptReviewIssue(
                field="view",
                problem="rear view mentioned a front-facing expression",
                fix="changed the view to over-the-shoulder",
            ),
        ),
    )
    return PromptGenerationResult(
        draft=draft,
        final=final,
        fields=final.fields,
        positive_prompt=final.positive_prompt,
        negative_prompt=final.negative_prompt,
        assumptions=final.assumptions,
        missing_input=final.missing_input,
        review_issues=final.review_issues,
        models=PromptGenerationModels(draft="draft-model", review="review-model"),
    )


def test_llm_prompt_cli_prints_review_output_by_default(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    requests: list[PromptGenerationRequest] = []
    configs: list[tuple[LmStudioConfig, PromptgenConfig]] = []

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        *,
        lmstudio_config: LmStudioConfig,
        promptgen_config: PromptgenConfig,
        field_catalog: PromptFieldCatalog,
    ) -> PromptGenerationResult:
        requests.append(request)
        configs.append((lmstudio_config, promptgen_config))
        assert "subject" in field_catalog.fields
        final = PromptPassResult(
            fields={"subject": "Character_Token", "composition": "front_view"},
            positive_prompt="final prompt",
            negative_prompt="bad hands",
            assumptions=("kept requested pose",),
        )
        return PromptGenerationResult(
            draft=final,
            final=final,
            fields=final.fields,
            positive_prompt=final.positive_prompt,
            negative_prompt=final.negative_prompt,
            assumptions=final.assumptions,
            missing_input=final.missing_input,
            review_issues=final.review_issues,
            models=PromptGenerationModels(draft="draft-model", review="review-model"),
        )

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "front view portrait",
            "--model",
            "shared-model",
            "--draft-model",
            "draft-model",
            "--review-model",
            "review-model",
            "--base-url",
            "http://lm.example/v1",
            "--field",
            "subject=Character_Token",
            "--field",
            "pose_and_behavior=standing,leaning",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert "First version:\nfinal prompt" in result.output
    assert "Second version:\nfinal prompt" in result.output
    assert "Comments:\n- Assumption: kept requested pose" in result.output
    assert '"models"' not in result.output
    assert requests[0].user_input == "front view portrait"
    assert requests[0].model == "shared-model"
    assert requests[0].draft_model == "draft-model"
    assert requests[0].review_model == "review-model"
    assert requests[0].base_url == "http://lm.example/v1"
    assert requests[0].field_overrides["subject"] == "Character_Token"
    assert requests[0].field_overrides["pose_and_behavior"] == (
        "standing",
        "leaning",
    )
    assert requests[0].save_exchanges is None
    assert requests[0].exchange_dir is None
    assert isinstance(configs[0][0], LmStudioConfig)
    assert isinstance(configs[0][1], PromptgenConfig)


def test_llm_prompt_cli_isolates_storage_env_between_invocations(
    monkeypatch: pytest.MonkeyPatch,
    tmp_path: Path,
) -> None:
    models: list[str] = []
    storage_roots = []
    for name in ("alpha", "beta"):
        root = tmp_path / name
        root.mkdir()
        (root / ".env.apprc-storage").write_text(
            f"KNF_LMSTUDIO_MODEL={name}-model\n",
            encoding="utf-8",
        )
        ensure_project_resources(
            root,
            activation_token=f"{name}_Character",
            species_token=f"{name}_species",
        )
        storage_roots.append(root)

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        *,
        lmstudio_config: LmStudioConfig,
        promptgen_config: PromptgenConfig,
        field_catalog: PromptFieldCatalog,
    ) -> PromptGenerationResult:
        models.append(lmstudio_config.model)
        assert "subject" in field_catalog.fields
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)
    runner = CliRunner()

    for storage_root in storage_roots:
        monkeypatch.setenv("KNF_STORAGE", str(storage_root))
        result = invoke_cli(
            runner,
            ["llm", "prompt", "stored scene", "--no-export"],
        )
        assert result.exit_code == 0, result.output

    assert models == ["alpha-model", "beta-model"]


def test_llm_prompt_cli_verbose_prints_structured_json(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "front view portrait",
            "--verbose",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    payload = json.loads(result.output)
    assert payload["positive_prompt"] == "final prompt"
    assert payload["models"]["draft"] == "draft-model"


def test_llm_prompt_cli_format_json_prints_structured_json(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "front view portrait",
            "--format",
            "json",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert json.loads(result.output)["positive_prompt"] == "final prompt"


def test_llm_prompt_cli_text_output_and_no_review(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    requests: list[PromptGenerationRequest] = []

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        requests.append(request)
        final = PromptPassResult(
            fields={},
            positive_prompt="copyable prompt",
            negative_prompt="blurry",
        )
        return PromptGenerationResult(
            draft=final,
            final=final,
            fields=final.fields,
            positive_prompt=final.positive_prompt,
            negative_prompt=final.negative_prompt,
            assumptions=(),
            missing_input=(),
            review_issues=(),
            models=PromptGenerationModels(draft="model"),
        )

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "standing in forest",
            "--text",
            "--no-review",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert result.output == "copyable prompt\n\nNegative prompt: blurry\n"
    assert requests[0].review_enabled is False


def test_llm_prompt_cli_format_text_matches_text_alias(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "standing in forest",
            "--format",
            "text",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert result.output == "final prompt\n\nNegative prompt: final negative\n"


def test_llm_prompt_cli_format_review_prints_versions_and_comments(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "rear view forest",
            "--format",
            "review",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert result.output == (
        "First version:\n"
        "draft prompt\n\n"
        "Negative prompt: draft negative\n\n"
        "Second version:\n"
        "final prompt\n\n"
        "Negative prompt: final negative\n\n"
        "Comments:\n"
        "- Review issue: view: rear view mentioned a front-facing expression "
        "Fix: changed the view to over-the-shoulder\n"
        "- Assumption: kept the requested forest setting\n"
        "- Missing input: lighting was not specified\n"
    )
    assert '"models"' not in result.output


def test_llm_prompt_cli_format_review_prints_empty_comments(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        final = PromptPassResult(fields={}, positive_prompt="plain prompt")
        return PromptGenerationResult(
            draft=final,
            final=final,
            fields=final.fields,
            positive_prompt=final.positive_prompt,
            negative_prompt=final.negative_prompt,
            assumptions=(),
            missing_input=(),
            review_issues=(),
            models=PromptGenerationModels(draft="model"),
        )

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "simple scene",
            "--format",
            "review",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert "Comments:\nNo review comments.\n" in result.output


def test_llm_prompt_cli_text_alias_and_format_conflict(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        raise AssertionError("generate_prompt should not run with invalid flags")

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "standing in forest",
            "--text",
            "--format",
            "review",
        ],
    )

    assert result.exit_code != 0
    assert "--text cannot be combined with --format" in result.output


def test_llm_prompt_cli_verbose_and_format_conflict(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        raise AssertionError("generate_prompt should not run with invalid flags")

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "standing in forest",
            "--verbose",
            "--format",
            "json",
        ],
    )

    assert result.exit_code != 0
    assert "--verbose cannot be combined with --format" in result.output


def test_llm_prompt_cli_text_alias_and_verbose_conflict(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        raise AssertionError("generate_prompt should not run with invalid flags")

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "standing in forest",
            "--text",
            "--verbose",
        ],
    )

    assert result.exit_code != 0
    assert "--text cannot be combined with --verbose" in result.output


def test_llm_prompt_cli_exports_result_files_under_storage_by_default(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        ["--skip-dotenv-layers", "llm", "prompt", "stored scene", "--verbose"],
    )

    assert result.exit_code == 0, result.output
    assert json.loads(result.stdout)["positive_prompt"] == "final prompt"
    result_dir = isolated_cli_storage / ".llm_promptgen" / "results"
    json_files = sorted(result_dir.glob("*-prompt-result.json"))
    text_files = sorted(result_dir.glob("*-prompt-result.txt"))
    assert len(json_files) == 1
    assert len(text_files) == 1
    assert (
        json.loads(json_files[0].read_text(encoding="utf-8"))["positive_prompt"]
        == "final prompt"
    )
    assert "First version:\ndraft prompt" in text_files[0].read_text(encoding="utf-8")
    assert "Saved prompt result JSON:" in result.stderr
    assert "Saved prompt result text:" in result.stderr


def test_llm_prompt_cli_no_export_suppresses_storage_result_files(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
) -> None:
    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        ["--skip-dotenv-layers", "llm", "prompt", "stored scene", "--no-export"],
    )

    assert result.exit_code == 0, result.output
    assert "Saved prompt result" not in result.stderr
    assert not (isolated_cli_storage / ".llm_promptgen" / "results").exists()


def test_llm_prompt_cli_save_exchanges_defaults_to_storage_dir(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
) -> None:
    requests: list[PromptGenerationRequest] = []

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        requests.append(request)
        final = PromptPassResult(fields={}, positive_prompt="stored prompt")
        return PromptGenerationResult(
            draft=final,
            final=final,
            fields=final.fields,
            positive_prompt=final.positive_prompt,
            negative_prompt=final.negative_prompt,
            assumptions=(),
            missing_input=(),
            review_issues=(),
            models=PromptGenerationModels(draft="model"),
        )

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "stored scene",
            "--save-exchanges",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert requests[0].save_exchanges is True
    assert requests[0].exchange_dir == isolated_cli_storage / ".llm_promptgen"


def test_llm_prompt_cli_resolves_exchange_dir_when_enabled_by_config(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
) -> None:
    requests: list[PromptGenerationRequest] = []
    prompt_configs: list[PromptgenConfig] = []
    monkeypatch.setenv("KNF_PROMPTGEN_SAVE_EXCHANGES", "true")

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        *,
        lmstudio_config: LmStudioConfig,
        promptgen_config: PromptgenConfig,
        field_catalog: PromptFieldCatalog,
    ) -> PromptGenerationResult:
        requests.append(request)
        prompt_configs.append(promptgen_config)
        assert "subject" in field_catalog.fields
        return _reviewed_prompt_result()

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "stored scene",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert prompt_configs[0].save_exchanges is True
    assert requests[0].save_exchanges is None
    assert requests[0].exchange_dir == isolated_cli_storage / ".llm_promptgen"


def test_llm_prompt_cli_relative_exchange_override_resolves_under_storage(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
) -> None:
    requests: list[PromptGenerationRequest] = []
    monkeypatch.setenv("KNF_PROMPTGEN_EXCHANGE_DIR", "debug/llm")

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        requests.append(request)
        final = PromptPassResult(fields={}, positive_prompt="stored prompt")
        return PromptGenerationResult(
            draft=final,
            final=final,
            fields=final.fields,
            positive_prompt=final.positive_prompt,
            negative_prompt=final.negative_prompt,
            assumptions=(),
            missing_input=(),
            review_issues=(),
            models=PromptGenerationModels(draft="model"),
        )

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "stored scene",
            "--save-exchanges",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert requests[0].exchange_dir == isolated_cli_storage / "debug" / "llm"


def test_llm_prompt_cli_absolute_exchange_override_is_preserved(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
    tmp_path: Path,
) -> None:
    requests: list[PromptGenerationRequest] = []
    exchange_root = tmp_path / "external-exchanges"
    assert isolated_cli_storage.is_dir()
    monkeypatch.setenv("KNF_PROMPTGEN_EXCHANGE_DIR", str(exchange_root))

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        requests.append(request)
        final = PromptPassResult(fields={}, positive_prompt="stored prompt")
        return PromptGenerationResult(
            draft=final,
            final=final,
            fields=final.fields,
            positive_prompt=final.positive_prompt,
            negative_prompt=final.negative_prompt,
            assumptions=(),
            missing_input=(),
            review_issues=(),
            models=PromptGenerationModels(draft="model"),
        )

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "stored scene",
            "--save-exchanges",
            "--no-export",
        ],
    )

    assert result.exit_code == 0, result.output
    assert requests[0].exchange_dir == exchange_root


def test_llm_prompt_cli_rejects_relative_exchange_escape(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
) -> None:
    outside = isolated_cli_storage.parent / "outside-exchanges"
    monkeypatch.setenv("KNF_PROMPTGEN_EXCHANGE_DIR", "../outside-exchanges")

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        raise AssertionError(f"Unexpected prompt request: {request}")

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    result = invoke_cli(
        CliRunner(),
        [
            "--skip-dotenv-layers",
            "llm",
            "prompt",
            "stored scene",
            "--save-exchanges",
            "--no-export",
        ],
    )

    assert result.exit_code != 0
    assert "must stay inside selected project root" in result.output
    assert not outside.exists()


def test_llm_prompt_cli_rejects_symlinked_result_directory(
    monkeypatch: pytest.MonkeyPatch,
    isolated_cli_storage: Path,
    tmp_path: Path,
) -> None:
    outside = tmp_path / "outside-results"
    outside.mkdir()
    (isolated_cli_storage / ".llm_promptgen").symlink_to(
        outside,
        target_is_directory=True,
    )

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        raise AssertionError(f"Unexpected prompt request: {request}")

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    result = invoke_cli(
        CliRunner(),
        ["--skip-dotenv-layers", "llm", "prompt", "stored scene"],
    )

    assert result.exit_code != 0
    assert "must not be a symlink" in result.output
    assert not tuple(outside.iterdir())


def test_llm_prompt_cli_requires_storage(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv("KNF_STORAGE")

    def fake_generate_prompt(
        request: PromptGenerationRequest,
        **_: object,
    ) -> PromptGenerationResult:
        raise AssertionError("generate_prompt should not run without storage")

    monkeypatch.setattr(cli_llm, "generate_prompt", fake_generate_prompt)

    runner = CliRunner()
    result = invoke_cli(
        runner,
        ["--skip-dotenv-layers", "llm", "prompt", "unstored scene"],
    )

    assert result.exit_code != 0
    assert "KNF_STORAGE is required" in result.output