SQL Select Envelope

Caution

The db_schema argument of the constructor is passed to the LLM. This is how the LLM knows how to construct the query. If this concerns you, limit the information that you include in the schema.

class PromptEnvelope(*, llm: LLMIntegration, db_schema: str, validators: Sequence[heimdallm.bifrosts.sql.validator.ConstraintValidator])

The purpose of the prompt envelope is to wrap the untrusted input in additional context for the LLM to produce the correct output. We do not do validation in the envelope, because it is impossible to prevent prompt injection.

While not necessary to subclass, you are recommended to do so if you want to customize the envelope.

Parameters:
  • llm (LLMIntegration) – The LLM integration being sent the human input. This can be used to tweak the wrap() and unwrap() methods to account for quirks of the specific LLM.

  • db_schema (str) – The database schema of the database being queried. It is passed to the LLM so that the LLM knows how the tables and columns are connected.

  • validators (Sequence[heimdallm.bifrosts.sql.validator.ConstraintValidator]) – The validators to use to validate the output of the LLM. They aren’t used to validate here, but some of the validator’s properties are added to the envelope to help guide the LLM to produce the correct output.

property params: dict

Returns a dictionary of additional parameters to be passed to the template. Override in a subclass for complete control over values that you want in the envelope.

Returns:

The extra parameters to pass to the template.

template(env: Environment) Template

Returns the template to use for the envelope. Override in a subclass for complete customization.

Parameters:

env (Environment) – The environment to use to load the template.

Returns:

The template to use for the envelope.

Return type:

Template

unwrap(untrusted_llm_output: str) str

Unpack the SQL query from the LLM output by finding it (hopefully) among the delimiters.

Parameters:

untrusted_llm_output (str) – The output from the LLM.

Returns:

The SQL query.

Return type:

str

wrap(untrusted_input: str) str

Performs the wrapping of the untrusted input with the envelope. Not intended to be overridden, but not foribben either. Consider overriding the template() and params() properties first instead.

Parameters:

untrusted_input (str) – The untrusted input from the user.

Returns:

The wrapped input to send to the LLM.

Return type:

str