feat: Add Prompt class basic offline functionality

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3932 from googleapis:release-please--branches--main 346f4c0
PiperOrigin-RevId: 642433403

tests/unit/vertexai/test_prompt.py

@@ -0,0 +1,186 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Unit tests for generative model prompts."""
+# pylint: disable=protected-access,bad-continuation
+
+from vertexai.generative_models._prompt import Prompt
+from vertexai.generative_models import Content, Part
+
+import pytest
+
+from typing import Any
+
+
+def _is_list_type(obj: Any, T: Any) -> bool:
+    return isinstance(obj, list) and all(isinstance(s, T) for s in obj)
+
+
+@pytest.mark.usefixtures("google_auth_mock")
+class TestPrompt:
+    """Unit tests for generative model prompts."""
+
+    def test_inferred_prompt_constructor(self):
+        # Create inferred prompt with string only variable values
+        prompt = Prompt(
+            prompt_data="Rate the movie {movie1}",
+            variables={
+                "movie1": "The Avengers",
+            },
+        )
+        # Inferred prompt data should remain as string before compilation
+        assert prompt.prompt_data == "Rate the movie {movie1}"
+        # Variables values should be converted to List[Part]
+        assert _is_list_type(prompt.variables["movie1"], Part)
+
+        # Create inferred prompt with List[Part] variable values
+        prompt = Prompt(
+            prompt_data="Rate the movie {movie1}",
+            variables={
+                "movie1": [Part.from_text("The Avengers")],
+            },
+        )
+        # Variables values should be converted to List[Part]
+        assert _is_list_type(prompt.variables["movie1"], Part)
+
+        # Inferred prompt variables must either be string or List[Part]
+        with pytest.raises(ValueError):
+            Prompt(
+                prompt_data="Rate the movie {movie1}",
+                variables={
+                    "movie1": Part.from_text("The Avengers"),
+                },
+            )
+
+    def test_inferred_prompt_to_content(self):
+        prompt = Prompt(
+            prompt_data="Which movie is better, {movie1} or {movie2}?",
+            variables={
+                "movie1": "The Avengers",
+                "movie2": "Frozen",
+            },
+        )
+        unassembled_prompt_content = prompt.to_content()
+        expected_content = [
+            Content(
+                parts=[
+                    Part.from_text("Which movie is better, {movie1} or {movie2}?"),
+                ],
+                role="user",
+            )
+        ]
+        assert unassembled_prompt_content[0].role == expected_content[0].role
+        for i in range(len(unassembled_prompt_content[0].parts)):
+            assert (
+                unassembled_prompt_content[0].parts[i].text
+                == expected_content[0].parts[i].text
+            )
+
+        # Check assembled prompt content
+        prompt.assemble()
+        assembled_prompt_content = prompt.to_content()
+        expected_content = [
+            Content(
+                parts=[
+                    Part.from_text("Which movie is better, "),
+                    Part.from_text("The Avengers"),
+                    Part.from_text(" or "),
+                    Part.from_text("Frozen"),
+                    Part.from_text("?"),
+                ],
+                role="user",
+            )
+        ]
+        assert assembled_prompt_content[0].role == expected_content[0].role
+        for i in range(len(assembled_prompt_content[0].parts)):
+            assert (
+                assembled_prompt_content[0].parts[i].text
+                == expected_content[0].parts[i].text
+            )
+
+    def test_inferred_prompt_assemble(self):
+        prompt = Prompt(
+            prompt_data="Which movie is better, {movie1} or {movie2}?",
+            variables={
+                "movie1": "The Avengers",
+            },
+        )
+
+        # Check partially assembled prompt content
+        prompt.assemble()
+        assembled1_prompt_content = prompt.to_content()
+        expected1_content = [
+            Content(
+                parts=[
+                    Part.from_text("Which movie is better, "),
+                    Part.from_text("The Avengers"),
+                    Part.from_text(" or "),
+                    Part.from_text("{movie2}"),
+                    Part.from_text("?"),
+                ],
+                role="user",
+            )
+        ]
+        assert assembled1_prompt_content[0].role == expected1_content[0].role
+        for i in range(len(assembled1_prompt_content[0].parts)):
+            assert (
+                assembled1_prompt_content[0].parts[i].text
+                == expected1_content[0].parts[i].text
+            )
+
+        # Check assembled prompt content
+        prompt.assemble(movie2="Frozen")
+        assembled2_prompt_content = prompt.to_content()
+        expected2_content = [
+            Content(
+                parts=[
+                    Part.from_text("Which movie is better, "),
+                    Part.from_text("The Avengers"),
+                    Part.from_text(" or "),
+                    Part.from_text("Frozen"),
+                    Part.from_text("?"),
+                ],
+                role="user",
+            )
+        ]
+        assert assembled2_prompt_content[0].role == expected2_content[0].role
+        for i in range(len(assembled2_prompt_content[0].parts)):
+            assert (
+                assembled2_prompt_content[0].parts[i].text
+                == expected2_content[0].parts[i].text
+            )
+
+        # Check assembled prompt content with override
+        prompt.assemble(movie1="Inception")
+        assembled3_prompt_content = prompt.to_content()
+        expected3_content = [
+            Content(
+                parts=[
+                    Part.from_text("Which movie is better, "),
+                    Part.from_text("Inception"),
+                    Part.from_text(" or "),
+                    Part.from_text("Frozen"),
+                    Part.from_text("?"),
+                ],
+                role="user",
+            )
+        ]
+        assert assembled3_prompt_content[0].role == expected3_content[0].role
+        for i in range(len(assembled3_prompt_content[0].parts)):
+            assert (
+                assembled3_prompt_content[0].parts[i].text
+                == expected3_content[0].parts[i].text
+            )