Squirreling SQL tool (#359)

Author: platypii

bin/chat.js

@@ -1,4 +1,4 @@
-import { tools } from './tools.js'
+import { tools } from './tools/tools.js'
 
 /** @type {'text' | 'tool'} */
 let outputMode = 'text' // default output mode
@@ -20,6 +20,12 @@ const colors = {
   normal: '\x1b[0m', // reset
 }
 
+const ignoreMessageTypes = [
+  'response.completed',
+  'response.output_item.added',
+  'response.function_call_arguments.delta',
+]
+
 /**
  * @import { ChatInput, Message } from './types.d.ts'
  * @param {ChatInput} chatInput
@@ -89,7 +95,7 @@ async function sendToServer(chatInput) {
             summary: chunk.item.summary,
           }
           incoming.push(reasoningItem)
-        } else if (chunk.key || chunk.type === 'response.completed') {
+        } else if (chunk.key || ignoreMessageTypes.includes(chunk.type)) {
           // ignore
         } else {
           console.log('\nUnknown chunk', chunk)
@@ -171,7 +177,7 @@ async function sendMessages(messages) {
       } catch (error) {
         const message = error instanceof Error ? error.message : String(error)
         const toolName = toolCall.name ?? toolCall.id
-        write(colors.error, `\nError calling tool ${toolName}: ${message}`, colors.normal)
+        write(colors.error, `\nError calling tool ${toolName}: ${message}\n`, colors.normal)
         incoming.push({ type: 'function_call_output', output: `Error calling tool ${toolName}: ${message}`, call_id })
       }
     }

bin/tools.js

@@ -0,0 +1,76 @@
+import fs from 'fs/promises'
+
+const fileLimit = 20 // limit to 20 files per page
+
+/**
+ * @import { ToolHandler } from '../types.d.ts'
+ * @type {ToolHandler}
+ */
+export const listFiles = {
+  emoji: '📂',
+  tool: {
+    type: 'function',
+    name: 'list_files',
+    description: `List the files in a directory. Files are listed recursively up to ${fileLimit} per page.`,
+    parameters: {
+      type: 'object',
+      properties: {
+        path: {
+          type: 'string',
+          description: 'The path to list files from. Optional, defaults to the current directory.',
+        },
+        filetype: {
+          type: 'string',
+          description: 'Optional file type to filter by, e.g. "parquet", "csv". If not provided, all files are listed.',
+        },
+        offset: {
+          type: 'number',
+          description: 'Skip offset number of files in the listing. Defaults to 0. Optional.',
+        },
+      },
+    },
+  },
+  /**
+   * @param {Record<string, unknown>} args
+   * @returns {Promise<string>}
+   */
+  async handleToolCall({ path, filetype, offset = 0 }) {
+    if (typeof path !== 'string') {
+      throw new Error('Expected path to be a string')
+    }
+    if (path.includes('..') || path.includes('~') || path.startsWith('/')) {
+      throw new Error('Invalid path: ' + path)
+    }
+    if (typeof filetype !== 'undefined' && typeof filetype !== 'string') {
+      throw new Error('Expected filetype to be a string or undefined')
+    }
+    const start = validateInteger('offset', offset, 0)
+    // list files in the directory
+    const filenames = (await fs.readdir(path || '.'))
+      .filter(key => !filetype || key.endsWith(`.${filetype}`)) // filter by file type if provided
+    const limited = filenames.slice(start, start + fileLimit)
+    const end = start + limited.length
+    return `Files ${start + 1}..${end} of ${filenames.length}:\n${limited.join('\n')}`
+  },
+}
+
+/**
+ * Validates that a value is an integer within the specified range. Max is inclusive.
+ * @param {string} name - The name of the value being validated.
+ * @param {unknown} value - The value to validate.
+ * @param {number} min - The minimum allowed value (inclusive).
+ * @param {number} [max] - The maximum allowed value (inclusive).
+ * @returns {number}
+ */
+function validateInteger(name, value, min, max) {
+  if (typeof value !== 'number' || isNaN(value)) {
+    throw new Error(`Invalid number for ${name}: ${value}`)
+  }
+  if (!Number.isInteger(value)) {
+    throw new Error(`Invalid number for ${name}: ${value}. Must be an integer.`)
+  }
+  if (value < min || max !== undefined && value > max) {
+    throw new Error(`Invalid number for ${name}: ${value}. Must be between ${min} and ${max}.`)
+  }
+  return value
+}

bin/tools/listFiles.js

@@ -0,0 +1,62 @@
+
+/**
+ * Convert an array of objects to a markdown table string.
+ * Truncates cell values to maxChars, appending '…' if truncated.
+ * Marks columns as "(truncated)" in the header if any cell was truncated.
+ * @param {Record<string, any>[]} data
+ * @param {number} maxChars
+ * @returns {string}
+ */
+export function markdownTable(data, maxChars) {
+  if (data.length === 0) {
+    return '(empty table)'
+  }
+  const columns = Object.keys(data[0] ?? {})
+  const truncated = columns.map(() => false) // is column truncated?
+  /** @type {string[]} */
+  const rows = []
+  for (const row of data) {
+    const rowStrings = new Array(columns.length)
+    for (let i = 0; i < columns.length; i++) {
+      const column = columns[i]
+      const value = column && row[column]
+      if (value === null || value === undefined) {
+        rowStrings[i] = ''
+      } else if (typeof value === 'object') {
+        const { json, isTruncated } = jsonStringify(value, maxChars)
+        if (isTruncated) truncated[i] = true
+        rowStrings[i] = json.replace(/\|/g, '\\|') // Escape pipe characters
+      } else {
+        let str = String(value)
+        if (str.length > maxChars) {
+          str = str.slice(0, maxChars) + '…'
+          truncated[i] = true
+        }
+        rowStrings[i] = str.replace(/\|/g, '\\|') // Escape pipe characters
+      }
+    }
+    rows.push(`| ${rowStrings.join(' | ')} |`)
+  }
+  // Show which columns were truncated in the header
+  const columnsWithTruncation = columns
+    .map((col, idx) => truncated[idx] ? `${col} (truncated)` : col)
+  const header = `| ${columnsWithTruncation.join(' | ')} |`
+  const separator = `| ${columns.map(() => '---').join(' | ')} |`
+
+  return [header, separator, ...rows].join('\n')
+}
+
+/**
+ * JSON stringify with truncation.
+ * @param {unknown} obj
+ * @param {number} maxChars
+ * @returns {{json: string, isTruncated: boolean}}
+ */
+function jsonStringify(obj, maxChars) {
+  const fullJson = JSON.stringify(obj)
+  if (fullJson.length <= maxChars) {
+    return { json: fullJson, isTruncated: false }
+  }
+  const sliced = fullJson.slice(0, maxChars) + '…'
+  return { json: sliced, isTruncated: true }
+}

bin/tools/markdownTable.js

@@ -0,0 +1,47 @@
+import { parquetPlan } from 'hyparquet/src/plan.js'
+import { asyncGroupToRows, readRowGroup } from 'hyparquet/src/rowgroup.js'
+import { whereToParquetFilter } from './parquetFilter.js'
+
+/**
+ * @import { AsyncBuffer, FileMetaData } from 'hyparquet'
+ * @import { AsyncDataSource, AsyncRow } from 'squirreling'
+ */
+
+/**
+ * Creates a parquet data source for use with squirreling SQL engine.
+ *
+ * @param {AsyncBuffer} file
+ * @param {FileMetaData} metadata
+ * @param {import('hyparquet-compressors').Compressors} compressors
+ * @returns {AsyncDataSource}
+ */
+export function parquetDataSource(file, metadata, compressors) {
+  return {
+    async *getRows(hints) {
+      const options = {
+        file,
+        metadata,
+        compressors,
+        columns: hints?.columns,
+        filter: whereToParquetFilter(hints?.where),
+      }
+      const plan = parquetPlan(options)
+      let count = 0
+      for (const subplan of plan.groups) {
+        const rg = readRowGroup(options, plan, subplan)
+        const rows = await asyncGroupToRows(rg, 0, rg.groupRows, undefined, 'object')
+        for (const asyncRow of rows) {
+          /** @type {AsyncRow} */
+          const row = {}
+          for (const [key, value] of Object.entries(asyncRow)) {
+            row[key] = () => Promise.resolve(value)
+          }
+          yield row
+          count++
+          // Check limit after each row
+          if (hints?.limit !== undefined && count >= hints.limit) return
+        }
+      }
+    },
+  }
+}