feat(py): handle store for tool results

pkg-py/src/commons/_data_source.py

@@ -15,6 +15,7 @@
 
 from . import _duckdb
 from ._backends import Backend, DuckDBBackend, EngineBackend
+from ._frames import is_frame
 from ._sql_guard import check_query
 
 if TYPE_CHECKING:
@@ -347,7 +348,7 @@ def _load_pins(self, labels: list[str]) -> None:
         for position, label in enumerate(labels):
             pin = self.pending.pins[label]
             value = self.pending.board.pin_read(pin)
-            if not _is_frame(value):
+            if not is_frame(value):
                 raise TypeError(
                     f"Pin {pin!r} is a {type(value).__name__}, not a data frame, "
                     f"so it cannot become the table {label!r}."
@@ -465,19 +466,13 @@ def _check_named_frames(frames: dict[str, Any]) -> None:
             "or a pins board."
         )
     for name, frame in frames.items():
-        if not _is_frame(frame):
+        if not is_frame(frame):
             raise TypeError(
                 f"{name} must be a pandas or polars data frame, got "
                 f"{type(frame).__name__}."
             )
 
 
-def _is_frame(value: Any) -> bool:
-    # Duck-typed rather than imported: pandas and polars are both optional at
-    # this boundary, and DuckDB accepts either through the same registration.
-    return hasattr(value, "__dataframe__") or hasattr(value, "columns")
-
-
 def normalize_table_registry(tables: Any) -> dict[str, TableId]:
     """Turn a `tables` argument into label -> `TableId`.
 

pkg-py/src/commons/_frames.py

@@ -0,0 +1,148 @@
+"""Recognizing and describing a data frame, whichever library it came from.
+
+pandas and polars are both optional at every boundary that accepts a frame,
+so neither is imported here: a frame is recognized and read through what it
+offers rather than through its class.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+__all__ = ["describe_frame", "is_frame"]
+
+
+# len() and [] are what describing a frame needs, so a value that merely has
+# columns — a database table, say — is not one.
+def is_frame(value: Any) -> bool:
+    return (
+        hasattr(value, "__dataframe__") or hasattr(value, "columns")
+    ) and hasattr(value, "__len__") and hasattr(value, "__getitem__")
+
+
+# ellmer's `df_schema()` describes a frame for the R agent; this describes one
+# for the Python agent, in the same terms, for whichever frame library the
+# result came from.
+MAX_SUMMARY_COLUMNS = 50
+
+
+def describe_frame(frame: Any, max_columns: int = MAX_SUMMARY_COLUMNS) -> str:
+    """A column-by-column description, so the model can write code against it."""
+    names = list(frame.columns)
+    shape = f"{_count(len(frame), 'row')} and {_count(len(names), 'column')}"
+    lines = [f"A data frame with {shape}:"]
+    lines += [
+        f"* {name}: {_describe_column(_column_at(frame, position))}"
+        for position, name in enumerate(names[:max_columns])
+    ]
+    if len(names) > max_columns:
+        lines.append(f"and {len(names) - max_columns} more columns")
+    return "\n".join(lines)
+
+
+# By position rather than by name, because pandas allows duplicate column
+# names, and a name then selects a frame rather than a column.
+def _column_at(frame: Any, position: int) -> Any:
+    iloc = getattr(frame, "iloc", None)
+    if iloc is not None:
+        return iloc[:, position]
+    return frame[:, position]
+
+
+def _count(number: int, noun: str) -> str:
+    return f"{number:,} {noun}" if number == 1 else f"{number:,} {noun}s"
+
+
+def _describe_column(column: Any) -> str:
+    kind = _kind(column.dtype)
+    missing = _missing(column)
+    described = f"{missing} missing"
+    if kind in ("numeric", "temporal"):
+        # A column with nothing left to take a range over says only how much
+        # is missing.
+        properties = (
+            [described]
+            if missing == len(column)
+            else [
+                f"range [{_value(column.min())}, {_value(column.max())}]",
+                described,
+            ]
+        )
+    elif kind == "boolean":
+        true = int(column.sum())
+        properties = [
+            f"{true} True",
+            f"{len(column) - missing - true} False",
+            described,
+        ]
+    else:
+        properties = [described]
+        unique = _describe_values(column)
+        if unique is not None:
+            properties.append(unique)
+    return f"{column.dtype} with {_flatten(properties)}"
+
+
+# pandas dtypes carry a numpy `kind` character; polars dtypes answer questions
+# about themselves instead. Neither library is imported here, because both are
+# optional wherever a frame reaches commons.
+def _kind(dtype: Any) -> str:
+    kind = getattr(dtype, "kind", None)
+    if kind is not None:
+        if kind in "iuf":
+            return "numeric"
+        if kind == "b":
+            return "boolean"
+        if kind in "Mm":
+            return "temporal"
+        return "other"
+    if dtype.is_numeric():
+        return "numeric"
+    if dtype.is_temporal():
+        return "temporal"
+    if str(dtype) == "Boolean":
+        return "boolean"
+    return "other"
+
+
+def _missing(column: Any) -> int:
+    if hasattr(column, "isna"):
+        return int(column.isna().sum())
+    return int(column.null_count())
+
+
+# Like ellmer: the values themselves only when there are few and they are
+# short, so a column of free text stays a count rather than a wall of prompt.
+def _describe_values(column: Any) -> str | None:
+    try:
+        values = _unique(column)
+    except TypeError:
+        # Unhashable values, like a column of lists, have no unique count.
+        return None
+    described = _count(len(values), "unique value")
+    quoted = [json.dumps(str(value), ensure_ascii=False) for value in values]
+    if 0 < len(values) <= 10 and sum(len(value) for value in quoted) < 200:
+        described = f"{described} ({', '.join(quoted)})"
+    return described
+
+
+def _unique(column: Any) -> list[Any]:
+    if hasattr(column, "dropna"):
+        return list(column.dropna().unique())
+    return column.drop_nulls().unique(maintain_order=True).to_list()
+
+
+# A timestamp at midnight is a date as far as the model is concerned, and the
+# time of day is noise in a range.
+def _value(value: Any) -> str:
+    isoformat = getattr(value, "isoformat", None)
+    if isoformat is None:
+        return str(value)
+    return isoformat().removesuffix("T00:00:00").replace("T", " ")
+
+
+def _flatten(properties: list[str]) -> str:
+    if len(properties) == 1:
+        return properties[0]
+    return f"{', '.join(properties[:-1])}, and {properties[-1]}"

