Enum: add tests

Author: irgolic

tests/integration_tests/mock_llm_outputs.py

@@ -28,6 +28,8 @@ def _invoke_llm(self, prompt, *args, **kwargs):
             pydantic.COMPILED_PROMPT_FULL_REASK_1: pydantic.LLM_OUTPUT_FULL_REASK_1,
             pydantic.COMPILED_PROMPT_REASK_2: pydantic.LLM_OUTPUT_REASK_2,
             pydantic.COMPILED_PROMPT_FULL_REASK_2: pydantic.LLM_OUTPUT_FULL_REASK_2,
+            pydantic.COMPILED_PROMPT_ENUM: pydantic.LLM_OUTPUT_ENUM,
+            pydantic.COMPILED_PROMPT_ENUM_2: pydantic.LLM_OUTPUT_ENUM_2,
             string.COMPILED_PROMPT: string.LLM_OUTPUT,
             string.COMPILED_PROMPT_REASK: string.LLM_OUTPUT_REASK,
             string.COMPILED_LIST_PROMPT: string.LIST_LLM_OUTPUT,

tests/integration_tests/test_assets/pydantic/__init__.py

@@ -31,12 +31,16 @@
 COMPILED_PROMPT_REASK_2 = reader("compiled_prompt_reask_2.txt")
 COMPILED_PROMPT_FULL_REASK_2 = reader("compiled_prompt_full_reask_2.txt")
 COMPILED_INSTRUCTIONS_REASK_2 = reader("compiled_instructions_reask_2.txt")
+COMPILED_PROMPT_ENUM = reader("compiled_prompt_enum.txt")
+COMPILED_PROMPT_ENUM_2 = reader("compiled_prompt_enum_2.txt")
 
 LLM_OUTPUT = reader("llm_output.txt")
 LLM_OUTPUT_REASK_1 = reader("llm_output_reask_1.txt")
 LLM_OUTPUT_FULL_REASK_1 = reader("llm_output_full_reask_1.txt")
 LLM_OUTPUT_REASK_2 = reader("llm_output_reask_2.txt")
 LLM_OUTPUT_FULL_REASK_2 = reader("llm_output_full_reask_2.txt")
+LLM_OUTPUT_ENUM = reader("llm_output_enum.txt")
+LLM_OUTPUT_ENUM_2 = reader("llm_output_enum_2.txt")
 
 RAIL_SPEC_WITH_REASK = reader("reask.rail")
 

tests/integration_tests/test_assets/pydantic/compiled_prompt_enum.txt

@@ -0,0 +1,4 @@
+What is the status of this task?
+
+Json Output:
+

tests/integration_tests/test_assets/pydantic/compiled_prompt_enum_2.txt

@@ -0,0 +1,4 @@
+What is the status of this task REALLY?
+
+Json Output:
+

tests/integration_tests/test_assets/pydantic/llm_output_enum.txt

@@ -0,0 +1 @@
+{"status": "not started"}

tests/integration_tests/test_assets/pydantic/llm_output_enum_2.txt

@@ -0,0 +1 @@
+{"status": "i dont know?"}

tests/integration_tests/test_guard.py

@@ -1,3 +1,4 @@
+import enum
 import json
 import os
 from typing import Optional, Union
@@ -782,3 +783,29 @@ def test_in_memory_validator_log_is_not_duplicated(mocker):
 
     finally:
         OneLine.run_in_separate_process = separate_proc_bak
+
+
+def test_enum_datatype(mocker):
+    mocker.patch("guardrails.llm_providers.OpenAICallable", new=MockOpenAICallable)
+
+    class TaskStatus(enum.Enum):
+        not_started = "not started"
+        on_hold = "on hold"
+        in_progress = "in progress"
+
+    class Task(BaseModel):
+        status: TaskStatus
+
+    guard = gd.Guard.from_pydantic(Task)
+    _, dict_o = guard(
+        get_static_openai_create_func(),
+        prompt="What is the status of this task?",
+    )
+    assert dict_o == {"status": "not started"}
+
+    guard = gd.Guard.from_pydantic(Task)
+    with pytest.raises(ValueError):
+        guard(
+            get_static_openai_create_func(),
+            prompt="What is the status of this task REALLY?",
+        )