sqlframe 1.3.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

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/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.dist-info → sqlframe-1.4.0.dist-info}/METADATA

@@ -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
@@ -62,19 +62,19 @@ Provides-Extra: spark
 Requires-Dist: pyspark (<3.6,>=2) ; extra == 'spark'
 
 <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:
 
@@ -97,6 +97,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.dist-info → sqlframe-1.4.0.dist-info}/RECORD

@@ -1,10 +1,10 @@
 sqlframe/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sqlframe/_version.py,sha256=HGwtpza1HCPtlyqElUvIyH97K44TO13CYiYVZNezQ1M,411
+sqlframe/_version.py,sha256=R8-T9fmURjcuoxYpHTAjyNAhgJPDtI2jogCjqYYkfCU,411
 sqlframe/base/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sqlframe/base/_typing.py,sha256=DuTay8-o9W-pw3RPZCgLunKNJLS9PkaV11G_pxXp9NY,1256
 sqlframe/base/catalog.py,sha256=ATDGirouUjal05P4ymL-wIi8rgjg_8w4PoACamiO64A,37245
 sqlframe/base/column.py,sha256=p3VrtATBmjAYHollFcsdps2UJTNC-Pvyg4Zt7y4CK9w,15358
-sqlframe/base/dataframe.py,sha256=NzvzC7LE5al1uXBchmdV4Ko3ZBis6UoKt51WanYxo-8,61259
+sqlframe/base/dataframe.py,sha256=9PuqC9dBficSE-Y1v_BHyk4gK-Hd43SaVBmxBeyNnD8,62939
 sqlframe/base/decorators.py,sha256=I5osMgx9BuCgbtp4jVM2DNwYJVLzCv-OtTedhQEik0g,1882
 sqlframe/base/exceptions.py,sha256=pCB9hXX4jxZWzNg3JN1i38cv3BmpUlee5NoLYx3YXIQ,208
 sqlframe/base/function_alternatives.py,sha256=to0kv3MTJmQFeVTMcitz0AxBIoUJC3cu5LkEY5aJpoo,31318
@@ -92,8 +92,8 @@ sqlframe/standalone/readwriter.py,sha256=EZNyDJ4ID6sGNog3uP4-e9RvchX4biJJDNtc5hk
 sqlframe/standalone/session.py,sha256=wQmdu2sv6KMTAv0LRFk7TY7yzlh3xvmsyqilEtRecbY,1191
 sqlframe/standalone/types.py,sha256=KwNyuXIo-2xVVd4bZED3YrQOobKCtemlxGrJL7DrTC8,34
 sqlframe/standalone/window.py,sha256=6GKPzuxeSapJakBaKBeT9VpED1ACdjggDv9JRILDyV0,35
-sqlframe-1.3.0.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
-sqlframe-1.3.0.dist-info/METADATA,sha256=duHVFV6F1x3jqfpChK7g5u4KA5DtqZUIjYhv2fC6r3M,6861
-sqlframe-1.3.0.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-sqlframe-1.3.0.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
-sqlframe-1.3.0.dist-info/RECORD,,
+sqlframe-1.4.0.dist-info/LICENSE,sha256=VZu79YgW780qxaFJMr0t5ZgbOYEh04xWoxaWOaqIGWk,1068
+sqlframe-1.4.0.dist-info/METADATA,sha256=nnz73ML6w8WyctFzwiaKVVNr9RQwmpmfckrcKqEX_PE,7219
+sqlframe-1.4.0.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+sqlframe-1.4.0.dist-info/top_level.txt,sha256=T0_RpoygaZSF6heeWwIDQgaP0varUdSK1pzjeJZRjM8,9
+sqlframe-1.4.0.dist-info/RECORD,,