litellm updates to support only messages llm kwarg

guardrails/actions/reask.py

@@ -247,7 +247,7 @@ def get_reask_setup_for_string(
     validation_response: Optional[Union[str, List, Dict, ReAsk]] = None,
     prompt_params: Optional[Dict[str, Any]] = None,
     exec_options: Optional[GuardExecutionOptions] = None,
-) -> Tuple[Dict[str, Any], Prompt, Instructions]:
+) -> Tuple[Dict[str, Any], Messages]:
     prompt_params = prompt_params or {}
     exec_options = exec_options or GuardExecutionOptions()
 
@@ -259,13 +259,11 @@ def get_reask_setup_for_string(
     )
 
     reask_prompt_template = None
-    if exec_options.reask_prompt:
-        reask_prompt_template = Prompt(exec_options.reask_prompt)
-    else:
-        reask_prompt_template = Prompt(
-            constants["high_level_string_reask_prompt"]
-            + constants["complete_string_suffix"]
-        )
+
+    reask_prompt_template = Prompt(
+        constants["high_level_string_reask_prompt"]
+        + constants["complete_string_suffix"]
+    )
 
     error_messages = "\n".join(
         [
@@ -286,10 +284,7 @@ def get_reask_setup_for_string(
     )
 
     instructions = None
-    if exec_options.reask_instructions:
-        instructions = Instructions(exec_options.reask_instructions)
-    if instructions is None:
-        instructions = Instructions("You are a helpful assistant.")
+    instructions = Instructions("You are a helpful assistant.")
     instructions = instructions.format(
         output_schema=schema_prompt_content,
         xml_output_schema=xml_output_schema,
@@ -298,9 +293,13 @@ def get_reask_setup_for_string(
     messages = None
     if exec_options.reask_messages:
         messages = Messages(exec_options.reask_messages)
+
     if messages is None:
         messages = Messages(
-            [{"role": "system", "content": "You are a helpful assistant."}]
+            [
+                {"role": "system", "content": instructions},
+                {"role": "user", "content": prompt},
+            ]
         )
 
     messages = messages.format(
@@ -309,21 +308,21 @@ def get_reask_setup_for_string(
         **prompt_params,
     )
 
-    return output_schema, prompt, instructions
+    return output_schema, messages
 
 
 def get_original_prompt(exec_options: Optional[GuardExecutionOptions] = None) -> str:
     exec_options = exec_options or GuardExecutionOptions()
-    original_msg_history = exec_options.msg_history or []
-    msg_history_prompt = next(
+    original_messages = exec_options.messages or []
+    messages_prompt = next(
         (
             h.get("content")
-            for h in original_msg_history
+            for h in original_messages
             if isinstance(h, dict) and h.get("role") == "user"
         ),
         "",
     )
-    original_prompt = exec_options.prompt or msg_history_prompt or ""
+    original_prompt = messages_prompt or ""
     return original_prompt
 
 
@@ -338,7 +337,7 @@ def get_reask_setup_for_json(
     use_full_schema: Optional[bool] = False,
     prompt_params: Optional[Dict[str, Any]] = None,
     exec_options: Optional[GuardExecutionOptions] = None,
-) -> Tuple[Dict[str, Any], Prompt, Instructions]:
+) -> Tuple[Dict[str, Any], Messages]:
     reask_schema = output_schema
     is_skeleton_reask = not any(isinstance(reask, FieldReAsk) for reask in reasks)
     is_nonparseable_reask = any(
@@ -351,8 +350,6 @@ def get_reask_setup_for_json(
     use_xml = prompt_uses_xml(original_prompt)
 
     reask_prompt_template = None
-    if exec_options.reask_prompt:
-        reask_prompt_template = Prompt(exec_options.reask_prompt)
 
     if is_nonparseable_reask:
         if reask_prompt_template is None:
@@ -462,30 +459,26 @@ def reask_decoder(obj: ReAsk):
     )
 
     instructions = None
-    if exec_options.reask_instructions:
-        instructions = Instructions(exec_options.reask_instructions)
-    else:
-        instructions_const = (
-            constants["high_level_xml_instructions"]
-            if use_xml
-            else constants["high_level_json_instructions"]
-        )
-        instructions = Instructions(instructions_const)
+    instructions_const = (
+        constants["high_level_xml_instructions"]
+        if use_xml
+        else constants["high_level_json_instructions"]
+    )
+    instructions = Instructions(instructions_const)
     instructions = instructions.format(**prompt_params)
 
-    # TODO: enable this in 0.6.0
-    # messages = None
-    # if exec_options.reask_messages:
-    #     messages = Messages(exec_options.reask_messages)
-    # else:
-    #     messages = Messages(
-    #         [
-    #             {"role": "system", "content": instructions},
-    #             {"role": "user", "content": prompt},
-    #         ]
-    #     )
+    messages = None
+    if exec_options.reask_messages:
+        messages = Messages(exec_options.reask_messages)
+    else:
+        messages = Messages(
+            [
+                {"role": "system", "content": instructions},
+                {"role": "user", "content": prompt},
+            ]
+        )
 
-    return reask_schema, prompt, instructions
+    return reask_schema, messages
 
 
 def get_reask_setup(
@@ -499,7 +492,7 @@ def get_reask_setup(
     use_full_schema: Optional[bool] = False,
     prompt_params: Optional[Dict[str, Any]] = None,
     exec_options: Optional[GuardExecutionOptions] = None,
-) -> Tuple[Dict[str, Any], Prompt, Instructions]:
+) -> Tuple[Dict[str, Any], Messages]:
     prompt_params = prompt_params or {}
     exec_options = exec_options or GuardExecutionOptions()
 

guardrails/applications/text2sql.py

@@ -70,7 +70,12 @@ def __init__(
         rail_spec: Optional[str] = None,
         rail_params: Optional[Dict] = None,
         example_formatter: Callable = example_formatter,
-        reask_prompt: str = REASK_PROMPT,
+        reask_messages: list[Dict[str, str]] = [
+            {
+                "role": "user",
+                "content": REASK_PROMPT,
+            }
+        ],
         llm_api: Optional[Callable] = None,
         llm_api_kwargs: Optional[Dict] = None,
         num_relevant_examples: int = 2,
@@ -108,7 +113,7 @@ def __init__(
             schema_file,
             rail_spec,
             rail_params,
-            reask_prompt,
+            reask_messages,
         )
 
         # Initialize the document store.
@@ -122,7 +127,12 @@ def _init_guard(
         schema_file: Optional[str] = None,
         rail_spec: Optional[str] = None,
         rail_params: Optional[Dict] = None,
-        reask_prompt: str = REASK_PROMPT,
+        reask_messages: list[Dict[str, str]] = [
+            {
+                "role": "user",
+                "content": REASK_PROMPT,
+            }
+        ],
     ):
         # Initialize the Guard class
         if rail_spec is None:
@@ -140,7 +150,7 @@ def _init_guard(
             rail_spec_str = Template(rail_spec_str).safe_substitute(**rail_params)
 
         guard = Guard.for_rail_string(rail_spec_str)
-        guard._exec_opts.reask_prompt = reask_prompt
+        guard._exec_opts.reask_messages = reask_messages
 
         return guard