Skyvern/tests/unit/test_text_prompt_block.py

import sys
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from jinja2 import StrictUndefined
from jinja2.sandbox import SandboxedEnvironment

import skyvern.forge.sdk.workflow.models.block as _block_mod
from skyvern.config import settings as base_settings
from skyvern.forge.prompts import prompt_engine
from skyvern.forge.sdk.settings_manager import SettingsManager
from skyvern.forge.sdk.workflow.context_manager import WorkflowRunContext
from skyvern.forge.sdk.workflow.models.block import TextPromptBlock
from skyvern.forge.sdk.workflow.models.parameter import OutputParameter, ParameterType
from skyvern.schemas.workflows import TextPromptBlockYAML, WorkflowRequest

block_module = sys.modules["skyvern.forge.sdk.workflow.models.block"]


def _make_workflow_run_context(values: dict | None = None) -> WorkflowRunContext:
    """Create a minimal WorkflowRunContext with given values for template rendering tests."""
    ctx = WorkflowRunContext(
        workflow_title="test",
        workflow_id="w_test",
        workflow_permanent_id="wpid_test",
        workflow_run_id="wr_test",
        aws_client=MagicMock(),
    )
    if values:
        ctx.values.update(values)
    return ctx


def _make_text_prompt_block(
    prompt: str = "test prompt",
    json_schema: dict | None = None,
) -> TextPromptBlock:
    now = datetime.now(timezone.utc)
    return TextPromptBlock(
        label="test-block",
        llm_key=None,
        prompt=prompt,
        parameters=[],
        json_schema=json_schema,
        output_parameter=OutputParameter(
            parameter_type=ParameterType.OUTPUT,
            key="text_prompt_output",
            description=None,
            output_parameter_id="output-test",
            workflow_id="workflow-1",
            created_at=now,
            modified_at=now,
            deleted_at=None,
        ),
        model=None,
    )


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "model_name",
    ["gemini-2.5-flash", "gemini-3-pro-preview"],
)
async def test_text_prompt_block_uses_selected_model(monkeypatch, model_name):
    # Reset SettingsManager to base settings so cloud overrides from earlier tests don't leak
    monkeypatch.setattr(SettingsManager, "_SettingsManager__instance", base_settings)
    expected_llm_key = base_settings.get_model_name_to_llm_key()[model_name]["llm_key"]
    now = datetime.now(timezone.utc)
    output_parameter = OutputParameter(
        parameter_type=ParameterType.OUTPUT,
        key="text_prompt_output",
        description=None,
        output_parameter_id="output-1",
        workflow_id="workflow-1",
        created_at=now,
        modified_at=now,
        deleted_at=None,
    )

    block = TextPromptBlock(
        label="text-block",
        llm_key="AZURE_OPENAI_GPT4_1",
        prompt="Explain the status.",
        parameters=[],
        json_schema=None,
        output_parameter=output_parameter,
        model={"model_name": model_name},
    )

    captured: dict[str, str] = {}
    fake_default_handler = AsyncMock()

    async def fake_resolve_default_llm_handler(*args, **kwargs):
        return fake_default_handler

    async def fake_handler(*, prompt: str, prompt_name: str, **kwargs):
        captured["prompt"] = prompt
        captured["prompt_name"] = prompt_name
        return {"llm_response": "ok"}

    def fake_get_override_handler(llm_key: str | None, *, default):
        captured["llm_key"] = llm_key if llm_key else "default"
        return fake_handler if llm_key else default

    block_module.app.LLM_API_HANDLER = fake_default_handler
    LLMAPIHandlerFactory = block_module.LLMAPIHandlerFactory
    monkeypatch.setattr(
        LLMAPIHandlerFactory,
        "get_override_llm_api_handler",
        fake_get_override_handler,
        raising=False,
    )
    monkeypatch.setattr(
        TextPromptBlock,
        "_resolve_default_llm_handler",
        fake_resolve_default_llm_handler,
        raising=False,
    )
    monkeypatch.setattr(
        prompt_engine,
        "load_prompt_from_string",
        lambda template, **kwargs: template,
    )

    response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")

    assert captured["llm_key"] == expected_llm_key
    assert captured["prompt_name"] == "text-prompt"
    assert response == {"llm_response": "ok"}


