sqlframe 1.3.0__tar.gz → 1.4.0__tar.gz

{sqlframe-1.3.0 → sqlframe-1.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 1.3.0
+Version: 1.4.0
 Summary: Taking the Spark out of PySpark by converting to SQL
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman
@@ -29,19 +29,19 @@ Provides-Extra: spark
 License-File: LICENSE
 
 <div align="center">
-  <img src="https://sqlframe.readthedocs.io/en/latest/docs/images/sqlframe_logo.png" alt="SQLFrame Logo" width="400"/>
+  <img src="https://sqlframe.readthedocs.io/en/stable/docs/images/sqlframe_logo.png" alt="SQLFrame Logo" width="400"/>
 </div>
 
 SQLFrame implements the PySpark DataFrame API in order to enable running transformation pipelines directly on database engines - no Spark clusters or dependencies required.
 
 SQLFrame currently supports the following engines (many more in development):
 
-* [BigQuery](https://sqlframe.readthedocs.io/en/latest/bigquery/)
-* [DuckDB](https://sqlframe.readthedocs.io/en/latest/duckdb)
-* [Postgres](https://sqlframe.readthedocs.io/en/latest/postgres)
+* [BigQuery](https://sqlframe.readthedocs.io/en/stable/bigquery/)
+* [DuckDB](https://sqlframe.readthedocs.io/en/stable/duckdb)
+* [Postgres](https://sqlframe.readthedocs.io/en/stable/postgres)
 
 SQLFrame also has a "Standalone" session that be used to generate SQL without any connection to a database engine.
-* [Standalone](https://sqlframe.readthedocs.io/en/latest/standalone)
+* [Standalone](https://sqlframe.readthedocs.io/en/stable/standalone)
 
 SQLFrame is great for:
 
@@ -64,6 +64,12 @@ pip install sqlframe
 
 See specific engine documentation for additional setup instructions.
 
+## Configuration
+
+SQLFrame generates consistently accurate yet complex SQL for engine execution. 
+However, when using df.sql(), it produces more human-readable SQL. 
+For details on how to configure this output and leverage OpenAI to enhance the SQL, see [Generated SQL Configuration](https://sqlframe.readthedocs.io/en/stable/configuration/#generated-sql).
+
 ## Example Usage
 
 ```python

{sqlframe-1.3.0 → sqlframe-1.4.0}/README.md

@@ -1,17 +1,17 @@
 <div align="center">
-  <img src="https://sqlframe.readthedocs.io/en/latest/docs/images/sqlframe_logo.png" alt="SQLFrame Logo" width="400"/>
+  <img src="https://sqlframe.readthedocs.io/en/stable/docs/images/sqlframe_logo.png" alt="SQLFrame Logo" width="400"/>
 </div>
 
 SQLFrame implements the PySpark DataFrame API in order to enable running transformation pipelines directly on database engines - no Spark clusters or dependencies required.
 
 SQLFrame currently supports the following engines (many more in development):
 
-* [BigQuery](https://sqlframe.readthedocs.io/en/latest/bigquery/)
-* [DuckDB](https://sqlframe.readthedocs.io/en/latest/duckdb)
-* [Postgres](https://sqlframe.readthedocs.io/en/latest/postgres)
+* [BigQuery](https://sqlframe.readthedocs.io/en/stable/bigquery/)
+* [DuckDB](https://sqlframe.readthedocs.io/en/stable/duckdb)
+* [Postgres](https://sqlframe.readthedocs.io/en/stable/postgres)
 
 SQLFrame also has a "Standalone" session that be used to generate SQL without any connection to a database engine.
-* [Standalone](https://sqlframe.readthedocs.io/en/latest/standalone)
+* [Standalone](https://sqlframe.readthedocs.io/en/stable/standalone)
 
 SQLFrame is great for:
 
@@ -34,6 +34,12 @@ pip install sqlframe
 
 See specific engine documentation for additional setup instructions.
 
+## Configuration
+
+SQLFrame generates consistently accurate yet complex SQL for engine execution. 
+However, when using df.sql(), it produces more human-readable SQL. 
+For details on how to configure this output and leverage OpenAI to enhance the SQL, see [Generated SQL Configuration](https://sqlframe.readthedocs.io/en/stable/configuration/#generated-sql).
+
 ## Example Usage
 
 ```python

{sqlframe-1.3.0 → sqlframe-1.4.0}/docs/configuration.md

@@ -184,46 +184,59 @@ The dialect of the generated SQL will be based on the session's dialect. However
 df.sql(dialect="bigquery")
 ```
 
-### OpenAI Enriched
+### OpenAI Enrichment
 
-OpenAI's models can be used to enrich the generated SQL to make it more human-like. 
-This is useful when you want to generate SQL that is more readable for humans.
-You must have `OPENAI_API_KEY` set in your environment variables to use this feature.
+OpenAI's models can be used to enrich the generated SQL to make it more human-like.
+You can have it just provide more readable CTE names or you can have it try to make the whole SQL statement more readable.
+
+#### Example
 
 ```python
 # create session and `df` like normal
 # The model to use defaults to `gpt-4o` but can be changed by passing a string to the `openai_model` parameter.
->>> df.sql(optimize=False, use_openai=True)
-WITH natality_data AS (
-  SELECT
-    year,
-    ever_born
-  FROM `bigquery-public-data`.`samples`.`natality`
-), single_child_families AS (
+>>> df.sql(openai_config={"mode": "cte_only", "model": "gpt-3.5-turbo"})
+WITH `single_child_families_by_year` AS (
   SELECT
-    year,
-    COUNT(*) AS num_single_child_families
-  FROM natality_data
-  WHERE ever_born = 1
-  GROUP BY year
-), lagged_families AS (
-  SELECT
-    year,
-    num_single_child_families,
-    LAG(num_single_child_families, 1) OVER (ORDER BY year) AS last_year_num_single_child_families
-  FROM single_child_families
-), percent_change_families AS (
+    `natality`.`year` AS `year`,
+    COUNT(*) AS `num_single_child_families`
+  FROM `bigquery-public-data`.`samples`.`natality` AS `natality`
+  WHERE
+    `natality`.`ever_born` = 1
+  GROUP BY
+    `natality`.`year`
+), `families_with_percent_change` AS (
   SELECT
-    year,
-    num_single_child_families,
-    ((num_single_child_families - last_year_num_single_child_families) / last_year_num_single_child_families) AS percent_change
-  FROM lagged_families
-  ORDER BY ABS(percent_change) DESC
+    `single_child_families_by_year`.`year` AS `year`,
+    `single_child_families_by_year`.`num_single_child_families` AS `num_single_child_families`,
+    LAG(`single_child_families_by_year`.`num_single_child_families`, 1) OVER (ORDER BY `single_child_families_by_year`.`year`) AS `last_year_num_single_child_families`
+  FROM `single_child_families_by_year` AS `single_child_families_by_year`
 )
 SELECT
-  year,
-  FORMAT('%\'.0f', ROUND(CAST(num_single_child_families AS FLOAT64), 0)) AS `new families single child`,
-  FORMAT('%\'.2f', ROUND(CAST((percent_change * 100) AS FLOAT64), 2)) AS `percent change`
-FROM percent_change_families
+  `families_with_percent_change`.`year` AS `year`,
+  FORMAT('%\'.0f', ROUND(CAST(`families_with_percent_change`.`num_single_child_families` AS FLOAT64), 0)) AS `new families single child`,
+  FORMAT(
+    '%\'.2f',
+    ROUND(
+      CAST((
+        (
+          (
+            `families_with_percent_change`.`num_single_child_families` - `families_with_percent_change`.`last_year_num_single_child_families`
+          ) / `families_with_percent_change`.`last_year_num_single_child_families`
+        ) * 100
+      ) AS FLOAT64),
+      2
+    )
+  ) AS `percent change`
+FROM `families_with_percent_change` AS `families_with_percent_change`
+ORDER BY
+  ABS(`percent_change`) DESC
 LIMIT 5
 ```
+
+#### Parameters
+
+| Parameter         | Description                                                           | Default    |
+|-------------------|-----------------------------------------------------------------------|------------|
+| `mode`            | The mode to use. Can be `cte_only` or `full`.                         | `cte_only` |
+| `model`           | The OpenAI model to use. Note: The default may change in new releases | `gpt-4o`   |
+| `prompt_override` | A string to use to override the default prompt.                       | None       |

{sqlframe-1.3.0 → sqlframe-1.4.0}/sqlframe/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.3.0'
-__version_tuple__ = version_tuple = (1, 3, 0)
+__version__ = version = '1.4.0'
+__version_tuple__ = version_tuple = (1, 4, 0)

{sqlframe-1.3.0 → sqlframe-1.4.0}/sqlframe/base/dataframe.py

@@ -2,13 +2,16 @@
 
 from __future__ import annotations
 
+import enum
 import functools
 import itertools
+import json
 import logging
 import sys
 import typing as t
 import zlib
 from copy import copy
+from dataclasses import dataclass
 
 import sqlglot
 from prettytable import PrettyTable
@@ -75,6 +78,46 @@ JOIN_HINTS = {
 DF = t.TypeVar("DF", bound="_BaseDataFrame")
 
 
+class OpenAIMode(enum.Enum):
+    CTE_ONLY = "cte_only"
+    FULL = "full"
+
+    @property
+    def is_cte_only(self) -> bool:
+        return self == OpenAIMode.CTE_ONLY
+
+    @property
+    def is_full(self) -> bool:
+        return self == OpenAIMode.FULL
+
+
+@dataclass
+class OpenAIConfig:
+    mode: OpenAIMode = OpenAIMode.CTE_ONLY
+    model: str = "gpt-4o"
+    prompt_override: t.Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, config: t.Dict[str, t.Any]) -> OpenAIConfig:
+        if "mode" in config:
+            config["mode"] = OpenAIMode(config["mode"].lower())
+        return cls(**config)
+
+    def get_prompt(self, dialect: Dialect) -> str:
+        if self.prompt_override:
+            return self.prompt_override
+        if self.mode.is_cte_only:
+            return f"You are a backend tool that creates unique CTE alias names match what a human would write and in snake case. You respond without code blocks and only a json payload with the key being the CTE name that is being replaced and the value being the new CTE human readable name."
+        return f"""
+        You are a backend tool that converts correct {dialect} SQL to simplified and more human readable version.
+        You respond without code block with rewritten {dialect} SQL.
+        You don't change any column names in the final select because the user expects those to remain the same.
+        You make unique CTE alias names match what a human would write and in snake case.
+        You improve formatting with spacing and line-breaks.
+        You remove redundant parenthesis and aliases.
+        When remove extra quotes, make sure to keep quotes around words that could be reserved words"""
+
+
 class _BaseDataFrameNaFunctions(t.Generic[DF]):
     def __init__(self, df: DF):
         self.df = df
@@ -476,8 +519,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         dialect: DialectType = None,
         optimize: bool = True,
         pretty: bool = True,
-        use_openai: bool = False,
-        openai_model: str = "gpt-4o",
+        openai_config: t.Optional[t.Union[t.Dict[str, t.Any], OpenAIConfig]] = None,
         as_list: bool = False,
         **kwargs,
     ) -> t.Union[str, t.List[str]]:
@@ -487,6 +529,11 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         select_expressions = df._get_select_expressions()
         output_expressions: t.List[t.Union[exp.Select, exp.Cache, exp.Drop]] = []
         replacement_mapping: t.Dict[exp.Identifier, exp.Identifier] = {}
+        openai_config = (
+            OpenAIConfig.from_dict(openai_config)
+            if openai_config is not None and isinstance(openai_config, dict)
+            else openai_config
+        )
 
         for expression_type, select_expression in select_expressions:
             select_expression = select_expression.transform(
@@ -497,7 +544,7 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
                 select_expression = t.cast(
                     exp.Select, self.session._optimize(select_expression, dialect=dialect)
                 )
-            elif use_openai:
+            elif openai_config:
                 qualify(select_expression, dialect=dialect, schema=self.session.catalog._schema)
                 pushdown_projections(select_expression, schema=self.session.catalog._schema)
 
@@ -556,35 +603,32 @@ class _BaseDataFrame(t.Generic[SESSION, WRITER, NA, STAT, GROUP_DATA]):
         results = []
         for expression in output_expressions:
             sql = expression.sql(dialect=dialect, pretty=pretty, **kwargs)
-            if use_openai:
+            if openai_config:
+                assert isinstance(openai_config, OpenAIConfig)
                 verify_openai_installed()
                 from openai import OpenAI
 
                 client = OpenAI()
-                prompt = f"""
-                You are a backend tool that converts correct {dialect} SQL to simplified and more human readable version.
-                You respond without code block with rewritten {dialect} SQL.
-                You don't change any column names in the final select because the user expects those to remain the same.
-                You make unique CTE alias names match what a human would write and in snake case.
-                You improve formatting with spacing and line-breaks.
-                You remove redundant parenthesis and aliases.
-                When remove extra quotes, make sure to keep quotes around words that could be reserved words
-                """
                 chat_completed = client.chat.completions.create(
                     messages=[
-                        {
+                        {  # type: ignore
                             "role": "system",
-                            "content": prompt,
+                            "content": openai_config.get_prompt(dialect),
                         },
                         {
                             "role": "user",
                             "content": sql,
                         },
                     ],
-                    model=openai_model,
+                    model=openai_config.model,
                 )
                 assert chat_completed.choices[0].message.content is not None
-                sql = chat_completed.choices[0].message.content
+                if openai_config.mode.is_cte_only:
+                    cte_replacement_mapping = json.loads(chat_completed.choices[0].message.content)
+                    for old_name, new_name in cte_replacement_mapping.items():
+                        sql = sql.replace(old_name, new_name)
+                else:
+                    sql = chat_completed.choices[0].message.content
             results.append(sql)
 
         if as_list:

{sqlframe-1.3.0 → sqlframe-1.4.0}/sqlframe.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sqlframe
-Version: 1.3.0
+Version: 1.4.0
 Summary: Taking the Spark out of PySpark by converting to SQL
 Home-page: https://github.com/eakmanrq/sqlframe
 Author: Ryan Eakman
@@ -29,19 +29,19 @@ Provides-Extra: spark
 License-File: LICENSE
 
 <div align="center">
-  <img src="https://sqlframe.readthedocs.io/en/latest/docs/images/sqlframe_logo.png" alt="SQLFrame Logo" width="400"/>
+  <img src="https://sqlframe.readthedocs.io/en/stable/docs/images/sqlframe_logo.png" alt="SQLFrame Logo" width="400"/>
 </div>
 
 SQLFrame implements the PySpark DataFrame API in order to enable running transformation pipelines directly on database engines - no Spark clusters or dependencies required.
 
 SQLFrame currently supports the following engines (many more in development):
 
-* [BigQuery](https://sqlframe.readthedocs.io/en/latest/bigquery/)
-* [DuckDB](https://sqlframe.readthedocs.io/en/latest/duckdb)
-* [Postgres](https://sqlframe.readthedocs.io/en/latest/postgres)
+* [BigQuery](https://sqlframe.readthedocs.io/en/stable/bigquery/)
+* [DuckDB](https://sqlframe.readthedocs.io/en/stable/duckdb)
+* [Postgres](https://sqlframe.readthedocs.io/en/stable/postgres)
 
 SQLFrame also has a "Standalone" session that be used to generate SQL without any connection to a database engine.
-* [Standalone](https://sqlframe.readthedocs.io/en/latest/standalone)
+* [Standalone](https://sqlframe.readthedocs.io/en/stable/standalone)
 
 SQLFrame is great for:
 
@@ -64,6 +64,12 @@ pip install sqlframe
 
 See specific engine documentation for additional setup instructions.
 
+## Configuration
+
+SQLFrame generates consistently accurate yet complex SQL for engine execution. 
+However, when using df.sql(), it produces more human-readable SQL. 
+For details on how to configure this output and leverage OpenAI to enhance the SQL, see [Generated SQL Configuration](https://sqlframe.readthedocs.io/en/stable/configuration/#generated-sql).
+
 ## Example Usage
 
 ```python