diff --git a/.changeset/quick-ways-report.md b/.changeset/quick-ways-report.md new file mode 100644 index 00000000..a88562b0 --- /dev/null +++ b/.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. diff --git a/packages/llm_manager/package.json b/packages/llm_manager/package.json index 334a0abb..e7eecde2 100644 --- a/packages/llm_manager/package.json +++ b/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", diff --git a/packages/llm_manager/src/model/supportedMethods/csv.ts b/packages/llm_manager/src/model/supportedMethods/csv.ts new file mode 100644 index 00000000..d80c4f3f --- /dev/null +++ b/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 diff --git a/packages/llm_manager/src/model/supportedMethods/index.ts b/packages/llm_manager/src/model/supportedMethods/index.ts index 74fb7211..b66709e6 100644 --- a/packages/llm_manager/src/model/supportedMethods/index.ts +++ b/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), } } diff --git a/packages/llm_manager/src/model/supportedMethods/json.ts b/packages/llm_manager/src/model/supportedMethods/json.ts new file mode 100644 index 00000000..0728d0bd --- /dev/null +++ b/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 diff --git a/packages/llm_manager/src/model/supportedMethods/ref.ts b/packages/llm_manager/src/model/supportedMethods/ref.ts index 5eaef42c..a0adaf89 100644 --- a/packages/llm_manager/src/model/supportedMethods/ref.ts +++ b/packages/llm_manager/src/model/supportedMethods/ref.ts @@ -43,7 +43,7 @@ const buildRefMethod = ({ refContext, ) - return `(${compiledSubPrompt.prompt})` + return compiledSubPrompt.prompt }, readMetadata: async () => { diff --git a/packages/llm_manager/src/model/supportedMethods/table.ts b/packages/llm_manager/src/model/supportedMethods/table.ts new file mode 100644 index 00000000..822af4ad --- /dev/null +++ b/packages/llm_manager/src/model/supportedMethods/table.ts @@ -0,0 +1,78 @@ +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 diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 554189de..0892908c 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1346,6 +1346,9 @@ importers: dotenv: specifier: ^16.4.5 version: 16.4.5 + json-2-csv: + specifier: ^5.5.1 + version: 5.5.1 openai: specifier: ^4.52.0 version: 4.52.0 @@ -1356,6 +1359,9 @@ importers: '@latitude-data/eslint-config': specifier: workspace:* version: link:../../tools/eslint-config + '@latitude-data/query_result': + specifier: workspace:* + version: link:../query_result '@latitude-data/typescript': specifier: workspace:* version: link:../../tools/typescript