125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
"""Class for representing a prompt entry."""
|
|
|
|
import re
|
|
from string import Template
|
|
from typing import List, Optional
|
|
|
|
import regex
|
|
|
|
from guardrails.classes.templating.namespace_template import NamespaceTemplate
|
|
from guardrails.utils.constants import constants
|
|
from guardrails.utils.templating_utils import get_template_variables
|
|
|
|
|
|
class BasePrompt:
|
|
"""Base class for representing an LLM prompt."""
|
|
|
|
def __init__(
|
|
self,
|
|
source: str,
|
|
output_schema: Optional[str] = None,
|
|
*,
|
|
xml_output_schema: Optional[str] = None,
|
|
):
|
|
"""Initialize and substitute constants in the prompt."""
|
|
self._source = source
|
|
self.format_instructions_start = self.get_format_instructions_idx(source)
|
|
|
|
# FIXME: Why is this happening on init instead of on format?
|
|
# Substitute constants in the prompt.
|
|
source = self.substitute_constants(source)
|
|
|
|
# FIXME: Why is this happening on init instead of on format?
|
|
# If an output schema is provided, substitute it in the prompt.
|
|
if output_schema or xml_output_schema:
|
|
self.source = Template(source).safe_substitute(
|
|
output_schema=output_schema, xml_output_schema=xml_output_schema
|
|
)
|
|
else:
|
|
self.source = source
|
|
|
|
def __repr__(self) -> str:
|
|
# Truncate the prompt to 50 characters and add ellipsis if it's longer.
|
|
truncated_prompt = self.source[:50]
|
|
if len(self.source) > 50:
|
|
truncated_prompt += "..."
|
|
return f"Prompt({truncated_prompt})"
|
|
|
|
def __str__(self) -> str:
|
|
return self.source
|
|
|
|
@property
|
|
def variable_names(self):
|
|
return get_template_variables(self.source)
|
|
|
|
@property
|
|
def format_instructions(self):
|
|
return self.source[self.format_instructions_start :]
|
|
|
|
def substitute_constants(self, text: str) -> str:
|
|
"""Substitute constants in the prompt."""
|
|
# Substitute constants by reading the constants file.
|
|
# Regex to extract all occurrences of ${gr.<constant_name>}
|
|
matches = re.findall(r"\${gr\.(\w+)}", text)
|
|
|
|
# Substitute all occurrences of ${gr.<constant_name>}
|
|
# with the value of the constant.
|
|
for match in matches:
|
|
template = NamespaceTemplate(text)
|
|
mapping = {f"gr.{match}": constants[match]}
|
|
text = template.safe_substitute(**mapping)
|
|
|
|
return text
|
|
|
|
def get_prompt_variables(self) -> List[str]:
|
|
return self.variable_names
|
|
|
|
def format(self, **kwargs) -> "BasePrompt":
|
|
raise NotImplementedError("Subclasses must implement this method.")
|
|
|
|
def make_vars_optional(self):
|
|
"""Make all variables in the prompt optional."""
|
|
for var in self.variable_names:
|
|
self.source = self.source.replace(f"{{{var}}}", f"{{{var}:}}")
|
|
|
|
def get_format_instructions_idx(self, text: str) -> Optional[int]:
|
|
"""Get the index of the first format instruction in the prompt.
|
|
|
|
It checks to see where the first instance of any constant is in the text.
|
|
Everything from then on is considered to be a format instruction.
|
|
|
|
Returns:
|
|
The index of the first format instruction in the prompt.
|
|
"""
|
|
# TODO(shreya): Optionally add support for special character demarcation.
|
|
|
|
# Regex to extract first occurrence of ${gr.<constant_name>}
|
|
|
|
matches = re.finditer(r"\${gr\.(\w+)}", text)
|
|
|
|
earliest_match_idx = None
|
|
earliest_match = None
|
|
|
|
# Find the earliest match where the match belongs to a constant.
|
|
for match in matches:
|
|
if match.group(1) in constants:
|
|
if earliest_match_idx is None or earliest_match_idx > match.start():
|
|
earliest_match_idx = match.start()
|
|
earliest_match = match
|
|
|
|
if earliest_match_idx is None:
|
|
return 0
|
|
|
|
if earliest_match is None:
|
|
return None
|
|
return earliest_match.start()
|
|
|
|
def escape(self) -> str:
|
|
"""Escape single curly braces into double curly braces."""
|
|
start_replaced = regex.sub(r"(?<!\$){", "{{", self.source)
|
|
# This variable length negative lookbehind is why we need `regex` over `re`
|
|
return regex.sub(r"(?<!\${.*)}", "}}", start_replaced)
|
|
|
|
def _to_request(self) -> str:
|
|
return self.source
|