typing and lint

Author: dtam

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.from_rail_string(rail_spec_str)
-        guard._exec_opts.reask_prompt = reask_prompt
+        guard._exec_opts.reask_messages = reask_messages
 
         return guard
 

guardrails/classes/history/call.py

@@ -78,21 +78,23 @@ def prompt_params(self) -> Optional[Dict]:
         return self.inputs.prompt_params
 
     @property
-    def messages(self) -> Optional[Messages]:
+    def messages(self) -> Optional[Union[Messages, list[dict[str, str]]]]:
         """The messages as provided by the user when initializing or calling the
         Guard."""
         return self.inputs.messages
 
     @property
-    def compiled_messages(self) -> Optional[str]:
+    def compiled_messages(self) -> Optional[list[dict[str, str]]]:
         """The initial compiled messages that were passed to the LLM on the
         first call."""
         if self.iterations.empty():
             return None
-        initial_inputs = self.iterations.first.inputs
-        messages: Messages = initial_inputs.messages
+        initial_inputs = self.iterations.first.inputs  # type: ignore
+        messages = initial_inputs.messages
         prompt_params = initial_inputs.prompt_params or {}
         compiled_messages = []
+        if messages is None:
+            return None
         for message in messages:
             content = message["content"].format(**prompt_params)
             if isinstance(content, (Prompt, Instructions)):
@@ -116,11 +118,11 @@ def reask_messages(self) -> Stack[Messages]:
             reasks = self.iterations.copy()
             initial_messages = reasks.first
             reasks.remove(initial_messages)  # type: ignore
-            initial_inputs = self.iterations.first.inputs
+            initial_inputs = self.iterations.first.inputs  # type: ignore
             prompt_params = initial_inputs.prompt_params or {}
             compiled_reasks = []
             for reask in reasks:
-                messages: Messages = reask.inputs.messages
+                messages = reask.inputs.messages
 
                 if messages is None:
                     compiled_reasks.append(None)
@@ -136,7 +138,7 @@ def reask_messages(self) -> Stack[Messages]:
                                 "content": content,
                             }
                         )
-                compiled_reasks.append(compiled_messages)
+                    compiled_reasks.append(compiled_messages)
             return Stack(*compiled_reasks)
 
         return Stack()

guardrails/classes/history/iteration.py

@@ -12,7 +12,7 @@
 from guardrails.classes.history.outputs import Outputs
 from guardrails.classes.generic.arbitrary_model import ArbitraryModel
 from guardrails.logger import get_scope_handler
-from guardrails.prompt.prompt import Prompt
+from guardrails.prompt import Prompt, Instructions
 from guardrails.classes.validation.validator_logs import ValidatorLogs
 from guardrails.actions.reask import ReAsk
 from guardrails.classes.validation.validation_result import ErrorSpan
@@ -189,7 +189,7 @@ def status(self) -> str:
     @property
     def rich_group(self) -> Group:
         def create_messages_table(
-            messages: Optional[List[Dict[str, Prompt]]],
+            messages: Optional[List[Dict[str, Union[str, Prompt, Instructions]]]],
         ) -> Union[str, Table]:
             if messages is None:
                 return "No messages."