@pytest.mark.asyncio
async def test_text_prompt_block_uses_workflow_handler_when_no_override(monkeypatch):
    now = datetime.now(timezone.utc)
    output_parameter = OutputParameter(
        parameter_type=ParameterType.OUTPUT,
        key="text_prompt_output",
        description=None,
        output_parameter_id="output-2",
        workflow_id="workflow-1",
        created_at=now,
        modified_at=now,
        deleted_at=None,
    )

    block = TextPromptBlock(
        label="text-block",
        llm_key=None,
        prompt="Summarize status.",
        parameters=[],
        json_schema=None,
        output_parameter=output_parameter,
        model=None,
    )

    captured: dict[str, str] = {}
    fake_secondary_handler = AsyncMock(return_value={"llm_response": "secondary"})

    async def fake_prompt_type_handler(*args, **kwargs):
        return None

    def fake_get_override_handler(llm_key: str | None, *, default):
        captured["llm_key"] = llm_key if llm_key else "default"
        captured["default_handler"] = default
        return default

    block_module.app.SECONDARY_LLM_API_HANDLER = fake_secondary_handler
    block_module.app.LLM_API_HANDLER = AsyncMock()
    LLMAPIHandlerFactory = block_module.LLMAPIHandlerFactory
    monkeypatch.setattr(
        LLMAPIHandlerFactory,
        "get_override_llm_api_handler",
        fake_get_override_handler,
        raising=False,
    )
    monkeypatch.setattr(
        block_module,
        "get_llm_handler_for_prompt_type",
        fake_prompt_type_handler,
        raising=False,
    )
    monkeypatch.setattr(
        prompt_engine,
        "load_prompt_from_string",
        lambda template, **kwargs: template,
    )

    response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")

    assert captured["llm_key"] == "default"
    assert captured["default_handler"] == fake_secondary_handler
    fake_secondary_handler.assert_awaited_once()
    assert response == {"llm_response": "secondary"}


@pytest.mark.asyncio
async def test_text_prompt_block_prefers_prompt_type_config_over_secondary(monkeypatch):
    now = datetime.now(timezone.utc)
    output_parameter = OutputParameter(
        parameter_type=ParameterType.OUTPUT,
        key="text_prompt_output",
        description=None,
        output_parameter_id="output-3",
        workflow_id="workflow-1",
        created_at=now,
        modified_at=now,
        deleted_at=None,
    )

    block = TextPromptBlock(
        label="text-block",
        llm_key=None,
        prompt="Provide summary.",
        parameters=[],
        json_schema=None,
        output_parameter=output_parameter,
        model=None,
    )

    captured: dict[str, str] = {}
    prompt_config_handler = AsyncMock(return_value={"llm_response": "config"})

    async def fake_prompt_type_handler(*args, **kwargs):
        return prompt_config_handler

    def fake_get_override_handler(llm_key: str | None, *, default):
        captured["default_handler"] = default
        return default

    block_module.app.SECONDARY_LLM_API_HANDLER = AsyncMock()
    block_module.app.LLM_API_HANDLER = AsyncMock()
    LLMAPIHandlerFactory = block_module.LLMAPIHandlerFactory
    monkeypatch.setattr(
        LLMAPIHandlerFactory,
        "get_override_llm_api_handler",
        fake_get_override_handler,
        raising=False,
    )
    monkeypatch.setattr(
        block_module,
        "get_llm_handler_for_prompt_type",
        fake_prompt_type_handler,
        raising=False,
    )
    monkeypatch.setattr(
        prompt_engine,
        "load_prompt_from_string",
        lambda template, **kwargs: template,
    )

    response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")

    assert captured["default_handler"] == prompt_config_handler
    prompt_config_handler.assert_awaited_once()
    assert response == {"llm_response": "config"}


