Utilities for generating and testing Postgres queries from natural text.
setup.py- prompt-engineering test approach/examples
A python package for issuing requests for SQL translation to OpenAI LLM APIs. This package includes three main components:
default_openai_config.yamlincludes the default configuration for the OpenAI model.db_schema.pyincludes methods for extracting database schema information from a Postgres database to include in prompts sent to the OpenAI models.gen_query.pyincludes methods for composing a final prompt from a combination of pre-specified instructions (e.g. instructions to return Postgres results), schema details, and a natural-language query (prompt.pyprovides underlying prompt construction utilities that can be used for further customization).
An app for generating and testing different combinations of schema information, initialization prompts, and user prompts for text-to-sql translation. This is a tool for rapid but unsystematic experimentation with different prompts and for building intuition on how different types of prompts affect model output.
importosfrompprintimportpprintimportbitdotiofromdotenvimportload_dotenvfrompg_text_queryimportget_db_schema, get_default_prompt, generate_query# Initialize OPENAI_API_KEY and BITIO_KEYload_dotenv() DB_NAME="bitdotio/palmerpenguins"b=bitdotio.bitdotio(os.getenv("BITIO_KEY")) # Extract a structured db schema from Postgreswithb.pooled_cursor(DB_NAME) ascur: db_schema=get_db_schema(cur, DB_NAME) pprint(db_schema)Output:
{"description": "# Palmer\xa0Archipelago (Antarctica) penguin data\n""\n""\n"">\xa0The goal of palmerpenguins is to provide a great dataset ""for data exploration & visualization, as an alternative to ""`iris`.\n""\n""\n""\\- ""[source](https://github.com/allisonhorst/palmerpenguins/)\n""name": "bitdotio/palmerpenguins", "schemata": [{"description": "standard public schema", "is_foreign": False, "name": "public", "tables": [{"columns": [{"character_maximum_length": None, "column_default": None, "data_type": "text", "description": "penguin species ""(Adélie, Chinstrap, ""or\xa0Gentoo)", "is_nullable": "YES", "name": "species", "ordinal_position": 2, }, ...TRUNCATED...{"character_maximum_length": None, "column_default": None, "data_type": "bigint", "description": "study year (2007, ""2008, or 2007)", "is_nullable": "YES", "name": "year", "ordinal_position": 9, }, ], "description": "Measurements of 344 penguins of ""three different species from three ""islands in the Palmer archipelago.", "name": "penguins", } ], "views": [], } ], }# Construct a prompt that includes text description of queryprompt=get_default_prompt( "most common species and island for each island", db_schema, ) # Note: prompt includes extra `SELECT 1` as a naive approach to hinting for# raw SQL continuationprint(prompt)Output (prompt includes extra SELECT 1 query as a naive approach to hinting for # raw SQL continuation):
-- Language PostgreSQL -- Table penguins, columns = [species text, island text, bill_length_mm double precision, bill_depth_mm double precision, flipper_length_mm bigint, body_mass_g bigint, sex text, year bigint] -- A PostgreSQL query to return 1 and a PostgreSQL query for most common species and island for each island SELECT 1;# Using default OpenAI request config, which can be overriden here w/ kwargsquery=generate_query(prompt) print(query)Output:
SELECT species, island, COUNT(*) FROM penguins GROUP BY species, islandQuery validation (using pglast)
query=raise_if_invalid_query("['not', 'valid', 'sql']")Output:
pg_text_query.errors.QueryGenError: Generated query is not valid PostgreSQLpip install streamlit streamlit playground/app.py