@@ -198,11 +198,11 @@ def create_messages_table(
             table.add_column("Content")
 
             for msg in messages:
-                table.add_row(str(msg["role"]), msg["content"])
+                table.add_row(str(msg["role"]), msg["content"])  # type: ignore
 
             return table
 
-        table = create_messages_table(self.inputs.messages)
+        table = create_messages_table(self.inputs.messages)  # type: ignore
 
         return Group(
             Panel(table, title="Messages", style="on #E7DFEB"),

guardrails/formatters/json_formatter.py

@@ -106,15 +106,15 @@ def fn(
                 **kwargs,
             ) -> str:
                 prompt = ""
-                for msg in messages:
+                for msg in messages:  # type: ignore
                     prompt += msg["content"]
 
                 return json.dumps(
                     Jsonformer(
                         model=model.model,
                         tokenizer=model.tokenizer,
                         json_schema=self.output_schema,
-                        prompt=prompt
+                        prompt=prompt,
                     )()
                 )
 
@@ -132,15 +132,15 @@ def fn(
                 **kwargs,
             ) -> str:
                 prompt = ""
-                for msg in messages:
+                for msg in messages:  # type: ignore
                     prompt += msg["content"]
-            
+
                 return json.dumps(
                     Jsonformer(
                         model=model,
                         tokenizer=tokenizer,
                         json_schema=self.output_schema,
-                        prompt=prompt
+                        prompt=prompt,
                     )()
                 )
 

guardrails/llm_providers.py

@@ -32,13 +32,16 @@
 
 
 # todo fix circular import
-def messages_string(messages: MessageHistory) -> str:
+def messages_string(
+    messages: Union[list[dict[str, Union[str, Prompt, Instructions]]], MessageHistory],
+) -> str:
     messages_copy = ""
     for msg in messages:
         content = (
-            msg["content"].source
+            msg["content"].source  # type: ignore
             if isinstance(msg["content"], Prompt)
-            else msg["content"]
+            or isinstance(msg["content"], Instructions)  # type: ignore
+            else msg["content"]  # type: ignore
         )
         messages_copy += content
     return messages_copy
@@ -269,7 +272,9 @@ def _invoke_llm(
         self,
         model_generate: Any,
         *args,
-        messages: list[dict[str, Union[str, Prompt, Instructions]]],
+        messages: Union[
+            list[dict[str, Union[str, Prompt, Instructions]]], MessageHistory
+        ],
         **kwargs,
     ) -> LLMResponse:
         try:

guardrails/prompt/messages.py

@@ -2,9 +2,9 @@
 
 import re
 from string import Template
-from typing import Dict, List, Optional
-
+from typing import Dict, List, Optional, Union
 
+from guardrails.prompt import Prompt, Instructions
 from guardrails.classes.templating.namespace_template import NamespaceTemplate
 from guardrails.utils.constants import constants
 from guardrails.utils.templating_utils import get_template_variables
@@ -13,7 +13,7 @@
 class Messages:
     def __init__(
         self,
-        source: List[Dict[str, str]],
+        source: List[Dict[str, Union[str, Prompt, Instructions]]],
         output_schema: Optional[str] = None,
         *,
         xml_output_schema: Optional[str] = None,
@@ -71,9 +71,7 @@ def format(
             filtered_kwargs = {k: v for k, v in kwargs.items() if k in vars}
 
             # Return another instance of the class with the formatted message.
-            formatted_message = Template(msg_str).safe_substitute(
-                **filtered_kwargs
-            )
+            formatted_message = Template(msg_str).safe_substitute(**filtered_kwargs)
             formatted_messages.append(
                 {"role": message["role"], "content": formatted_message}
             )

guardrails/run/runner.py

@@ -16,7 +16,8 @@
 )
 from guardrails.logger import set_scope
 from guardrails.prompt import Prompt
-from guardrails.run.utils import messages_source, messages_string
+from guardrails.prompt.messages import Messages
+from guardrails.run.utils import messages_source
 from guardrails.schema.rail_schema import json_schema_to_rail_output
 from guardrails.schema.validator import schema_validation
 from guardrails.hub_telemetry.hub_tracing import trace
@@ -35,6 +36,7 @@
 from guardrails.actions.reask import NonParseableReAsk, ReAsk, introspect
 from guardrails.telemetry import trace_call, trace_step
 
+
 class Runner:
     """Runner class that calls an LLM API with a prompt, and performs input and
     output validation.
@@ -293,9 +295,11 @@ def validate_messages(
                 else msg["content"]
             )
             inputs = Inputs(
-                        llm_output=content,
-                    )
-            iteration = Iteration(call_id=call_log.id, index=attempt_number, inputs=inputs)
+                llm_output=content,
+            )
+            iteration = Iteration(
+                call_id=call_log.id, index=attempt_number, inputs=inputs
+            )
             call_log.iterations.insert(0, iteration)
             value, _metadata = validator_service.validate(
                 value=content,
@@ -309,7 +313,7 @@ def validate_messages(
             validated_msg = validator_service.post_process_validation(
                 value, attempt_number, iteration, OutputTypes.STRING
             )
-            
+
             iteration.outputs.validation_response = validated_msg
 
             if isinstance(validated_msg, ReAsk):
@@ -501,7 +505,7 @@ def prepare_to_loop(
         prompt_params: Optional[Dict] = None,
     ) -> Tuple[
         Dict[str, Any],
-        Optional[List[Dict]],
+        Optional[Union[List[Dict], Messages]],
     ]:
         """Prepare to loop again."""
         prompt_params = prompt_params or {}

guardrails/utils/misc.py

@@ -49,11 +49,11 @@ def generate_test_artifacts(
             ext = f"_reask_{i}"
 
         # Save the compiled prompt.
-        compiled_prompt = logs.inputs.prompt
+        compiled_messages = logs.inputs.messages
         with open(
             os.path.join(artifact_dir, f"compiled_prompt_{on_fail_type}{ext}.txt"), "w"
         ) as f:
-            f.write(str(compiled_prompt or ""))
+            f.write(str(compiled_messages or ""))
 
         # Save the llm output.
         llm_output = logs.raw_output