update repr_llm and add DataFrameSummarizer with customizable summarizing function

pyproject.toml

@@ -24,7 +24,7 @@ mkdocstrings = { version = ">=0.19,<0.22", optional = true }
 mkdocstrings-python = { version = ">=0.7.1,<0.10.0", optional = true }
 duckdb-engine = "^0.9.2"
 exceptiongroup = "^1.0.4"
-repr-llm = "^0.2.1"
+repr-llm = "^0.3.0"
 structlog = "^23.2.0"
 
 [tool.poetry.group.dev.dependencies]

src/dx/formatters/main.py

@@ -9,8 +9,8 @@
 from IPython.core.interactiveshell import InteractiveShell
 from IPython.display import display as ipydisplay
 from pandas.io.json import build_table_schema
-from repr_llm.pandas import summarize_dataframe
 
+from dx.formatters.summarizing import make_df_summary
 from dx.sampling import get_column_string_lengths, get_df_dimensions, sample_if_too_big
 from dx.settings import get_settings
 from dx.types.main import DXDisplayMode
@@ -216,7 +216,7 @@ def format_output(
  # add additional payload for LLM consumption; if any parsing/summarizing errors occur, we
  # shouldn't block displaying the bundle
  try:
- payload["text/llm+plain"] = summarize_dataframe(df)
+ payload["text/llm+plain"] = make_df_summary(df)
  except Exception as e:
  logger.debug(f"Error in summarize_dataframe: {e}")
 

src/dx/formatters/summarizing.py

@@ -0,0 +1,53 @@
+from typing import Callable, Optional
+
+import pandas as pd
+
+
+class DataFrameSummarizer:
+ _instance: "DataFrameSummarizer" = None
+ summarizing_func: Optional[Callable] = None
+
+ def __init__(self, summarizing_func: Optional[Callable] = None):
+ self.summarizing_func = summarizing_func
+
+ self._load_repr_llm()
+
+ def _load_repr_llm(self) -> None:
+ """Load repr_llm's summarize_dataframe into the summarizing_func if it's available."""
+ try:
+ from repr_llm.pandas import summarize_dataframe
+
+ self.summarizing_func = summarize_dataframe
+ except ImportError:
+ return
+
+ @classmethod
+ def instance(cls) -> "DataFrameSummarizer":
+ if cls._instance is None:
+ cls._instance = cls()
+ return cls._instance
+
+ def summarize(self, df: pd.DataFrame) -> str:
+ """Generate a summary of a dataframe using the configured summarizing_func."""
+ if not isinstance(df, pd.DataFrame):
+ raise ValueError("`df` must be a pandas DataFrame")
+
+ if self.summarizing_func is None:
+ return df.describe().to_string()
+
+ return self.summarizing_func(df)
+
+
+def get_summarizing_function() -> Optional[Callable]:
+ """Get the function to use for summarizing dataframes."""
+ return DataFrameSummarizer.instance().summarizing_func
+
+
+def set_summarizing_function(func: Callable) -> None:
+ """Set the function to use for summarizing dataframes."""
+ DataFrameSummarizer.instance().summarizing_func = func
+
+
+def make_df_summary(df: pd.DataFrame) -> str:
+ """Generate a summary of a dataframe using the configured summarizing_func."""
+ return DataFrameSummarizer.instance().summarize(df)