@probelabs/probe 0.6.0-rc100

package/build/extract.js

@@ -0,0 +1,149 @@
+/**
+ * Extract functionality for the probe package
+ * @module extract
+ */
+
+import { exec } from 'child_process';
+import { promisify } from 'util';
+import { getBinaryPath, buildCliArgs, escapeString } from './utils.js';
+
+const execAsync = promisify(exec);
+
+/**
+ * Flag mapping for extract options
+ * Maps option keys to command-line flags
+ */
+const EXTRACT_FLAG_MAP = {
+	allowTests: '--allow-tests',
+	contextLines: '--context',
+	format: '--format',
+	inputFile: '--input-file'
+};
+
+/**
+ * Extract code blocks from files
+ * 
+ * @param {Object} options - Extract options
+ * @param {string[]} [options.files] - Files to extract from (can include line numbers with colon, e.g., "/path/to/file.rs:10")
+ * @param {string} [options.inputFile] - Path to a file containing unstructured text to extract file paths from
+ * @param {boolean} [options.allowTests] - Include test files
+ * @param {number} [options.contextLines] - Number of context lines to include
+ * @param {string} [options.format] - Output format ('markdown', 'plain', 'json')
+ * @param {Object} [options.binaryOptions] - Options for getting the binary
+ * @param {boolean} [options.binaryOptions.forceDownload] - Force download even if binary exists
+ * @param {string} [options.binaryOptions.version] - Specific version to download
+ * @param {boolean} [options.json] - Return results as parsed JSON instead of string
+ * @returns {Promise<string|Object>} - Extracted code as string or parsed JSON
+ * @throws {Error} If the extraction fails
+ */
+export async function extract(options) {
+	if (!options) {
+		throw new Error('Options object is required');
+	}
+
+	// Either files or inputFile must be provided
+	if ((!options.files || !Array.isArray(options.files) || options.files.length === 0) && !options.inputFile) {
+		throw new Error('Either files array or inputFile must be provided');
+	}
+
+	// Get the binary path
+	const binaryPath = await getBinaryPath(options.binaryOptions || {});
+
+	// Build CLI arguments from options
+	const cliArgs = buildCliArgs(options, EXTRACT_FLAG_MAP);
+
+	// If json option is true, override format to json
+	if (options.json && !options.format) {
+		cliArgs.push('--format', 'json');
+	}
+
+	// Add files as positional arguments if provided
+	if (options.files && Array.isArray(options.files) && options.files.length > 0) {
+		for (const file of options.files) {
+			cliArgs.push(escapeString(file));
+		}
+	}
+
+	// Create a single log record with all extract parameters (only in debug mode)
+	if (process.env.DEBUG === '1') {
+		let logMessage = `\nExtract:`;
+		if (options.files && options.files.length > 0) {
+			logMessage += ` files="${options.files.join(', ')}"`;
+		}
+		if (options.inputFile) logMessage += ` inputFile="${options.inputFile}"`;
+		if (options.allowTests) logMessage += " allowTests=true";
+		if (options.contextLines) logMessage += ` contextLines=${options.contextLines}`;
+		if (options.format) logMessage += ` format=${options.format}`;
+		if (options.json) logMessage += " json=true";
+		console.error(logMessage);
+	}
+
+	// Execute command
+	const command = `${binaryPath} extract ${cliArgs.join(' ')}`;
+
+	try {
+		const { stdout, stderr } = await execAsync(command);
+
+		if (stderr) {
+			console.error(`stderr: ${stderr}`);
+		}
+
+		// Parse the output to extract token usage information
+		let tokenUsage = {
+			requestTokens: 0,
+			responseTokens: 0,
+			totalTokens: 0
+		};
+
+		// Calculate approximate request tokens
+		if (options.files && Array.isArray(options.files)) {
+			tokenUsage.requestTokens = options.files.join(' ').length / 4;
+		} else if (options.inputFile) {
+			tokenUsage.requestTokens = options.inputFile.length / 4;
+		}
+
+		// Try to extract token information from the output
+		if (stdout.includes('Total tokens returned:')) {
+			const tokenMatch = stdout.match(/Total tokens returned: (\d+)/);
+			if (tokenMatch && tokenMatch[1]) {
+				tokenUsage.responseTokens = parseInt(tokenMatch[1], 10);
+				tokenUsage.totalTokens = tokenUsage.requestTokens + tokenUsage.responseTokens;
+			}
+		}
+
+		// Add token usage information to the output
+		let output = stdout;
+
+		// Add token usage information at the end if not already present
+		if (!output.includes('Token Usage:')) {
+			output += `\nToken Usage:\n  Request tokens: ${tokenUsage.requestTokens}\n  Response tokens: ${tokenUsage.responseTokens}\n  Total tokens: ${tokenUsage.totalTokens}\n`;
+		}
+
+		// Parse JSON if requested or if format is json
+		if (options.json || options.format === 'json') {
+			try {
+				const jsonOutput = JSON.parse(stdout);
+
+				// Add token usage to JSON output
+				if (!jsonOutput.token_usage) {
+					jsonOutput.token_usage = {
+						request_tokens: tokenUsage.requestTokens,
+						response_tokens: tokenUsage.responseTokens,
+						total_tokens: tokenUsage.totalTokens
+					};
+				}
+
+				return jsonOutput;
+			} catch (error) {
+				console.error('Error parsing JSON output:', error);
+				return output; // Fall back to string output with token usage
+			}
+		}
+
+		return output;
+	} catch (error) {
+		// Enhance error message with command details
+		const errorMessage = `Error executing extract command: ${error.message}\nCommand: ${command}`;
+		throw new Error(errorMessage);
+	}
+}