pkg-py/src/commons/_handles.py

@@ -0,0 +1,65 @@
+"""Conversation-scoped store of tool results.
+
+A later `run_python` call reaches an earlier result as a plain variable
+(`r1`, `r2`, ...), so a tool's output can be built on rather than repeated.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from ._frames import describe_frame, is_frame
+
+__all__ = ["HandleStore"]
+
+
+# Enough rows to work with, few enough that a runaway query cannot fill the
+# conversation's memory.
+MAX_HANDLE_ROWS = 10_000
+
+
+@dataclass
+class HandleStore:
+    max_rows: int = MAX_HANDLE_ROWS
+    # Frames in the store would flood a repr, and == on one raises.
+    _values: dict[str, Any] = field(default_factory=dict, repr=False, compare=False)
+
+    def register(self, value: Any) -> str | None:
+        """Store a result and return the note telling the model how to reach it.
+
+        Values that are not frames are stored too, so a scalar measure result
+        stays available for further derivation.
+        """
+        if value is None:
+            return None
+        handle = f"r{len(self._values) + 1}"
+        if not is_frame(value):
+            self._values[handle] = value
+            return _note(handle)
+
+        try:
+            truncated = len(value) > self.max_rows
+            if truncated:
+                value = value.head(self.max_rows)
+            description = describe_frame(value)
+        except (TypeError, AttributeError):
+            # is_frame is duck-typed, so a value that quacks like a frame but
+            # cannot be read like one is still stored, only undescribed.
+            self._values[handle] = value
+            return _note(handle)
+        self._values[handle] = value
+        capped = (
+            f" Only the first {self.max_rows:,} rows are stored." if truncated else ""
+        )
+        return f"{_note(handle)}{capped}\n{description}"
+
+    def ids(self) -> list[str]:
+        return list(self._values)
+
+    def get(self, handle: str) -> Any:
+        return self._values[handle]
+
+
+def _note(handle: str) -> str:
+    return f"Available to `run_python` as `{handle}`."