@pytest.mark.asyncio
async def test_text_prompt_block_bad_llm_key_uses_same_runtime_path_as_no_override(monkeypatch):
    monkeypatch.setattr(SettingsManager, "_SettingsManager__instance", base_settings)
    now = datetime.now(timezone.utc)

    normalized_bad = TextPromptBlockYAML(
        label="bad_key",
        prompt="Summarize status.",
        llm_key="ANTHROPIC_CLAUDE_3_5_SONNET",
    )
    no_override = TextPromptBlockYAML(
        label="no_override",
        prompt="Summarize status.",
        llm_key=None,
    )

    blocks = []
    for idx, yaml_block in enumerate((normalized_bad, no_override), start=1):
        output_parameter = OutputParameter(
            parameter_type=ParameterType.OUTPUT,
            key=f"text_prompt_output_{idx}",
            description=None,
            output_parameter_id=f"output-{idx}",
            workflow_id="workflow-1",
            created_at=now,
            modified_at=now,
            deleted_at=None,
        )
        blocks.append(
            TextPromptBlock(
                label=yaml_block.label,
                llm_key=yaml_block.llm_key,
                prompt=yaml_block.prompt,
                parameters=[],
                json_schema=None,
                output_parameter=output_parameter,
                model=yaml_block.model,
            )
        )

    captured: list[tuple[str | None, object]] = []
    fake_secondary_handler = AsyncMock(return_value={"llm_response": "secondary"})

    async def fake_prompt_type_handler(*args, **kwargs):
        return None

    def fake_get_override_handler(llm_key: str | None, *, default):
        captured.append((llm_key, default))
        return default

    block_module.app.SECONDARY_LLM_API_HANDLER = fake_secondary_handler
    block_module.app.LLM_API_HANDLER = AsyncMock()
    LLMAPIHandlerFactory = block_module.LLMAPIHandlerFactory
    monkeypatch.setattr(
        LLMAPIHandlerFactory,
        "get_override_llm_api_handler",
        fake_get_override_handler,
        raising=False,
    )
    monkeypatch.setattr(
        block_module,
        "get_llm_handler_for_prompt_type",
        fake_prompt_type_handler,
        raising=False,
    )
    monkeypatch.setattr(
        prompt_engine,
        "load_prompt_from_string",
        lambda template, **kwargs: template,
    )

    for block in blocks:
        response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")
        assert response == {"llm_response": "secondary"}

    assert captured == [
        (None, fake_secondary_handler),
        (None, fake_secondary_handler),
    ]


@pytest.mark.asyncio
async def test_text_prompt_block_uses_explicit_internal_llm_key_override(monkeypatch):
    now = datetime.now(timezone.utc)
    output_parameter = OutputParameter(
        parameter_type=ParameterType.OUTPUT,
        key="text_prompt_output_internal",
        description=None,
        output_parameter_id="output-internal",
        workflow_id="workflow-1",
        created_at=now,
        modified_at=now,
        deleted_at=None,
    )

    block = TextPromptBlock(
        label="text-block",
        llm_key="SPECIAL_INTERNAL_KEY",
        prompt="Summarize status.",
        parameters=[],
        json_schema=None,
        output_parameter=output_parameter,
        model=None,
    )

    captured: dict[str, object] = {}
    fake_default_handler = AsyncMock()
    fake_override_handler = AsyncMock(return_value={"llm_response": "override"})

    async def fake_resolve_default_llm_handler(*args, **kwargs):
        return fake_default_handler

    def fake_get_override_handler(llm_key: str | None, *, default):
        captured["llm_key"] = llm_key
        captured["default_handler"] = default
        return fake_override_handler

    block_module.app.LLM_API_HANDLER = fake_default_handler
    LLMAPIHandlerFactory = block_module.LLMAPIHandlerFactory
    monkeypatch.setattr(
        LLMAPIHandlerFactory,
        "get_override_llm_api_handler",
        fake_get_override_handler,
        raising=False,
    )
    monkeypatch.setattr(
        TextPromptBlock,
        "_resolve_default_llm_handler",
        fake_resolve_default_llm_handler,
        raising=False,
    )
    monkeypatch.setattr(
        prompt_engine,
        "load_prompt_from_string",
        lambda template, **kwargs: template,
    )

    response = await block.send_prompt(block.prompt, {}, workflow_run_id="workflow-run", organization_id="org-1")

    assert captured["llm_key"] == "SPECIAL_INTERNAL_KEY"
    assert captured["default_handler"] == fake_default_handler
    fake_override_handler.assert_awaited_once()
    assert response == {"llm_response": "override"}


