Prompt data interpolation functions

Author: csansoon

.changeset/quick-ways-report.md

@@ -0,0 +1,5 @@
+---
+"@latitude-data/llm-manager": minor
+---
+
+Added `table`, `json` and `csv` functions to print query results into the query in different formats.

packages/llm_manager/package.json

@@ -19,11 +19,13 @@
     "@latitude-data/source-manager": "workspace:*",
     "@latitude-data/sql-compiler": "workspace:^",
     "dotenv": "^16.4.5",
+    "json-2-csv": "^5.5.1",
     "openai": "^4.52.0",
     "yaml": "^2.3.4"
   },
   "devDependencies": {
     "@latitude-data/eslint-config": "workspace:*",
+    "@latitude-data/query_result": "workspace:*",
     "@latitude-data/typescript": "workspace:*",
     "@rollup/plugin-typescript": "^11.1.6",
     "@types/mock-fs": "^4.13.4",

packages/llm_manager/src/model/supportedMethods/csv.ts

@@ -0,0 +1,34 @@
+import { BuildSupportedMethodsArgs } from '$/types'
+import {
+  emptyMetadata,
+  type SupportedMethod,
+} from '@latitude-data/sql-compiler'
+import { QueryResultArray } from '@latitude-data/query_result'
+import { json2csv } from 'json-2-csv'
+
+const buildCastMethod = (_: BuildSupportedMethodsArgs): SupportedMethod => ({
+  requirements: {
+    interpolationPolicy: 'require', // Results can be interpolated
+    interpolationMethod: 'raw', // When interpolating, use the raw value
+    requireStaticArguments: false, // Can be used with variables or logic expressions
+  },
+
+  resolve: async (value: QueryResultArray) => {
+    // Check value is the correct type (array of objects)
+    if (!Array.isArray(value) || !value.every((v) => typeof v === 'object')) {
+      throw new Error('Value must be a query result.')
+    }
+
+    return json2csv(value, {
+      keys: value[0] ? Object.keys(value[0]) : [],
+      expandArrayObjects: false,
+      expandNestedObjects: false,
+    })
+  },
+
+  readMetadata: async () => {
+    return emptyMetadata()
+  },
+})
+
+export default buildCastMethod

packages/llm_manager/src/model/supportedMethods/index.ts

@@ -6,6 +6,9 @@ import { default as buildParam } from './param'
 import { default as buildRef } from './ref'
 import { default as buildRunQuery } from './runQuery'
 import { default as buildReadQuery } from './readQuery'
+import { default as table } from './table'
+import { default as json } from './json'
+import { default as csv } from './csv'
 
 export default function buildSupportedMethods(
   args: BuildSupportedMethodsArgs,
@@ -16,5 +19,8 @@ export default function buildSupportedMethods(
     runQuery: buildRunQuery(args),
     cast: buildCast(args),
     readQuery: buildReadQuery(args),
+    table: table(args),
+    json: json(args),
+    csv: csv(args),
   }
 }

packages/llm_manager/src/model/supportedMethods/json.ts

@@ -0,0 +1,40 @@
+import { BuildSupportedMethodsArgs } from '$/types'
+import {
+  emptyMetadata,
+  type SupportedMethod,
+} from '@latitude-data/sql-compiler'
+import { QueryResultArray } from '@latitude-data/query_result'
+
+function removeBigInts(results: QueryResultArray) {
+  results.forEach((row, i) => {
+    Object.entries(row).forEach(([key, value]) => {
+      if (typeof value === 'bigint') {
+        results[i]![key] = Number(value)
+      }
+    })
+  })
+}
+
+const buildCastMethod = (_: BuildSupportedMethodsArgs): SupportedMethod => ({
+  requirements: {
+    interpolationPolicy: 'require', // Results can be interpolated
+    interpolationMethod: 'raw', // When interpolating, use the raw value
+    requireStaticArguments: false, // Can be used with variables or logic expressions
+  },
+
+  resolve: async (value: QueryResultArray) => {
+    // Check value is the correct type (array of objects)
+    if (!Array.isArray(value) || !value.every((v) => typeof v === 'object')) {
+      throw new Error('Value must be a query result.')
+    }
+
+    removeBigInts(value)
+    return JSON.stringify(value)
+  },
+
+  readMetadata: async () => {
+    return emptyMetadata()
+  },
+})
+
+export default buildCastMethod

packages/llm_manager/src/model/supportedMethods/ref.ts

@@ -43,7 +43,7 @@ const buildRefMethod = ({
       refContext,
     )
 
-    return `(${compiledSubPrompt.prompt})`
+    return compiledSubPrompt.prompt
   },
 
   readMetadata: async () => {

packages/llm_manager/src/model/supportedMethods/table.ts

@@ -0,0 +1,74 @@
+import { BuildSupportedMethodsArgs } from '$/types'
+import {
+  emptyMetadata,
+  type SupportedMethod,
+} from '@latitude-data/sql-compiler'
+import { QueryResultArray } from '@latitude-data/query_result'
+
+type Col = {
+  name: string
+  values: string[]
+  maxLength: number
+}
+
+function padSpaces(str: string, length: number) {
+  return str + ' '.repeat(Math.max(0, length - str.length))
+}
+
+const buildCastMethod = (_: BuildSupportedMethodsArgs): SupportedMethod => ({
+  requirements: {
+    interpolationPolicy: 'require', // Results can be interpolated
+    interpolationMethod: 'raw', // When interpolating, use the raw value
+    requireStaticArguments: false, // Can be used with variables or logic expressions
+  },
+
+  resolve: async (value: QueryResultArray) => {
+    // Check value is the correct type (array of objects)
+    if (!Array.isArray(value) || !value.every((v) => typeof v === 'object')) {
+      throw new Error('Value must be a query result.')
+    }
+
+    const cols: Col[] = []
+
+    value.forEach((row, i) => {
+      Object.entries(row).forEach(([col, val]) => {
+        let colIndex = cols.findIndex((c) => c.name === col)
+        if (colIndex === -1) {
+          colIndex = cols.length
+          cols.push({
+            name: col,
+            values: Array(value.length).fill(''), // Fill with empty strings
+            maxLength: 0,
+          })
+        }
+
+        const valStr = String(val).replace(/\n/g, '\\n')
+        cols[colIndex]!.values[i] = valStr
+        
+        if (valStr.length > cols[colIndex]!.maxLength) {
+          cols[colIndex]!.maxLength = valStr.length
+        }
+      })
+    })
+
+    let table = ''
+
+    // Add header row
+    table += '| ' + cols.map(c => padSpaces(c.name, c.maxLength)).join(' | ') + ' |\n'
+    table += '|-' + cols.map(c => '-'.repeat(c.maxLength)).join('-|-') + '-|\n'
+
+    // Add data rows
+    value.forEach((_, i) => {
+      const cells = cols.map((c) => padSpaces(c.values[i]!, c.maxLength)) 
+      table += '| ' + cells.join(' | ') + ' |\n'
+    })
+
+    return table
+  },
+
+  readMetadata: async () => {
+    return emptyMetadata()
+  },
+})
+
+export default buildCastMethod