chore: import upstream snapshot with attribution

tests/backend/text_llm/test_text_llm_api_models.py

@@ -0,0 +1,54 @@
+"""Tests for TextLLM API request/response models and validation."""
+
+import pytest
+from pydantic import ValidationError
+
+from invokeai.app.api.routers.utilities import ExpandPromptRequest, ExpandPromptResponse, ImageToPromptRequest
+
+
+class TestExpandPromptRequest:
+    def test_defaults(self):
+        req = ExpandPromptRequest(prompt="a cat", model_key="abc-123")
+        assert req.max_tokens == 300
+        assert req.system_prompt is None
+
+    def test_max_tokens_upper_bound(self):
+        """max_tokens should be capped at 2048."""
+        with pytest.raises(ValidationError):
+            ExpandPromptRequest(prompt="a cat", model_key="abc-123", max_tokens=2049)
+
+    def test_max_tokens_lower_bound(self):
+        """max_tokens must be at least 1."""
+        with pytest.raises(ValidationError):
+            ExpandPromptRequest(prompt="a cat", model_key="abc-123", max_tokens=0)
+
+    def test_max_tokens_valid_range(self):
+        req = ExpandPromptRequest(prompt="a cat", model_key="abc-123", max_tokens=2048)
+        assert req.max_tokens == 2048
+        req2 = ExpandPromptRequest(prompt="a cat", model_key="abc-123", max_tokens=1)
+        assert req2.max_tokens == 1
+
+    def test_custom_system_prompt(self):
+        req = ExpandPromptRequest(prompt="a cat", model_key="abc-123", system_prompt="Be brief.")
+        assert req.system_prompt == "Be brief."
+
+
+class TestImageToPromptRequest:
+    def test_defaults(self):
+        req = ImageToPromptRequest(image_name="img.png", model_key="abc-123")
+        assert "Describe" in req.instruction
+
+    def test_custom_instruction(self):
+        req = ImageToPromptRequest(image_name="img.png", model_key="abc-123", instruction="What is this?")
+        assert req.instruction == "What is this?"
+
+
+class TestExpandPromptResponse:
+    def test_success_response(self):
+        resp = ExpandPromptResponse(expanded_prompt="A detailed scene")
+        assert resp.expanded_prompt == "A detailed scene"
+        assert resp.error is None
+
+    def test_error_response(self):
+        resp = ExpandPromptResponse(expanded_prompt="", error="Model failed")
+        assert resp.error == "Model failed"

tests/backend/text_llm/test_text_llm_pipeline.py