def test_workflow_request_deserialization_strips_invalid_text_prompt_llm_key() -> None:
    """Verify FastAPI deserialization (not just explicit model_validate) strips bad keys.

    Moved from tests/scenario/ — this test only validates Pydantic behavior and
    does not require a database connection.
    """
    raw_request = {
        "json_definition": {
            "title": "Deserialization test",
            "workflow_definition": {
                "blocks": [
                    {
                        "block_type": "text_prompt",
                        "label": "summarize",
                        "prompt": "Summarize the result",
                        "llm_key": "ANTHROPIC_CLAUDE_3_5_SONNET",
                    }
                ],
                "parameters": [],
            },
        }
    }
    workflow_request = WorkflowRequest.model_validate(raw_request)

    block = workflow_request.json_definition.workflow_definition.blocks[0]
    assert block.llm_key is None
    assert block.model is None


# --- json_schema Jinja template rendering tests (SKY-6479) ---


def test_render_schema_templates_resolves_variables():
    """json_schema description fields with {{ }} should be rendered."""
    block = _make_text_prompt_block(
        json_schema={
            "type": "object",
            "properties": {
                "result": {
                    "type": "string",
                    "description": "Add 14 days to {{ start_date }}",
                }
            },
        },
    )
    ctx = _make_workflow_run_context({"start_date": "2025-01-15"})
    rendered = block._render_schema_templates(block.json_schema, ctx)

    assert rendered["properties"]["result"]["description"] == "Add 14 days to 2025-01-15"
    # non-template values preserved
    assert rendered["type"] == "object"
    assert rendered["properties"]["result"]["type"] == "string"


def test_render_schema_templates_nested_list():
    """Templates inside lists within the schema should also be rendered."""
    block = _make_text_prompt_block(
        json_schema={
            "type": "object",
            "required": ["{{ field_name }}"],
        },
    )
    ctx = _make_workflow_run_context({"field_name": "invoice_date"})
    rendered = block._render_schema_templates(block.json_schema, ctx)

    assert rendered["required"] == ["invoice_date"]


def test_render_schema_templates_failure_preserves_original():
    """When a single value fails to render, it should be preserved and a warning logged."""
    block = _make_text_prompt_block(
        json_schema={
            "type": "object",
            "properties": {
                "good": {"description": "value is {{ known }}"},
                "bad": {"description": "value is {{ unknown }}"},
            },
        },
    )
    strict_env = SandboxedEnvironment(undefined=StrictUndefined)
    with (
        patch.object(base_settings, "WORKFLOW_TEMPLATING_STRICTNESS", "strict"),
        patch.object(_block_mod, "jinja_sandbox_env", strict_env),
    ):
        ctx = _make_workflow_run_context({"known": "hello"})
        rendered = block._render_schema_templates(block.json_schema, ctx)

        # good value rendered
        assert rendered["properties"]["good"]["description"] == "value is hello"
        # bad value preserved as-is
        assert rendered["properties"]["bad"]["description"] == "value is {{ unknown }}"


def test_format_potential_template_parameters_renders_json_schema():
    """format_potential_template_parameters should render json_schema templates."""
    block = _make_text_prompt_block(
        prompt="Calculate dates for {{ start_date }}",
        json_schema={
            "type": "object",
            "properties": {
                "date_plus_14": {
                    "type": "string",
                    "description": "14 days after {{ start_date }}",
                }
            },
        },
    )
    ctx = _make_workflow_run_context({"start_date": "2025-01-15"})
    block.format_potential_template_parameters(ctx)

    assert block.prompt == "Calculate dates for 2025-01-15"
    assert block.json_schema["properties"]["date_plus_14"]["description"] == "14 days after 2025-01-15"


def test_format_potential_template_parameters_no_json_schema():
    """When json_schema is None, format_potential_template_parameters should not fail."""
    block = _make_text_prompt_block(prompt="simple prompt", json_schema=None)
    ctx = _make_workflow_run_context()
    block.format_potential_template_parameters(ctx)

    assert block.json_schema is None
    assert block.prompt == "simple prompt"