update some tests

Author: dtam

guardrails/llm_providers.py

@@ -379,7 +379,15 @@ def _invoke_llm(
 
 
 class HuggingFacePipelineCallable(PromptCallableBase):
-    def _invoke_llm(self, prompt: str, pipeline: Any, *args, **kwargs) -> LLMResponse:
+    def _invoke_llm(
+        self,
+        pipeline: Any,
+        *args,
+        messages: Union[
+            list[dict[str, Union[str, Prompt, Instructions]]], MessageHistory
+        ],
+        **kwargs,
+    ) -> LLMResponse:
         try:
             import transformers  # noqa: F401 # type: ignore
         except ImportError:
@@ -400,7 +408,7 @@ def _invoke_llm(self, prompt: str, pipeline: Any, *args, **kwargs) -> LLMRespons
         temperature = kwargs.pop("temperature", None)
         if temperature == 0:
             temperature = None
-
+        prompt = messages_string(messages)
         trace_operation(
             input_mime_type="application/json",
             input_value={

tests/integration_tests/test_guard.py

@@ -994,15 +994,17 @@ def custom_llm(
     guard = gd.Guard.for_pydantic(Task)
     _, dict_o, *rest = guard(
         custom_llm,
-        prompt="What is the status of this task?",
+        messages=[{"role": "user", "content": "What is the status of this task?"}],
     )
     assert dict_o == {"status": "not started"}
 
     return_value = pydantic.LLM_OUTPUT_ENUM_2
     guard = gd.Guard.for_pydantic(Task)
     result = guard(
         custom_llm,
-        prompt="What is the status of this task REALLY?",
+        messages=[
+            {"role": "user", "content": "What is the status of this task REALLY?"}
+        ],
         num_reasks=0,
     )
 
@@ -1351,7 +1353,9 @@ def test_guard_for_pydantic_with_mock_hf_pipeline():
 
     pipe = make_mock_pipeline()
     guard = Guard()
-    _ = guard(pipe, prompt="Don't care about the output.")
+    _ = guard(
+        pipe, messages=[{"role": "user", "content": "Don't care about the output."}]
+    )
 
 
 @pytest.mark.skipif(

tests/unit_tests/test_llm_providers.py

@@ -258,7 +258,9 @@ def test_hugging_face_pipeline_callable():
     from guardrails.llm_providers import HuggingFacePipelineCallable
 
     hf_model_callable = HuggingFacePipelineCallable()
-    response = hf_model_callable("Hello", pipeline=pipeline)
+    response = hf_model_callable(
+        pipeline=pipeline, messages=[{"role": "user", "content": "Hello"}]
+    )
 
     assert isinstance(response, LLMResponse) is True
     assert response.output == "Hello there!"