@@ -0,0 +1,130 @@
+"""Tests for the TextLLMPipeline class."""
+
+from unittest.mock import MagicMock
+
+import torch
+
+from invokeai.backend.text_llm_pipeline import DEFAULT_SYSTEM_PROMPT, TextLLMPipeline
+
+
+def _make_mock_tokenizer(has_chat_template: bool = True) -> MagicMock:
+    """Create a mock tokenizer with configurable chat template support."""
+    tokenizer = MagicMock()
+    if has_chat_template:
+        tokenizer.chat_template = "{% for m in messages %}{{ m.content }}{% endfor %}"
+        tokenizer.apply_chat_template.return_value = "<|system|>You are helpful<|user|>hello<|assistant|>"
+    else:
+        tokenizer.chat_template = None
+
+    # Simulate tokenizer __call__ returning dict with input_ids
+    input_ids = torch.tensor([[1, 2, 3, 4, 5]])
+    tokenizer_output = MagicMock()
+    tokenizer_output.__getitem__ = lambda self, key: {"input_ids": input_ids}[key]
+    tokenizer_output.to.return_value = tokenizer_output
+    tokenizer.return_value = tokenizer_output
+
+    tokenizer.decode.return_value = "A detailed landscape with mountains"
+    return tokenizer
+
+
+def _make_mock_model() -> MagicMock:
+    """Create a mock causal LM model."""
+    model = MagicMock()
+    # generate returns tensor that includes input + generated tokens
+    model.generate.return_value = torch.tensor([[1, 2, 3, 4, 5, 10, 11, 12]])
+    return model
+
+
+def test_pipeline_uses_chat_template_when_available():
+    """Pipeline should use apply_chat_template when the tokenizer supports it."""
+    tokenizer = _make_mock_tokenizer(has_chat_template=True)
+    model = _make_mock_model()
+    pipeline = TextLLMPipeline(model, tokenizer)
+
+    pipeline.run(prompt="a cat", device=torch.device("cpu"), dtype=torch.float32)
+
+    tokenizer.apply_chat_template.assert_called_once()
+    call_args = tokenizer.apply_chat_template.call_args
+    messages = call_args[0][0]
+    assert any(m["role"] == "system" for m in messages)
+    assert any(m["role"] == "user" and m["content"] == "a cat" for m in messages)
+
+
+def test_pipeline_fallback_without_chat_template():
+    """Pipeline should use fallback formatting when no chat template exists."""
+    tokenizer = _make_mock_tokenizer(has_chat_template=False)
+    model = _make_mock_model()
+    pipeline = TextLLMPipeline(model, tokenizer)
+
+    pipeline.run(prompt="a cat", system_prompt="Be helpful", device=torch.device("cpu"), dtype=torch.float32)
+
+    tokenizer.apply_chat_template.assert_not_called()
+    # Check that the tokenizer was called with the fallback format
+    call_args = tokenizer.call_args[0][0]
+    assert "Be helpful" in call_args
+    assert "a cat" in call_args
+    assert "Assistant:" in call_args
+
+
+def test_pipeline_no_system_prompt():
+    """Pipeline should work without a system prompt."""
+    tokenizer = _make_mock_tokenizer(has_chat_template=True)
+    model = _make_mock_model()
+    pipeline = TextLLMPipeline(model, tokenizer)
+
+    pipeline.run(prompt="a dog", system_prompt="", device=torch.device("cpu"), dtype=torch.float32)
+
+    call_args = tokenizer.apply_chat_template.call_args
+    messages = call_args[0][0]
+    # No system message when system_prompt is empty
+    assert not any(m["role"] == "system" for m in messages)
+    assert any(m["role"] == "user" and m["content"] == "a dog" for m in messages)
+
+
+def test_pipeline_decodes_only_generated_tokens():
+    """Pipeline should strip input tokens and only decode newly generated ones."""
+    tokenizer = _make_mock_tokenizer(has_chat_template=True)
+    model = _make_mock_model()
+    pipeline = TextLLMPipeline(model, tokenizer)
+
+    pipeline.run(prompt="test", device=torch.device("cpu"), dtype=torch.float32)
+
+    # The mock model returns [1,2,3,4,5,10,11,12], input is [1,2,3,4,5]
+    # So decode should be called with [10, 11, 12]
+    decode_call = tokenizer.decode.call_args
+    decoded_tokens = decode_call[0][0]
+    assert decoded_tokens.tolist() == [10, 11, 12]
+    assert decode_call[1]["skip_special_tokens"] is True
+
+
+def test_pipeline_passes_generation_params():
+    """Pipeline should pass max_new_tokens and sampling params to model.generate."""
+    tokenizer = _make_mock_tokenizer(has_chat_template=True)
+    model = _make_mock_model()
+    pipeline = TextLLMPipeline(model, tokenizer)
+
+    pipeline.run(prompt="test", max_new_tokens=100, device=torch.device("cpu"), dtype=torch.float32)
+
+    generate_kwargs = model.generate.call_args[1]
+    assert generate_kwargs["max_new_tokens"] == 100
+    assert generate_kwargs["do_sample"] is True
+    assert generate_kwargs["temperature"] == 0.7
+    assert generate_kwargs["top_p"] == 0.9
+
+
+def test_pipeline_returns_stripped_string():
+    """Pipeline should return a stripped string from the decoded output."""
+    tokenizer = _make_mock_tokenizer(has_chat_template=True)
+    tokenizer.decode.return_value = "  generated text with spaces  "
+    model = _make_mock_model()
+    pipeline = TextLLMPipeline(model, tokenizer)
+
+    result = pipeline.run(prompt="test", device=torch.device("cpu"), dtype=torch.float32)
+
+    assert result == "generated text with spaces"
+
+
+def test_default_system_prompt_content():
+    """The default system prompt should mention image generation."""
+    assert "image generation" in DEFAULT_SYSTEM_PROMPT.lower()
+    assert "prompt" in DEFAULT_SYSTEM_PROMPT.lower()