Assistant: Initial pass at implementing a data summary tool for Python (#8208)

First pass at #7114 

Provides Assistant with a `getDataSummary` tool, currently only
implemented for Python, that provides a JSON structured summary of a
data object by using the Positron API to communicate with the Variables
Comm. I updated the variable's python backend to reuse existing
functionality from the data explorer.

I used the `inspectVariables` tool as a guide for retrieving info from
the variables comm.

---------

Signed-off-by: Wes McKinney <wesm@apache.org>
Co-authored-by: Wes McKinney <wes@posit.co>
Co-authored-by: sharon <sharon-wang@users.noreply.github.com>

Author: 3 people

extensions/positron-assistant/package.json

@@ -360,6 +360,39 @@
           "positron-assistant"
         ]
       },
+      {
+        "name": "getTableSummary",
+        "displayName": "Get Table Summary",
+        "modelDescription": "Get structured information about table variables in the current session.",
+        "inputSchema": {
+          "type": "object",
+          "properties": {
+            "sessionIdentifier": {
+              "type": "string",
+              "description": "The identifier of the session that contains the tables."
+            },
+            "accessKeys": {
+              "type": "array",
+              "description": "An array of table 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

@@ -20,6 +20,29 @@ If the user asks you _how_ to do something, or asks for code rather than
 results, generate the code and return it directly without trying to execute it.
 </communication>
 
+<data-querying>
+
+**Data Object Information Workflow:**
+
+When the user asks questions that require detailed information about tabular
+data objects (DataFrames, arrays, matrices, etc.), use the `getTableSummary`
+tool to retrieve structured information such as data summaries and statistics.
+Currently, this tool is only available for Python so in R sessions you will need
+to execute code to query in-memory data.
+
+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)
+
+</data-querying>
+
 <package-management>
 You adhere to the following workflow when dealing with package management:
 

extensions/positron-assistant/src/participants.ts

@@ -283,6 +283,11 @@ abstract class PositronAssistantParticipant implements IPositronAssistantPartici
 					// Only include the documentCreate tool in the chat pane in edit or agent mode.
 					case PositronAssistantToolName.DocumentCreate:
 						return inChatPane && (isEditMode || isAgentMode);
+					// Only include the getTableSummary tool for Python sessions until supported in R
+					case PositronAssistantToolName.GetTableSummary:
+						// TODO: Remove this restriction when the tool is supported in R https://github.com/posit-dev/positron/issues/8343
+						// The logic above with TOOL_TAG_REQUIRES_ACTIVE_SESSION will handle checking for active sessions once this is removed.
+						return activeSessions.has('python');
 					// Otherwise, include the tool if it is tagged for use with Positron Assistant.
 					// Allow all tools in Agent mode.
 					default:

extensions/positron-assistant/src/tools.ts

@@ -299,6 +299,56 @@ export function registerAssistantTools(
 
 	context.subscriptions.push(inspectVariablesTool);
 
+	const getTableSummaryTool = vscode.lm.registerTool<{ sessionIdentifier: string; accessKeys: Array<Array<string>> }>(PositronAssistantToolName.GetTableSummary, {
+		/**
+		 * Called to get a summary information for one or more tabular datasets 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.querySessionTables(
+				options.input.sessionIdentifier,
+				options.input.accessKeys,
+				['summary_stats']);
+
+			// Return the result as a JSON string to the model
+			return new vscode.LanguageModelToolResult([
+				new vscode.LanguageModelTextPart(JSON.stringify(result))
+			]);
+		}
+	});
+	context.subscriptions.push(getTableSummaryTool);
+
 	const installPythonPackageTool = vscode.lm.registerTool<{
 		packages: string[];
 	}>(PositronAssistantToolName.InstallPythonPackage, {

extensions/positron-assistant/src/types.ts

@@ -7,6 +7,7 @@ export enum PositronAssistantToolName {
 	DocumentEdit = 'documentEdit',
 	EditFile = 'positron_editFile_internal',
 	ExecuteCode = 'executeCode',
+	GetTableSummary = 'getTableSummary',
 	GetPlot = 'getPlot',
 	InstallPythonPackage = 'installPythonPackage',
 	InspectVariables = 'inspectVariables',