package/build/index.js

@@ -0,0 +1,70 @@
+/**
+ * @probelabs/probe - Node.js wrapper for the probe code search tool
+ *
+ * This module provides JavaScript functions that wrap the probe binary functionality,
+ * making it easy to use probe's powerful code search capabilities in Node.js scripts.
+ *
+ * @module @probelabs/probe
+ */
+
+import { search } from './search.js';
+import { query } from './query.js';
+import { extract } from './extract.js';
+import { delegate } from './delegate.js';
+import { getBinaryPath, setBinaryPath } from './utils.js';
+import * as tools from './tools/index.js';
+import { listFilesByLevel } from './utils/file-lister.js';
+import { DEFAULT_SYSTEM_MESSAGE } from './tools/system-message.js';
+import {
+	searchSchema,
+	querySchema,
+	extractSchema,
+	delegateSchema,
+	attemptCompletionSchema,
+	searchToolDefinition,
+	queryToolDefinition,
+	extractToolDefinition,
+	delegateToolDefinition,
+	attemptCompletionToolDefinition,
+	parseXmlToolCall
+} from './tools/common.js';
+import { searchTool, queryTool, extractTool, delegateTool } from './tools/vercel.js';
+import { ProbeAgent } from './agent/ProbeAgent.js';
+import { SimpleTelemetry, SimpleAppTracer, initializeSimpleTelemetryFromOptions } from './agent/simpleTelemetry.js';
+
+export {
+	search,
+	query,
+	extract,
+	delegate,
+	getBinaryPath,
+	setBinaryPath,
+	listFilesByLevel,
+	tools,
+	DEFAULT_SYSTEM_MESSAGE,
+	// Export AI Agent (NEW!)
+	ProbeAgent,
+	// Export telemetry classes
+	SimpleTelemetry,
+	SimpleAppTracer,
+	initializeSimpleTelemetryFromOptions,
+	// Export tool generators directly
+	searchTool,
+	queryTool,
+	extractTool,
+	delegateTool,
+	// Export schemas
+	searchSchema,
+	querySchema,
+	extractSchema,
+	delegateSchema,
+	attemptCompletionSchema,
+	// Export tool definitions
+	searchToolDefinition,
+	queryToolDefinition,
+	extractToolDefinition,
+	delegateToolDefinition,
+	attemptCompletionToolDefinition,
+	// Export parser function
+	parseXmlToolCall
+};