diff --git a/CHANGELOG.md b/CHANGELOG.md index f3a67a9..6adde2a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## 1.0.1 - 9/12/23 + +- Simplify docs +- Add missing Postgres api docs +- Generalizing `validation_only` to the `Bifrost` superclass + ## 1.0.0 - 9/11/23 - Subquery support diff --git a/README.md b/README.md index 2d52d11..69470fb 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,11 @@ # HeimdaLLM -> Heimdall, the watchman of the gods, dwelt at its entrance, where he guarded Bifrost, -> the shimmering path connecting the realms. +Pronounced `[ˈhaΙͺm.dΙ”l.Ι™m]` or _HEIM-dall-EM_ + +HeimdaLLM is a robust static analysis framework for validating that LLM-generated +structured output is safe. It currently supports SQL. + +In simple terms, it helps makes sure that AI won't wreck your systems. [![Heimdall](https://raw.githubusercontent.com/amoffat/HeimdaLLM/main/docs/source/images/heimdall.png)](https://heimdallm.ai) [![Build status](https://github.com/amoffat/HeimdaLLM/actions/workflows/main.yml/badge.svg?branch=main)](https://github.com/amoffat/HeimdaLLM/actions) @@ -12,77 +16,63 @@ [![License: AGPL v3](https://img.shields.io/badge/License-AGPL_v3-blue.svg)](https://www.gnu.org/licenses/agpl-3.0) [![Coverage Status](https://coveralls.io/repos/github/amoffat/HeimdaLLM/badge.svg?branch=dev)](https://coveralls.io/github/amoffat/HeimdaLLM?branch=dev) -HeimdaLLM safely bridges the gap between untrusted human input and trusted -machine-readable output by augmenting LLMs with a robust validation framework. This -enables you externalize LLM technology to your users, so that you can do things like -execute trusted SQL queries from their untrusted input. +Consider the following natural-language database query: -To accomplish this, HeimdaLLM introduces a new technology, the 🌈✨ -[Bifrost](https://docs.heimdallm.ai/en/latest/bifrost.html), composed of 4 parts: an LLM -prompt envelope, an LLM integration, a grammar, and a constraint validator. These 4 -components operate as a single unitβ€”a Bifrostβ€”which is capable of translating untrusted -human input into trusted machine output. +``` +how much have i spent renting movies, broken down by month? +``` -✨ **This allows you to perform magic** ✨ +From this query (and a little bit of context), an LLM can produce the following SQL +query: + +```sql +SELECT + strftime('%Y-%m', payment.payment_date) AS month, + SUM(payment.amount) AS total_amount +FROM payment +JOIN rental ON payment.rental_id=rental.rental_id +JOIN customer ON payment.customer_id=customer.customer_id +WHERE customer.customer_id=:customer_id +GROUP BY month +LIMIT 10; +``` -Imagine giving your users natural language access to their data in your database, -without having to worry about dangerous queries. This is an actual query on the [Sakila -Sample -Database](https://www.kaggle.com/datasets/atanaskanev/sqlite-sakila-sample-database): +But how can you ensure the LLM-generated query is safe and that it only accesses +authorized data? -```python -traverse("Show me the movies I rented the longest, and the number of days I had them for.") -``` +HeimdaLLM performs static analysis on the generated SQL to ensure that only certain +columns, tables, and functions are used. It also automatically edits the query to add a +`LIMIT` and to remove forbidden columns. Lastly, it ensures that there is a column +constraint that would restrict the results to only the user's data. + +It does all of this locally, without AI, using good ol' fashioned grammars and parsers: ``` βœ… Ensuring SELECT statement... βœ… Resolving column and table aliases... βœ… Allowlisting selectable columns... - βœ… Removing 4 forbidden columns... + βœ… Removing 2 forbidden columns... βœ… Ensuring correct row LIMIT exists... - βœ… Lowering row LIMIT to 5... + βœ… Lowering row LIMIT to 10... βœ… Checking JOINed tables and conditions... βœ… Checking required WHERE conditions... βœ… Ensuring query is constrained to requester's identity... βœ… Allowlisting SQL functions... + βœ… strftime + βœ… SUM ``` -| Title | Rental Date | Return Date | Rental Days | -| --------------- | ----------------------- | ----------------------- | ----------- | -| OUTLAW HANKY | 2005-08-19 05:48:12.000 | 2005-08-28 10:10:12.000 | 9.181944 | -| BOULEVARD MOB | 2005-08-19 07:06:51.000 | 2005-08-28 10:35:51.000 | 9.145139 | -| MINDS TRUMAN | 2005-08-02 17:42:49.000 | 2005-08-11 18:14:49.000 | 9.022222 | -| AMERICAN CIRCUS | 2005-07-12 16:37:55.000 | 2005-07-21 16:04:55.000 | 8.977083 | -| LADY STAGE | 2005-07-28 10:07:04.000 | 2005-08-06 08:16:04.000 | 8.922917 | - -You can safely run this example here: - -[![Open in GitHub Codespaces](https://img.shields.io/badge/Open%20in-Codespaces-purple.svg)](https://github.com/codespaces/new?hide_repo_select=true&ref=main&repo=656570421) - -or [view the read-only notebook](./notebooks/demo.ipynb) - -# πŸ“‹ Explanation - -So, what is actually happening above? - -1. Unsafe free-form input is provided, presumably from some front end user interface. -1. That unsafe input is wrapped in a prompt envelope, producing a prompt with additional - context to help an LLM produce a correct query. -1. The unsafe prompt is sent to an LLM of your choice, which then produces an unsafe - SQL query. -1. The LLM response is parsed by a strict grammar which defines only the SQL features - that are allowed. -1. If parsing succeeds, we know at the very least we're dealing with a valid SQL query - albeit an untrusted one. -1. Different features of the parsed query are extracted for validation. -1. A soft validation pass is performed on the extracted features, and we potentially - modify the query to be compliant, for example, to add a `LIMIT` clause, or to remove - disallowed columns. -1. A hard validation pass is performed with your custom constraints to ensure that the - query is only accessing allowed tables, columns, and functions, while containing - required conditions. -1. If validation succeeds, the resulting SQL query can then be sent to the database. -1. If validation fails, you'll see a helpful exception explaining exactly why. +The validated query can then be executed: + +| month | total_amount | +| ------- | ------------ | +| 2005-05 | 4.99 | +| 2005-06 | 22.95 | +| 2005-07 | 100.78 | +| 2005-08 | 87.82 | + +Want to get started quickly? Go +[here](https://docs.heimdallm.ai/en/latest/quickstart/index.html). # πŸ₯½ Safety @@ -94,7 +84,7 @@ me](https://github.com/sponsors/amoffat) or [inquire about interest in a commerc license](https://forms.gle/frEPeeJx81Cmwva78). To understand some of the potential vulnerabilities, take a look at the [attack -surface](https://docs.heimdallm.ai/en/latest/attack_surface.html) to see the risks and +surface](https://docs.heimdallm.ai/en/latest/attack-surface.html) to see the risks and the mitigations. # πŸ“š Database support diff --git a/docs/source/api/abc/index.rst b/docs/source/api/abc/index.rst index 93c5a14..20bd003 100644 --- a/docs/source/api/abc/index.rst +++ b/docs/source/api/abc/index.rst @@ -9,7 +9,7 @@ intended for direct use. bifrost envelope validator - llm_integration + llm-integration context sql/index \ No newline at end of file diff --git a/docs/source/api/abc/llm_integration.rst b/docs/source/api/abc/llm-integration.rst similarity index 100% rename from docs/source/api/abc/llm_integration.rst rename to docs/source/api/abc/llm-integration.rst diff --git a/docs/source/api/bifrosts/index.rst b/docs/source/api/bifrosts/index.rst index c582e1b..a7943f8 100644 --- a/docs/source/api/bifrosts/index.rst +++ b/docs/source/api/bifrosts/index.rst @@ -1,8 +1,8 @@ Bifrosts ======== -:doc:`Bifrosts ` are the fundamental unit of translating untrusted input into -trusted output. This document will expand as we add more Bifrosts. +:doc:`Bifrosts ` are the fundamental unit of translating +untrusted input into trusted output. This document will expand as we add more Bifrosts. .. toctree:: diff --git a/docs/source/api/bifrosts/sql/index.rst b/docs/source/api/bifrosts/sql/index.rst index 16e5e20..3299a49 100644 --- a/docs/source/api/bifrosts/sql/index.rst +++ b/docs/source/api/bifrosts/sql/index.rst @@ -6,9 +6,11 @@ database, please participate in `this poll. `_ .. toctree:: + :maxdepth: 3 sqlite/index mysql/index + postgres/index exceptions common \ No newline at end of file diff --git a/docs/source/api/bifrosts/sql/postgres/index.rst b/docs/source/api/bifrosts/sql/postgres/index.rst new file mode 100644 index 0000000..762d12a --- /dev/null +++ b/docs/source/api/bifrosts/sql/postgres/index.rst @@ -0,0 +1,6 @@ +Postgres +======== + +.. toctree:: + + select/index \ No newline at end of file diff --git a/docs/source/api/bifrosts/sql/postgres/select/bifrost.rst b/docs/source/api/bifrosts/sql/postgres/select/bifrost.rst new file mode 100644 index 0000000..88876c4 --- /dev/null +++ b/docs/source/api/bifrosts/sql/postgres/select/bifrost.rst @@ -0,0 +1,13 @@ +SQL Select Bifrost +================== + +The SQL Select Bifrost produces a trusted SQL Select statement. It uses the following +components: + +* :class:`SQLPromptEnvelope ` +* :class:`SQLConstraintValidator ` +* `Grammar `_ + +.. autoclass:: heimdallm.bifrosts.sql.postgres.select.bifrost.Bifrost + :members: + :inherited-members: \ No newline at end of file diff --git a/docs/source/api/bifrosts/sql/postgres/select/envelope.rst b/docs/source/api/bifrosts/sql/postgres/select/envelope.rst new file mode 100644 index 0000000..0cce825 --- /dev/null +++ b/docs/source/api/bifrosts/sql/postgres/select/envelope.rst @@ -0,0 +1,12 @@ +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. + +.. autoclass:: heimdallm.bifrosts.sql.postgres.select.envelope.PromptEnvelope + :members: + :inherited-members: \ No newline at end of file diff --git a/docs/source/api/bifrosts/sql/postgres/select/index.rst b/docs/source/api/bifrosts/sql/postgres/select/index.rst new file mode 100644 index 0000000..2bbd493 --- /dev/null +++ b/docs/source/api/bifrosts/sql/postgres/select/index.rst @@ -0,0 +1,8 @@ +Select +====== + +.. toctree:: + + bifrost + envelope + validator \ No newline at end of file diff --git a/docs/source/api/bifrosts/sql/postgres/select/validator.rst b/docs/source/api/bifrosts/sql/postgres/select/validator.rst new file mode 100644 index 0000000..df42f8c --- /dev/null +++ b/docs/source/api/bifrosts/sql/postgres/select/validator.rst @@ -0,0 +1,8 @@ +SQL Select Validator +==================== + +.. autoclass:: heimdallm.bifrosts.sql.postgres.select.validator.ConstraintValidator + :members: + :inherited-members: + + \ No newline at end of file diff --git a/docs/source/api/index.rst b/docs/source/api/index.rst index 90be0c6..a3aba2e 100644 --- a/docs/source/api/index.rst +++ b/docs/source/api/index.rst @@ -7,5 +7,5 @@ :maxdepth: 4 bifrosts/index - llm_providers/index + llm-providers/index abc/index \ No newline at end of file diff --git a/docs/source/api/llm_providers/index.rst b/docs/source/api/llm-providers/index.rst similarity index 100% rename from docs/source/api/llm_providers/index.rst rename to docs/source/api/llm-providers/index.rst diff --git a/docs/source/api/llm_providers/providers/openai.rst b/docs/source/api/llm-providers/providers/openai.rst similarity index 100% rename from docs/source/api/llm_providers/providers/openai.rst rename to docs/source/api/llm-providers/providers/openai.rst diff --git a/docs/source/bifrost.rst b/docs/source/architecture/bifrost.rst similarity index 91% rename from docs/source/bifrost.rst rename to docs/source/architecture/bifrost.rst index 95cf213..524b6ae 100644 --- a/docs/source/bifrost.rst +++ b/docs/source/architecture/bifrost.rst @@ -1,6 +1,11 @@ 🌈 Bifrost ========== +.. DANGER:: + + Constructing Bifrosts manually is an advanced topic. Most of the time, you want to + use :meth:`Bifrost.validation_only `. + The Bifrost is the key technology that enables the translation of untrusted human input into trusted machine-readable input. It is composed of 4 parts: @@ -42,7 +47,7 @@ out the structured data from the delimiters that you instructed the LLM to use. ********************** The :term:`LLM ` itself is the brains of the Bifrost. We view it as a black box -with a :doc:`well-defined interface `. Because of this, +with a :doc:`well-defined interface `. Because of this, HeimdaLLM aims to make it easy to swap out LLMs in your Bifrost, so that as LLM capabilities and prices change, your system can adapt to use them with minimal effort. @@ -51,7 +56,7 @@ Current LLM integrations: .. toctree:: :maxdepth: 2 - api/llm_providers/index + /api/llm-providers/index πŸ“œ The grammar ************** diff --git a/docs/source/architecture/index.rst b/docs/source/architecture/index.rst new file mode 100644 index 0000000..c31589f --- /dev/null +++ b/docs/source/architecture/index.rst @@ -0,0 +1,7 @@ +πŸ“ Architecture +=============== + +.. toctree:: + :maxdepth: 4 + + bifrost \ No newline at end of file diff --git a/docs/source/attack-surface/index.rst b/docs/source/attack-surface/index.rst new file mode 100644 index 0000000..c792007 --- /dev/null +++ b/docs/source/attack-surface/index.rst @@ -0,0 +1,14 @@ +πŸ›‘οΈ Attack Surface +================= + +:term:`LLMs ` are vulnerable to :term:`prompt injection` attacks, which can be used +to construct responses that are dangerous to the system. This is the primary reason that +LLMs have not seen widespread adoption as :term:`externalized ` products. + +Prompt injection can have different consequences for different types of structured +outputs. + +.. toctree:: + :glob: + + sql \ No newline at end of file diff --git a/docs/source/attack_surface/sql.rst b/docs/source/attack-surface/sql.rst similarity index 100% rename from docs/source/attack_surface/sql.rst rename to docs/source/attack-surface/sql.rst diff --git a/docs/source/attack_surface.rst b/docs/source/attack_surface.rst deleted file mode 100644 index 6c33abb..0000000 --- a/docs/source/attack_surface.rst +++ /dev/null @@ -1,14 +0,0 @@ -πŸ›‘οΈ Attack Surface -================= - -:term:`LLMs ` are vulnerable to :term:`prompt injection` attacks, which can be used to construct -responses that are dangerous to the system. This is the primary reason that LLMs have -not seen widespread adoption as :term:`externalized ` products. - -Prompt injection can have different consequences within different contexts, and this -page will reflect how different :doc:`Bifrosts ` manage the attack surface. - -.. toctree:: - :glob: - - attack_surface/* \ No newline at end of file diff --git a/docs/source/blog/posts/safe-sql-execution.rst b/docs/source/blog/posts/safe-sql-execution.rst index 36e857a..57dd2c6 100644 --- a/docs/source/blog/posts/safe-sql-execution.rst +++ b/docs/source/blog/posts/safe-sql-execution.rst @@ -359,23 +359,23 @@ authoritative id. -🧠 Constraint validation ------------------------- +🧠 Static analysis +------------------ .. figure:: /images/smiley.jpg - "Where did it come from? What's the access?" + "What's the access?" -Constraint validation uses a real grammar to parse SQL queries into an AST. Static -analysis can then be performed on this parse tree to determine which tables and columns -are being used, how they're being used, if required conditions are present, and a range -of other features. +Static analysis uses a real grammar to parse SQL queries into an AST, which can then be +analyzed to determine which tables and columns are being used, how they're being used, +if required conditions are present, and a range of other features. -Additionally, these frameworks may automatically add nodes or replace nodes on the AST -to help ensure the SQL query conforms to constraint validation. In other words, the -query may be automatically edited to be compliant. Examples of this are to ensure a -correct ``LIMIT`` on the query, or remove a forbidden column from the ``SELECT``. +Additionally, these static analysis frameworks may automatically add nodes or replace +nodes on the AST to help ensure the SQL query conforms to constraints. In +other words, the query may be automatically edited to be compliant. Examples of this are +to ensure a correct ``LIMIT`` on the query, or remove a forbidden column from the +``SELECT``. These frameworks can be treated as denylists or allowlists. You can list which tables, columns, joins, and functions are allowed, or which are denied. This allows for a higher @@ -423,10 +423,10 @@ are playing an increasing role in the future of UI and UX, and relational databa not going away any time soon. For them to work together effectively, tooling needs to bridge the gap to make them safer. -The most promising solutions are cloned databases and constraint validators, because -they are theoretically complete solutions that can offer the highest levels of security. -They vary primarily in their complexity and flexibility: cloned databases views are a -high-complexity allowlist, while constraint validators are a low-complexity allowlist or +The most promising solutions are cloned databases and static analysis, because they are +theoretically complete solutions that can offer the highest levels of security. They +vary primarily in their complexity and flexibility: cloned databases views are a +high-complexity allowlist, while static analysis is a low-complexity allowlist or denylist. Other, non-complete solutions should not be considered if you value the safety of your diff --git a/docs/source/faq.rst b/docs/source/faq.rst index 3110937..2b6315d 100644 --- a/docs/source/faq.rst +++ b/docs/source/faq.rst @@ -25,11 +25,19 @@ What databases are supported? * Sqlite * MySQL +* Postgres There is rapid development for the other top relational SQL databases. To help us prioritize, please `vote here `_ on which database you would like to see supported: +Do HeimdaLLM use an LLM? +************************ + +For static analysis, no it does not. It uses good old fashioned grammars and parsers. +However, we do include a lightweight framework to build a complete +natural-language-to-safe-SQL workflow. See :doc:`this quickstart `. + Do I need to purchase a commercial license? ******************************************* diff --git a/docs/source/glossary.rst b/docs/source/glossary.rst index dfb7a09..85eb097 100644 --- a/docs/source/glossary.rst +++ b/docs/source/glossary.rst @@ -6,9 +6,8 @@ autofixing An alias for :term:`reconstruction`. - :doc:`Bifrost ` - An object that can traverse untrusted human input into untrusted - machine-readable input. + :doc:`Bifrost ` + An object that can convert untrusted input into trusted output. :doc:`constraint validator ` An object that is capable of analyzing an LLM's structured output to determine diff --git a/docs/source/index.rst b/docs/source/index.rst index d0b38b5..19e19cb 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -1,8 +1,14 @@ HeimdaLLM ========= - Heimdall, the watchman of the gods, dwelt at its entrance, where he guarded Bifrost, - the shimmering path connecting the realms. +.. role:: phonetic + +Pronounced :phonetic:`[ˈhaΙͺm.dΙ”l.Ι™m]` or HEIM-dall-em + +HeimdaLLM is a robust static analysis framework for validating that LLM-generated +structured output is safe. It currently supports SQL. + +In simple terms, it helps makes sure that AI won't wreck your systems. .. image:: https://raw.githubusercontent.com/amoffat/HeimdaLLM/main/docs/source/images/heimdall.png :target: https://github.com/amoffat/HeimdaLLM @@ -32,78 +38,88 @@ HeimdaLLM :target: https://github.com/amoffat/HeimdaLLM :alt: GitHub Repo stars -.. ATTENTION:: +Consider the following natural-language database query: - These docs are under active development. See an issue? Report it `here. - `__ - Want to make sure something is included? Please request it `here. - `__ +.. code-block:: text -Welcome to the HeimdaLLM documentation! + how much have i spent renting movies, broken down by month? -HeimdaLLM safely bridges the gap between untrusted human input and trusted -machine-readable output by augmenting :term:`LLMs ` with a robust validation -framework. This enables you to :term:`externalize ` LLM technology to -your users, so that you can do things like execute trusted SQL queries from their -untrusted input. +From this query (and a little bit of context), an LLM can produce the following SQL +query: -Imagine giving your users natural language access to their data in your database, -without having to worry about dangerous queries. +.. code-block:: sql -.. code-block:: python + SELECT + strftime('%Y-%m', payment.payment_date) AS month, + SUM(payment.amount) AS total_amount + FROM payment + JOIN rental ON payment.rental_id=rental.rental_id + JOIN customer ON payment.customer_id=customer.customer_id + WHERE customer.customer_id=:customer_id + GROUP BY month + LIMIT 10; - traverse("Show me the movies I rented the longest, and the number of days I had them for.") +But how can you ensure the LLM-generated query is safe and that it only accesses +authorized data? + +HeimdaLLM performs static analysis on the generated SQL to ensure that only certain +columns, tables, and functions are used. It also automatically edits the query to add a +``LIMIT`` and to remove forbidden columns. Lastly, it ensures that there is a column +constraint that would restrict the results to only the user's data. + +It does all of this locally, without AI, using good ol' fashioned grammars and parsers: .. code-block:: text βœ… Ensuring SELECT statement... βœ… Resolving column and table aliases... βœ… Allowlisting selectable columns... - βœ… Removing 4 forbidden columns... + βœ… Removing 2 forbidden columns... βœ… Ensuring correct row LIMIT exists... - βœ… Lowering row LIMIT to 5... + βœ… Lowering row LIMIT to 10... βœ… Checking JOINed tables and conditions... βœ… Checking required WHERE conditions... βœ… Ensuring query is constrained to requester's identity... βœ… Allowlisting SQL functions... + βœ… strftime + βœ… SUM -+-----------------+------------------------+------------------------+--------------+ -| Title | Rental Date | Return Date | Rental Days | -+=================+========================+========================+==============+ -| OUTLAW HANKY | 2005-08-19 05:48:12.000| 2005-08-28 10:10:12.000| 9.181944 | -+-----------------+------------------------+------------------------+--------------+ -| BOULEVARD MOB | 2005-08-19 07:06:51.000| 2005-08-28 10:35:51.000| 9.145139 | -+-----------------+------------------------+------------------------+--------------+ -| MINDS TRUMAN | 2005-08-02 17:42:49.000| 2005-08-11 18:14:49.000| 9.022222 | -+-----------------+------------------------+------------------------+--------------+ -| AMERICAN CIRCUS | 2005-07-12 16:37:55.000| 2005-07-21 16:04:55.000| 8.977083 | -+-----------------+------------------------+------------------------+--------------+ -| LADY STAGE | 2005-07-28 10:07:04.000| 2005-08-06 08:16:04.000| 8.922917 | -+-----------------+------------------------+------------------------+--------------+ +The validated query can then be executed: -.. TIP:: ++---------+--------------+ +| month | total_amount | ++---------+--------------+ +| 2005-05 | 4.99 | ++---------+--------------+ +| 2005-06 | 22.95 | ++---------+--------------+ +| 2005-07 | 100.78 | ++---------+--------------+ +| 2005-08 | 87.82 | ++---------+--------------+ - Run this example safely in Github Codespaces |CodespacesLink|_ - -Interested in getting started quickly? Check out the :doc:`quickstart`. Otherwise, -browse the navigation on the left. +Want to get started quickly? :doc:`quickstart/index`. .. toctree:: :hidden: :glob: :maxdepth: 5 - quickstart + quickstart/index blog/index - bifrost api/index reconstruction - attack_surface - tutorials - llm_quirks + attack-surface/index + tutorials/index + llm-quirks/index glossary roadmap + architecture/index faq -.. |CodespacesLink| image:: https://img.shields.io/badge/Open%20in-Codespaces-purple.svg -.. _CodespacesLink: https://github.com/codespaces/new?hide_repo_select=true&ref=main&repo=656570421 \ No newline at end of file +.. ATTENTION:: + + These docs are under active development. See an issue? Report it `here. + `__ + Want to make sure something is included? Please request it `here. + `__ \ No newline at end of file diff --git a/docs/source/llm_quirks.rst b/docs/source/llm-quirks/index.rst similarity index 92% rename from docs/source/llm_quirks.rst rename to docs/source/llm-quirks/index.rst index 8215d97..cfc4337 100644 --- a/docs/source/llm_quirks.rst +++ b/docs/source/llm-quirks/index.rst @@ -7,4 +7,4 @@ quirks discovered while building Bifrosts. .. toctree:: :glob: - llm_quirks/* \ No newline at end of file + sql \ No newline at end of file diff --git a/docs/source/llm_quirks/sql.rst b/docs/source/llm-quirks/sql.rst similarity index 91% rename from docs/source/llm_quirks/sql.rst rename to docs/source/llm-quirks/sql.rst index f2d69ee..b1df98c 100644 --- a/docs/source/llm_quirks/sql.rst +++ b/docs/source/llm-quirks/sql.rst @@ -23,9 +23,7 @@ Arbitrary limits **************** An LLM will sometimes include a ``LIMIT`` on a query, and other times it won't. This can -have performance impacts and is generally undesirable. Originally, we attempted to coax -the LLM into including a ``LIMIT`` on every query, by using the :ref:`prompt envelope -`, but this proved to be unreliable. +have performance impacts and is generally undesirable. Instead, our constraint validator is capable of rebuilding the query to include a ``LIMIT`` if one is not present, or to adjust an existing ``LIMIT``, if the current one diff --git a/docs/source/quickstart/index.rst b/docs/source/quickstart/index.rst new file mode 100644 index 0000000..a8f0ce7 --- /dev/null +++ b/docs/source/quickstart/index.rst @@ -0,0 +1,15 @@ +πŸš€ Quickstart +============= + +There are two primary ways of using HeimdaLLM, depending on your unique circumstances. +The first is if you are doing the LLM integration yourself, and you only need to +validate the structured output. The second is if you want to use HeimdaLLM to create +a complete workflow to convert natural language to validated, structured output, via an +LLM. + +.. toctree:: + :glob: + :maxdepth: 5 + + validation + llm \ No newline at end of file diff --git a/docs/source/quickstart.rst b/docs/source/quickstart/llm.rst similarity index 92% rename from docs/source/quickstart.rst rename to docs/source/quickstart/llm.rst index 0fd4d00..65cc9b6 100644 --- a/docs/source/quickstart.rst +++ b/docs/source/quickstart/llm.rst @@ -1,14 +1,17 @@ -πŸš€ Quickstart -============= +LLM integration +=============== -This quickstart will walk you through setting up a SQL Bifrost with OpenAI's LLM. The +This quickstart will walk you through setting up HeimdaLLM with OpenAI's LLM. The end result is a function that takes natural language input and returns a trusted SQL ``SELECT`` query, constrained to your requirements. +If you wish to also use HeimdaLLM only for validating an existing SQL query, +see :doc:`this quickstart `. + .. TIP:: You can also find this quickstart code in a Jupyter Notebook `here. - `_ + `_ First let's set up our imports. @@ -110,7 +113,7 @@ that the LLM is guided to produce a correct response. validators=[validator], ) -Now we can bring everything together into a :doc:`/bifrost` +Now we can bring everything together into a :doc:`/architecture/bifrost` .. code-block:: python diff --git a/docs/source/quickstart/validation.rst b/docs/source/quickstart/validation.rst new file mode 100644 index 0000000..592fb6b --- /dev/null +++ b/docs/source/quickstart/validation.rst @@ -0,0 +1,97 @@ +Validation only +=============== + +This quickstart will walk you through setting up a HeimdaLLM to validate a SQL query +from an untrusted source, presumably from an LLM. The end result is a function that +takes natural language input and returns a trusted SQL ``SELECT`` query, constrained to +your requirements. + +If you wish to also use HeimdaLLM to talk to an LLM for you to generate the SQL query, +see :doc:`this quickstart `. + +.. TIP:: + + You can also find this quickstart code in a Jupyter Notebook `here. + `_ + + +First let's set up our imports. + +.. code-block:: python + + import logging + from typing import Sequence + + import structlog + + from heimdallm.bifrosts.sql.sqlite.select.bifrost import Bifrost + from heimdallm.bifrosts.sql.sqlite.select.validator import ConstraintValidator + from heimdallm.bifrosts.sql.common import FqColumn, JoinCondition, ParameterizedConstraint + + logging.basicConfig(level=logging.ERROR) + structlog.configure(logger_factory=structlog.stdlib.LoggerFactory()) + + +Let's define our constraint validator(s). These are used to constrain the SQL query so +that it only has access to tables and columns that you allow. For more information on +the methods that you can override in the derived class, look :doc:`here. +` + +.. code-block:: python + + class CustomerConstraintValidator(ConstraintValidator): + def requester_identities(self) -> Sequence[ParameterizedConstraint]: + return [ + ParameterizedConstraint( + column="customer.customer_id", + placeholder="customer_id", + ), + ] + + def parameterized_constraints(self) -> Sequence[ParameterizedConstraint]: + return [] + + def select_column_allowed(self, column: FqColumn) -> bool: + return True + + def allowed_joins(self) -> Sequence[JoinCondition]: + return [ + JoinCondition("payment.rental_id", "rental.rental_id"), + JoinCondition( + "customer.customer_id", + "payment.customer_id", + identity="customer_id", + ), + ] + + def max_limit(self) -> int | None: + return 10 + + + validator = CustomerConstraintValidator() + +Now let's construct a Bifrost that validates SQL: + +.. code-block:: python + + bifrost = Bifrost.validation_only( + constraint_validators=[validator], + ) + +You can now validate constraints on SQL: + +.. code-block:: python + + query = """ + SELECT + strftime('%Y-%m', payment.payment_date) AS month, + SUM(payment.amount) AS total_amount + FROM payment + JOIN rental ON payment.rental_id=rental.rental_id + JOIN customer ON payment.customer_id=customer.customer_id + WHERE customer.customer_id=:customer_id + GROUP BY month + """ + + query = bifrost.traverse(query) + print(query) \ No newline at end of file diff --git a/docs/source/reconstruction.rst b/docs/source/reconstruction.rst index fc5976b..7c3188e 100644 --- a/docs/source/reconstruction.rst +++ b/docs/source/reconstruction.rst @@ -3,7 +3,7 @@ A Bifrost is capable of reconstructing the structured output from an LLM, to help ensure that it will pass validation. In the code and docs, this is called "reconstruction" or -"autofixing." Alghtough reconstruction is a general HeimdaLLM concept, This page will +"autofixing." Although reconstruction is a general HeimdaLLM concept, This page will explain it concretely in the context of the :class:`SQL SELECT Bifrost. ` diff --git a/docs/source/roadmap.rst b/docs/source/roadmap.rst index 16d2876..c13bbd0 100644 --- a/docs/source/roadmap.rst +++ b/docs/source/roadmap.rst @@ -8,26 +8,6 @@ Below is a very high-level roadmap for HeimdaLLM. Want to know when these features get implemented? Add your email to `this form `_ and I will contact you. -More LLM integrations -********************* - -Currently we support OpenAI, but I intend to add support for all major LLM API services, -and private LLMs, as they become more capable. - -More databases -************** - -I will be adding support for more SQL-based databases: - -* SQL Server -* Oracle -* Snowflake - -.. NOTE:: - - Want us to prioritize a specific database? Let us know by `voting here. - `_ - More SQL statement types ************************ @@ -54,6 +34,28 @@ constraint validators are defined by subclassing a Python class. A future implem could be language agnostic by providing an api and a JSON or YAML spec for constraining LLM output. + +More LLM integrations +********************* + +Currently we support OpenAI, but I intend to add support for all major LLM API services, +and private LLMs, as they become more capable. + +More databases +************** + +I will be adding support for more SQL-based databases: + +* SQL Server +* Oracle +* Snowflake + +.. NOTE:: + + Want us to prioritize a specific database? Let us know by `voting here. + `_ + + More Bifrosts ************* diff --git a/docs/source/tutorials.rst b/docs/source/tutorials/index.rst similarity index 100% rename from docs/source/tutorials.rst rename to docs/source/tutorials/index.rst diff --git a/heimdallm/bifrost.py b/heimdallm/bifrost.py index 277ce3f..aa80947 100644 --- a/heimdallm/bifrost.py +++ b/heimdallm/bifrost.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Callable, Sequence +from typing import TYPE_CHECKING, Any, Callable, Sequence, Union import structlog from lark import Lark, ParseTree @@ -49,6 +49,20 @@ def __init__( self.constraint_validators = constraint_validators self.ctx = TraverseContext() + @classmethod + def validation_only( + cls, + constraint_validators: Union[Any, Sequence[Any]], + ): + """A convenience method for doing just static analysis. This creates a + Bifrost that assumes its untrusted input was already produced by an LLM, so we + only need to parse and validate it. + + :param constraint_validators: A constraint validator or sequence of constraint + validators to run on the untrusted input. + """ + raise NotImplementedError + def traverse( self, untrusted_human_input: str, diff --git a/heimdallm/bifrosts/sql/bifrost.py b/heimdallm/bifrosts/sql/bifrost.py index a80ce04..4434437 100644 --- a/heimdallm/bifrosts/sql/bifrost.py +++ b/heimdallm/bifrosts/sql/bifrost.py @@ -41,7 +41,7 @@ def validation_only( Sequence["heimdallm.bifrosts.sql.validator.ConstraintValidator"], ], ): - """A convenience method for doing just constraint validation. This creates a + """A convenience method for doing just static analysis. This creates a Bifrost that assumes its untrusted input is a SQL query already, so it does not need to communicate with the LLM, only parse and validate it. diff --git a/notebooks/pentest.ipynb b/notebooks/pentest.ipynb index af49b56..36dd70d 100644 --- a/notebooks/pentest.ipynb +++ b/notebooks/pentest.ipynb @@ -12,7 +12,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 2, "metadata": {}, "outputs": [], "source": [ @@ -22,7 +22,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 3, "metadata": {}, "outputs": [], "source": [ @@ -58,7 +58,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 4, "metadata": {}, "outputs": [], "source": [ @@ -68,7 +68,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ @@ -80,7 +80,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 6, "metadata": {}, "outputs": [], "source": [ @@ -112,7 +112,7 @@ "metadata": {}, "outputs": [], "source": [ - "untrusted_input = \"DROP TABLE customer;\"\n", + "untrusted_input = \"SELECT id from film;\"\n", "query(untrusted_input)" ] } @@ -133,7 +133,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.6" + "version": "3.10.11" }, "orig_nbformat": 4 }, diff --git a/notebooks/quickstart.ipynb b/notebooks/quickstart/llm.ipynb similarity index 100% rename from notebooks/quickstart.ipynb rename to notebooks/quickstart/llm.ipynb diff --git a/notebooks/quickstart/validation.ipynb b/notebooks/quickstart/validation.ipynb new file mode 100644 index 0000000..5b6e833 --- /dev/null +++ b/notebooks/quickstart/validation.ipynb @@ -0,0 +1,163 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This quickstart will walk you through setting up a SQL Bifrost with OpenAI's LLM. The\n", + "end result is a function that takes natural language input and returns a trusted SQL\n", + "`SELECT` query, constrained to your requirements.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "import logging\n", + "from typing import Sequence\n", + "\n", + "import structlog\n", + "\n", + "from heimdallm.bifrosts.sql.sqlite.select.bifrost import Bifrost\n", + "from heimdallm.bifrosts.sql.sqlite.select.validator import ConstraintValidator\n", + "from heimdallm.bifrosts.sql.common import (\n", + " FqColumn,\n", + " JoinCondition,\n", + " ParameterizedConstraint,\n", + ")\n", + "\n", + "logging.basicConfig(level=logging.ERROR)\n", + "structlog.configure(logger_factory=structlog.stdlib.LoggerFactory())" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Let's define our constraint validator(s). These are used to constrain the SQL query\n", + "so that it only has access to tables and columns that you allow. For more information on\n", + "the methods that you can override in the derived class, look [here.](https://docs.heimdallm.ai/en/latest/api/bifrosts/sql/sqlite/select/validator.html)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "class CustomerConstraintValidator(ConstraintValidator):\n", + " def requester_identities(self) -> Sequence[ParameterizedConstraint]:\n", + " return [\n", + " ParameterizedConstraint(\n", + " column=\"customer.customer_id\",\n", + " placeholder=\"customer_id\",\n", + " ),\n", + " ]\n", + "\n", + " def parameterized_constraints(self) -> Sequence[ParameterizedConstraint]:\n", + " return []\n", + "\n", + " def select_column_allowed(self, column: FqColumn) -> bool:\n", + " return True\n", + "\n", + " def allowed_joins(self) -> Sequence[JoinCondition]:\n", + " return [\n", + " JoinCondition(\"payment.rental_id\", \"rental.rental_id\"),\n", + " JoinCondition(\n", + " \"customer.customer_id\",\n", + " \"payment.customer_id\",\n", + " identity=\"customer_id\",\n", + " ),\n", + " ]\n", + "\n", + " def max_limit(self) -> int | None:\n", + " return 10\n", + "\n", + "\n", + "validator = CustomerConstraintValidator()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now let's construct a Bifrost that validates SQL:\n" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "bifrost = Bifrost.validation_only(\n", + " constraint_validators=[validator],\n", + ")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can now validate constraints on SQL:\n" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "SELECT strftime('%Y-%m',payment.payment_date) as month,SUM(payment.amount) as total_amount FROM payment JOIN rental on payment.rental_id=rental.rental_id JOIN customer on payment.customer_id=customer.customer_id WHERE customer.customer_id=:customer_id GROUP BY payment.payment_date LIMIT 10\n" + ] + } + ], + "source": [ + "query = \"\"\"\n", + "SELECT\n", + " strftime('%Y-%m', payment.payment_date) AS month,\n", + " SUM(payment.amount) AS total_amount\n", + "FROM payment\n", + "JOIN rental ON payment.rental_id=rental.rental_id\n", + "JOIN customer ON payment.customer_id=customer.customer_id\n", + "WHERE customer.customer_id=:customer_id\n", + "GROUP BY month\n", + "\"\"\"\n", + "\n", + "query = bifrost.traverse(query)\n", + "print(query)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.11" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/pyproject.toml b/pyproject.toml index 6ef63c2..ee68df2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "heimdallm" -version = "1.0.0" +version = "1.0.1" description = "Construct trusted SQL queries from untrusted input" homepage = "https://github.com/amoffat/HeimdaLLM" repository = "https://github.com/amoffat/HeimdaLLM"