Deprecate Genie Truncation (#171)

Author: aravind-segu

src/databricks_ai_bridge/genie.py

@@ -1,14 +1,12 @@
 import bisect
 import logging
 import time
-import warnings
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Optional, Union
 
 import mlflow
 import pandas as pd
-import tiktoken
 from databricks.sdk import WorkspaceClient
 
 MAX_TOKENS_OF_DATA = 20000
@@ -17,6 +15,8 @@
 
 # Define a function to count tokens
 def _count_tokens(text):
+    import tiktoken
+
     encoding = tiktoken.encoding_for_model("gpt-4o")
     return len(encoding.encode(text))
 
@@ -62,6 +62,16 @@ def _parse_query_result(resp, truncate_results) -> Union[str, pd.DataFrame]:
         rows.append(row)
 
     dataframe = pd.DataFrame(rows, columns=header)
+
+    if truncate_results:
+        query_result = _truncate_result(dataframe)
+    else:
+        query_result = dataframe.to_markdown()
+
+    return query_result.strip()
+
+
+def _truncate_result(dataframe):
     query_result = dataframe.to_markdown()
     tokens_used = _count_tokens(query_result)
 
@@ -88,15 +98,7 @@ def is_too_big(n):
     # Double-check edge case if we overshot by one
     if _count_tokens(truncated_result) > MAX_TOKENS_OF_DATA:
         truncated_result = truncated_df.iloc[:-1].to_markdown()
-
-    if not truncate_results:
-        warnings.warn(
-            "Detected large Genie output, truncating it to better fit in LLM context windows. Automatic result truncation in Genie is deprecated and will be disabled by default in a future release; we recommend truncating large Genie results in your agent code instead, if needed. To keep automatic truncation for large Genie outputs enabled, set truncate_results=True when initializing the Genie class.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-
-    return truncated_result.strip()
+    return truncated_result
 
 
 class Genie:

tests/databricks_ai_bridge/test_genie.py

@@ -188,7 +188,8 @@ def test_parse_query_result_with_null_values():
     assert result == expected_df.to_markdown()
 
 
-def test_parse_query_result_trims_data():
+@pytest.mark.parametrize("truncate_results", [True, False])
+def test_parse_query_result_trims_data(truncate_results):
     # patch MAX_TOKENS_OF_DATA to 100 for this test
     with patch("databricks_ai_bridge.genie.MAX_TOKENS_OF_DATA", 100):
         resp = {
@@ -216,22 +217,57 @@ def test_parse_query_result_trims_data():
                 ]
             },
         }
-        result = _parse_query_result(resp, truncate_results=True)
-        assert (
-            result
-            == pd.DataFrame(
-                {
-                    "id": [1, 2, 3],
-                    "name": ["Alice", "Bob", "Charlie"],
-                    "created_at": [
-                        datetime(2023, 10, 1).date(),
-                        datetime(2023, 10, 2).date(),
-                        datetime(2023, 10, 3).date(),
-                    ],
-                }
-            ).to_markdown()
-        )
-        assert _count_tokens(result) <= 100
+        result = _parse_query_result(resp, truncate_results=truncate_results)
+
+        if truncate_results:
+            assert (
+                result
+                == pd.DataFrame(
+                    {
+                        "id": [1, 2, 3],
+                        "name": ["Alice", "Bob", "Charlie"],
+                        "created_at": [
+                            datetime(2023, 10, 1).date(),
+                            datetime(2023, 10, 2).date(),
+                            datetime(2023, 10, 3).date(),
+                        ],
+                    }
+                ).to_markdown()
+            )
+            assert _count_tokens(result) <= 100
+        else:
+            assert (
+                result
+                == pd.DataFrame(
+                    {
+                        "id": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+                        "name": [
+                            "Alice",
+                            "Bob",
+                            "Charlie",
+                            "David",
+                            "Eve",
+                            "Frank",
+                            "Grace",
+                            "Hank",
+                            "Ivy",
+                            "Jack",
+                        ],
+                        "created_at": [
+                            datetime(2023, 10, 1).date(),
+                            datetime(2023, 10, 2).date(),
+                            datetime(2023, 10, 3).date(),
+                            datetime(2023, 10, 4).date(),
+                            datetime(2023, 10, 5).date(),
+                            datetime(2023, 10, 6).date(),
+                            datetime(2023, 10, 7).date(),
+                            datetime(2023, 10, 8).date(),
+                            datetime(2023, 10, 9).date(),
+                            datetime(2023, 10, 10).date(),
+                        ],
+                    }
+                ).to_markdown()
+            )
 
 
 def markdown_to_dataframe(markdown_str: str) -> pd.DataFrame: