initial pass at implementing a data summary tool for Python

Author: melissa-barca

extensions/positron-assistant/package.json

@@ -262,6 +262,39 @@
           "positron-assistant"
         ]
       },
+      {
+        "name": "getDataSummary",
+        "displayName": "Get Data Summary",
+        "modelDescription": "Get structured information about data objects in the current session.",
+        "inputSchema": {
+          "type": "object",
+          "properties": {
+            "sessionIdentifier": {
+              "type": "string",
+              "description": "The identifier of the session that contains the data."
+            },
+            "accessKeys": {
+              "type": "array",
+              "description": "An array of data variables to summarize.",
+              "items": {
+                "type": "array",
+                "description": "A list of access keys that identify a variable by specifying its path.",
+                "items": {
+                  "type": "string",
+                  "description": "An access key that uniquely identifies a variable among its siblings."
+                }
+              }
+            }
+          },
+          "required": [
+            "sessionIdentifier",
+            "accessKeys"
+          ]
+        },
+        "tags": [
+          "positron-assistant"
+        ]
+      },
       {
         "name": "getProjectTree",
         "displayName": "Get Project Tree",

extensions/positron-assistant/src/md/prompts/chat/agent.md

@@ -23,6 +23,21 @@ results, generate the code and return it directly without trying to execute it.
 <package-management>
 You adhere to the following workflow when dealing with package management:
 
+**Data Object Information Workflow:**
+
+When the user asks questions that require detailed information about data objects (DataFrames, arrays, matrices, etc.), use the `getDataSummary` tool to retrieve structured information such as data summaries and statistics.
+
+To use the tool effectively:
+
+1. First ensure you have the correct `sessionIdentifier` from the user context
+2. Provide the `accessKeys` array with the path to the specific data objects
+   - Each access key is an array of strings representing the path to the variable
+   - If the user references a variable by name, determine the access key from context or previous tool results
+3. Do not call this tool when:
+   - The variables do not appear in the user context
+   - There is no active session
+   - The user only wants to see the structure/children of objects (use `inspectVariables` instead)
+
 **Package Management Workflow:**
 
 1. Before generating code that requires packages, you must first use the appropriate tool to check if each required package is installed. To do so, first determine the target language from the user's request or context
@@ -44,4 +59,4 @@ You adhere to the following workflow when dealing with package management:
 4. Never use Python tools when generating R code, or R tools when generating Python code
 5. Never instruct users to install, load, or import packages that are already loaded in their session
 6. Do not generate conditional code (if/then statements) to check package availability. Use the provided tools to determine package status and generate only the necessary installation or loading code based on the tool results
-</package-management>
+   </package-management>

extensions/positron-assistant/src/tools.ts

@@ -283,6 +283,55 @@ export function registerAssistantTools(
 	context.subscriptions.push(inspectVariablesTool);
 
 	context.subscriptions.push(ProjectTreeTool);
+
+	const getDataSummaryTool = vscode.lm.registerTool<{ sessionIdentifier: string; accessKeys: Array<Array<string>> }>(PositronAssistantToolName.GetDataSummary, {
+		/**
+		 * Called to get a data summary of one or more variables in the current session.
+		 * @param options The options for the tool invocation.
+		 * @param token The cancellation token.
+		 * @returns A vscode.LanguageModelToolResult containing the data summary.
+		 */
+		invoke: async (options, token) => {
+
+			// If no session identifier is provided, return an empty array.
+			if (!options.input.sessionIdentifier || options.input.sessionIdentifier === 'undefined') {
+				return new vscode.LanguageModelToolResult([
+					new vscode.LanguageModelTextPart('[[]]')
+				]);
+			}
+
+			// temporarily only enable for Python sessions
+			let session: positron.LanguageRuntimeSession | undefined;
+			const sessions = await positron.runtime.getActiveSessions();
+			if (sessions && sessions.length > 0) {
+				session = sessions.find(
+					(session) => session.metadata.sessionId === options.input.sessionIdentifier,
+				);
+			}
+			if (!session) {
+				return new vscode.LanguageModelToolResult([
+					new vscode.LanguageModelTextPart('[[]]')
+				]);
+			}
+
+			if (session.runtimeMetadata.languageId !== 'python') {
+				return new vscode.LanguageModelToolResult([
+					new vscode.LanguageModelTextPart('[[]]')
+				]);
+			}
+
+			// Call the Positron API to get the session variable data summaries
+			const result = await positron.runtime.getSessionVariableDataSummaries(
+				options.input.sessionIdentifier,
+				options.input.accessKeys);
+
+			// Return the result as a JSON string to the model
+			return new vscode.LanguageModelToolResult([
+				new vscode.LanguageModelTextPart(JSON.stringify(result))
+			]);
+		}
+	});
+	context.subscriptions.push(getDataSummaryTool);
 }
 
 /**

extensions/positron-assistant/src/types.ts

@@ -7,6 +7,7 @@ export enum PositronAssistantToolName {
 	DocumentEdit = 'documentEdit',
 	EditFile = 'positron_editFile_internal',
 	ExecuteCode = 'executeCode',
+	GetDataSummary = 'getDataSummary',
 	GetPlot = 'getPlot',
 	InspectVariables = 'inspectVariables',
 	SelectionEdit = 'selectionEdit',

extensions/positron-python/python_files/posit/positron/variables.py

@@ -31,6 +31,7 @@
     InspectRequest,
     ListRequest,
     RefreshParams,
+    SummarizeDataRequest,
     UpdateParams,
     Variable,
     VariableKind,
@@ -137,6 +138,10 @@ def handle_msg(
         elif isinstance(request, ViewRequest):
             self._perform_view_action(request.params.path)
 
+        elif isinstance(request, SummarizeDataRequest):
+            for path in request.params.paths:
+                self._perform_get_variable_summary(path)
+
         else:
             logger.warning(f"Unhandled request: {request}")
 
@@ -707,6 +712,88 @@ def _send_details(self, _path: list[str], value: Any = None):
         msg = InspectedVariable(children=children, length=len(children))
         self._send_result(msg.dict())
 
+    def _perform_get_variable_summary(self, path: list[str]) -> None:
+        """RPC handler for getting variable data summary."""
+        try:
+            self._summarize_data(path)
+        except Exception as err:
+            self._send_error(
+                JsonRpcErrorCode.INTERNAL_ERROR,
+                f"Error summarizing variable at '{path}': {err}",
+            )
+
+    def _summarize_data(self, path: list[str]):
+        """Compute statistical summary for a variable without opening a data explorer."""
+        from .data_explorer import DataExplorerState, _get_table_view, _value_type_is_supported
+        from .data_explorer_comm import FormatOptions
+
+        is_found, value = self._find_var(path)
+        if not is_found:
+            raise ValueError(f"Cannot find variable at '{path}' to summarize")
+
+        if not _value_type_is_supported(value):
+            raise ValueError(f"Variable at '{path}' is not supported for summary")
+
+        # Create a temporary table view without a comm
+        temp_state = DataExplorerState("temp_summary")
+        try:
+            table_view = _get_table_view(value, None, temp_state, self.kernel.job_queue)
+        except Exception as e:
+            raise ValueError(f"Failed to create table view: {e}") from e
+
+        # Get the number of columns and build schema manually
+        try:
+            num_columns = table_view.table.shape[1]
+            num_rows = table_view.table.shape[0]
+
+            # Get column schemas directly using the internal method
+            column_schemas = []
+            for i in range(num_columns):
+                column_schema = table_view._get_single_column_schema(i)  # noqa: SLF001
+                column_schemas.append(column_schema)
+
+            # Create schema object manually
+            from .data_explorer_comm import TableSchema
+
+            schema = TableSchema(columns=column_schemas)
+        except Exception as e:
+            raise ValueError(f"Failed to get schema: {e}") from e
+
+        # Create default format options
+        format_options = FormatOptions(
+            large_num_digits=4,
+            small_num_digits=6,
+            max_integral_digits=7,
+            max_value_length=1000,
+            thousands_sep=None,
+        )
+
+        profiles = []
+        for i, column in enumerate(schema.columns):
+            try:
+                summary_stats = table_view._prof_summary_stats(i, format_options)  # noqa: SLF001
+                profiles.append(
+                    {
+                        "column_name": column.column_name,
+                        "type_display": column.type_display,
+                        "summary_stats": summary_stats,
+                    }
+                )
+            except Exception as e:
+                # Skip columns that can't be summarized
+                logger.warning(f"Skipping summary stats for column {i} ({column.column_name}): {e}")
+                continue
+
+        self._send_result(
+            {
+                "schema": {
+                    "num_rows": num_rows,
+                    "num_columns": num_columns,
+                },
+                "column_profiles": profiles,
+            }
+        )
+
 
 def _summarize_variable(key: Any, value: Any, display_name: str | None = None) -> Variable | None:
     """

extensions/positron-python/python_files/posit/positron/variables_comm.py

@@ -183,6 +183,9 @@ class VariablesBackendRequest(str, enum.Enum):
     # Request a viewer for a variable
     View = "view"
 
+    # Summarize data
+    SummarizeData = "summarize_data"
+
 
 class ListRequest(BaseModel):
     """
@@ -361,6 +364,7 @@ class VariablesBackendMessageContent(BaseModel):
         InspectRequest,
         ClipboardFormatRequest,
         ViewRequest,
+        SummarizeDataRequest,
     ] = Field(..., discriminator="method")
 
 
@@ -417,6 +421,33 @@ class RefreshParams(BaseModel):
     )
 
 
+class SummarizeDataParams(BaseModel):
+    paths: List[List[StrictStr]] = Field(
+        description="Array of paths to variables to summarize, each path is an array of access keys.",
+    )
+
+
+class SummarizeDataRequest(BaseModel):
+    """
+    Request that the runtime summarize data in a variable.
+    """
+
+    params: SummarizeDataParams = Field(
+        description="Parameters to the SummarizeData method",
+    )
+
+    method: Literal[VariablesBackendRequest.SummarizeData] = Field(
+        description="The JSON-RPC method name (summarize_data)",
+    )
+
+    jsonrpc: str = (
+        Field(
+            default="2.0",
+            description="The JSON-RPC version specifier",
+        ),
+    )
+
+
 VariableList.update_forward_refs()
 
 InspectedVariable.update_forward_refs()
@@ -450,3 +481,9 @@ class RefreshParams(BaseModel):
 UpdateParams.update_forward_refs()
 
 RefreshParams.update_forward_refs()
+
+SummarizeDataParams.update_forward_refs()
+
+SummarizeDataRequest.update_forward_refs()
+
+VariablesBackendMessageContent.update_forward_refs()

src/positron-dts/positron.d.ts

@@ -1788,6 +1788,11 @@ declare module 'positron' {
 			accessKeys?: Array<Array<string>>):
 			Thenable<Array<Array<RuntimeVariable>>>;
 
+		export function getSessionVariableDataSummaries(
+			sessionId: string,
+			accessKeys: Array<Array<string>>):
+			Thenable<Array<string>>;
+
 		/**
 		 * Register a handler for runtime client instances. This handler will be called
 		 * whenever a new client instance is created by a language runtime of the given

src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts

@@ -1583,6 +1583,32 @@ export class MainThreadLanguageRuntime
 		}
 	}
 
+	$getSessionVariableDataSummaries(handle: number, accessKeys: Array<Array<string>>): Promise<Array<string>> {
+		const sessionId = this.findSession(handle).sessionId;
+		const instances = this._positronVariablesService.positronVariablesInstances;
+		for (const instance of instances) {
+			if (instance.session.sessionId === sessionId) {
+				return this.getSessionVariableDataSummaries(instance, accessKeys);
+			}
+		}
+		throw new Error(`No variables provider found for session ${sessionId}`);
+	}
+
+	async getSessionVariableDataSummaries(instance: IPositronVariablesInstance, accessKeys: Array<Array<string>>): Promise<Array<string>> {
+		const client = instance.getClientInstance();
+		if (!client) {
+			throw new Error(`No variables provider available for session ${instance.session.sessionId}`);
+		}
+		if (accessKeys.length === 0) {
+			throw new Error('No access keys provided for variable data retrieval');
+		}
+		const result = [];
+		for (const accessKey of accessKeys) {
+			result.push((await client.comm.summarizeData(accessKey)));
+		}
+		return result;
+	}
+
 	// Signals that language runtime discovery is complete.
 	$completeLanguageRuntimeDiscovery(): void {
 		this._runtimeStartupService.completeDiscovery(this._id);

src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts

@@ -141,6 +141,10 @@ export function createPositronApiFactoryAndRegisterActors(accessor: ServicesAcce
 				Thenable<Array<Array<positron.RuntimeVariable>>> {
 				return extHostLanguageRuntime.getSessionVariables(sessionId, accessKeys);
 			},
+			getSessionVariableDataSummaries(sessionId: string, accessKeys: Array<Array<string>>):
+				Thenable<Array<string>> {
+				return extHostLanguageRuntime.getSessionVariableDataSummaries(sessionId, accessKeys);
+			},
 			registerClientHandler(handler: positron.RuntimeClientHandler): vscode.Disposable {
 				return extHostLanguageRuntime.registerClientHandler(handler);
 			},

src/vs/workbench/api/common/positron/extHost.positron.protocol.ts

@@ -53,6 +53,7 @@ export interface MainThreadLanguageRuntimeShape extends IDisposable {
 	$interruptSession(handle: number): Promise<void>;
 	$focusSession(handle: number): void;
 	$getSessionVariables(handle: number, accessKeys?: Array<Array<string>>): Promise<Array<Array<Variable>>>;
+	$getSessionVariableDataSummaries(handle: number, accessKeys: Array<Array<string>>): Promise<Array<string>>;
 	$emitLanguageRuntimeMessage(handle: number, handled: boolean, message: SerializableObjectWithBuffers<ILanguageRuntimeMessage>): void;
 	$emitLanguageRuntimeState(handle: number, clock: number, state: RuntimeState): void;
 	$emitLanguageRuntimeExit(handle: number, exit: ILanguageRuntimeExit): void;