data queries

THIRD_PARTY_LICENSES.md

@@ -7091,6 +7091,30 @@ THIS SOFTWARE.
 
 -----------
 
+The following npm package may be included in this product:
+
+ - [email protected]
+
+This package contains the following license:
+
+ISC License
+
+Copyright (c) 2020, Stefan Terdell
+
+Permission to use, copy, modify, and/or distribute this software for any
+purpose with or without fee is hereby granted, provided that the above
+copyright notice and this permission notice appear in all copies.
+
+THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
+WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
+MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
+ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
+WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
+ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
+OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
+
+-----------
+
 The following npm package may be included in this product:
 
  - [email protected]

docs/src/content/docs/reference/scripts/context.md

@@ -257,6 +257,7 @@
 - `sliceTail`, number of rows to include from the end
 - `sliceSample`, number of rows to pick at random
 - `distinct`, list of column names to deduplicate the data based on
+- `query`, a [jq](https://jqlang.github.io/jq/) query to filter the data
 
 ```js
 defData("DATA", data, {

docs/src/content/docs/reference/scripts/parsers.md

@@ -16,7 +16,7 @@
 In general, parsing a JSON file as JSON5 does not cause harm, but it might be more forgiving
 to syntactic errors. In addition to JSON5, [JSON repair](https://www.npmjs.com/package/jsonrepair) is applied if the initial parse fails.
 
--   JSON5 example
+- JSON5 example
 
 ```json5
 {
@@ -272,8 +272,8 @@
 Parses error, warning annotations in various formats
 into a list of objects.
 
--   [GitHub Actions](https://docs.github.com/en/actions/using-workflows/workflow-commands-for-github-actions)
--   [Azure DevOps Pipeline](https://learn.microsoft.com/en-us/azure/devops/pipelines/scripts/logging-commands?view=azure-devops&tabs=bash#example-log-a-warning-about-a-specific-place-in-a-file)
+- [GitHub Actions](https://docs.github.com/en/actions/using-workflows/workflow-commands-for-github-actions)
+- [Azure DevOps Pipeline](https://learn.microsoft.com/en-us/azure/devops/pipelines/scripts/logging-commands?view=azure-devops&tabs=bash#example-log-a-warning-about-a-specific-place-in-a-file)
 -
 
 ```js
@@ -322,6 +322,14 @@
 const d = parsers.tidyData(rows, { sliceSample: 100, sort: "name" })
 ```
 
+## jq
+
+Apply a [jq](https://jqlang.github.io/jq/) query to a JSON object.
+
+```js
+const d = parsers.jq(rows, "map({ a })")
+```
+
 ## hash
 
 Utility to hash an object, array into a string that is appropriate for hashing purposes.

packages/core/package.json

@@ -21,6 +21,7 @@
     "@anthropic-ai/bedrock-sdk": "^0.12.0",
     "@anthropic-ai/sdk": "^0.33.1",
     "@azure/identity": "^4.5.0",
+    "@elastic/micro-jq": "^1.8.0",
     "@huggingface/jinja": "^0.3.2",
     "@huggingface/transformers": "^3.2.1",
     "@modelcontextprotocol/sdk": "^1.0.4",
@@ -60,6 +61,7 @@
     "inflection": "^3.0.0",
     "ini": "^5.0.0",
     "jimp": "^1.6.0",
+    "jqts": "^0.0.8",
     "json5": "^2.2.3",
     "jsonrepair": "^3.11.2",
     "magic-string": "^0.30.17",

packages/core/src/jq.test.ts

@@ -0,0 +1,45 @@
+import { jq } from "./jq"
+import { describe, test } from "node:test"
+import assert from "node:assert/strict"
+
+describe("jq", () => {
+    test("returns undefined when input is undefined", () => {
+        const result = jq(undefined, ".")
+        assert.strictEqual(result, undefined)
+    })
+
+    test("applies JQ transformation to input data", () => {
+        const input = { name: "John", age: 30 }
+        const query = ".name"
+        const result = jq(input, query)
+        assert.strictEqual(result, "John")
+    })
+
+    test("handles nested objects correctly", () => {
+        const input = { person: { name: "John", age: 30 } }
+        const query = ".person.name"
+        const result = jq(input, query)
+        assert.strictEqual(result, "John")
+    })
+
+    test("returns null for non-existent keys", () => {
+        const input = { name: "John", age: 30 }
+        const query = ".address"
+        const result = jq(input, query)
+        assert.strictEqual(result, null)
+    })
+
+    test("handles arrays correctly", () => {
+        const input = { people: [{ name: "John" }, { name: "Jane" }] }
+        const query = ".people[1].name"
+        const result = jq(input, query)
+        assert.strictEqual(result, "Jane")
+    })
+
+    test("returns entire input when query is '.'", () => {
+        const input = { name: "John", age: 30 }
+        const query = "."
+        const result = jq(input, query)
+        assert.deepStrictEqual(result, input)
+    })
+})

packages/core/src/jq.ts

@@ -0,0 +1,13 @@
+import _jq from "jqts"
+
+/**
+ * Loads and applies JQ transformation to the input data
+ * @param input
+ */
+export function jq(input: any, query: string): any {
+    if (input === undefined) return input
+
+    const pattern = _jq.compile(query)
+    const res = pattern.evaluate(input)[0]
+    return res
+}

packages/core/src/parsers.ts

@@ -28,6 +28,7 @@ import { jinjaRender } from "./jinja"
 import { createDiff, llmifyDiff } from "./diff"
 import { tidyData } from "./tidy"
 import { hash } from "./crypto"
+import { jq } from "./jq"
 
 export async function createParsers(options: {
     trace: MarkdownTrace
@@ -122,5 +123,6 @@ export async function createParsers(options: {
         tidyData: (rows, options) => tidyData(rows, options),
         hash: async (text, options) => await hash(text, options),
         unfence: unfence,
+        jq: jq,
     })
 }

packages/core/src/promptdom.ts

@@ -42,6 +42,7 @@ import { runtimeHost } from "./host"
 import { hash } from "./crypto"
 import { startMcpServer } from "./mcp"
 import { tryZodToJsonSchema } from "./zod"
+import { jq } from "./jq"
 
 // Definition of the PromptNode interface which is an essential part of the code structure.
 export interface PromptNode extends ContextExpansionOptions {
@@ -57,6 +58,7 @@ export interface PromptNode extends ContextExpansionOptions {
         | "assistant"
         | "system"
         | "def"
+        | "defData"
         | "chatParticipant"
         | "fileOutput"
         | "importTemplate"
@@ -98,6 +100,13 @@ export interface PromptDefNode extends PromptNode, DefOptions {
     resolved?: WorkspaceFile // Resolved file content
 }
 
+export interface PromptDefDataNode extends PromptNode, DefDataOptions {
+    type: "defData"
+    name: string // Name of the definition
+    value: Awaitable<object | object[]> // Data associated with the definition
+    resolved?: object | object[]
+}
+
 export interface PromptPrediction {
     type: "content"
     content: string
@@ -317,6 +326,48 @@ function renderDefNode(def: PromptDefNode): string {
     return res
 }
 
+function renderDefDataNode(n: PromptDefDataNode): string {
+    const { name, headers, priority, ephemeral, query } = n
+    let data = n.resolved
+    let format = n.format
+    const cacheControl = n.cacheControl ?? (ephemeral ? "ephemeral" : undefined)
+    if (
+        !format &&
+        Array.isArray(data) &&
+        data.length &&
+        (headers?.length || haveSameKeysAndSimpleValues(data))
+    )
+        format = "csv"
+    else if (!format) format = "yaml"
+
+    if (Array.isArray(data)) data = tidyData(data as object[], n)
+    if (query) data = jq(data, query)
+
+    let text: string
+    let lang: string
+    if (Array.isArray(data) && format === "csv") {
+        text = CSVToMarkdown(data)
+    } else if (format === "json") {
+        text = JSON.stringify(data)
+        lang = "json"
+    } else {
+        text = YAMLStringify(data)
+        lang = "yaml"
+    }
+
+    const value = lang
+        ? `${name}:
+\`\`\`${lang}
+${trimNewlines(text)}
+\`\`\`
+`
+        : `${name}:
+${trimNewlines(text)}
+`
+    // TODO maxTokens does not work well with data
+    return value
+}
+
 // Function to create an assistant node.
 export function createAssistantNode(
     value: Awaitable<string>,
@@ -468,50 +519,16 @@ function haveSameKeysAndSimpleValues(data: object[]): boolean {
 // Function to create a text node with data.
 export function createDefData(
     name: string,
-    data: object | object[],
+    value: Awaitable<object | object[]>,
     options?: DefDataOptions
-) {
-    if (data === undefined) return undefined
-    let { format, headers, priority, cacheControl } = options || {}
-    cacheControl =
-        cacheControl ?? (options?.ephemeral ? "ephemeral" : undefined)
-    if (
-        !format &&
-        Array.isArray(data) &&
-        data.length &&
-        (headers?.length || haveSameKeysAndSimpleValues(data))
-    )
-        format = "csv"
-    else if (!format) format = "yaml"
-
-    if (Array.isArray(data)) data = tidyData(data as object[], options)
-
-    let text: string
-    let lang: string
-    if (Array.isArray(data) && format === "csv") {
-        text = CSVToMarkdown(data)
-    } else if (format === "json") {
-        text = JSON.stringify(data)
-        lang = "json"
-    } else {
-        text = YAMLStringify(data)
-        lang = "yaml"
+): PromptDefDataNode {
+    if (value === undefined) return undefined
+    return {
+        type: "defData",
+        name,
+        value,
+        ...(options || {}),
     }
-
-    const value = lang
-        ? `${name}:
-\`\`\`${lang}
-${trimNewlines(text)}
-\`\`\`
-`
-        : `${name}:
-${trimNewlines(text)}
-`
-    // TODO maxTokens does not work well with data
-    return createTextNode(value, {
-        priority,
-        ephemeral: cacheControl === "ephemeral",
-    })
 }
 
 // Function to append a child node to a parent node.
@@ -529,6 +546,7 @@ export interface PromptNodeVisitor {
     afterNode?: (node: PromptNode) => Awaitable<void> // Post node visitor
     text?: (node: PromptTextNode) => Awaitable<void> // Text node visitor
     def?: (node: PromptDefNode) => Awaitable<void> // Definition node visitor
+    defData?: (node: PromptDefDataNode) => Awaitable<void> // Definition data node visitor
     image?: (node: PromptImageNode) => Awaitable<void> // Image node visitor
     schema?: (node: PromptSchemaNode) => Awaitable<void> // Schema node visitor
     tool?: (node: PromptToolNode) => Awaitable<void> // Function node visitor
@@ -553,6 +571,9 @@ export async function visitNode(node: PromptNode, visitor: PromptNodeVisitor) {
         case "def":
             await visitor.def?.(node as PromptDefNode)
             break
+        case "defData":
+            await visitor.defData?.(node as PromptDefDataNode)
+            break
         case "image":
             await visitor.image?.(node as PromptImageNode)
             break
@@ -666,6 +687,19 @@ async function resolvePromptNode(
                 n.error = e
             }
         },
+        defData: async (n) => {
+            try {
+                names.add(n.name)
+                const value = await n.value
+                n.resolved = value
+                const rendered = renderDefDataNode(n)
+                n.preview = rendered
+                n.tokens = estimateTokens(rendered, encoder)
+                n.children = [createTextNode(rendered)]
+            } catch (e) {
+                n.error = e
+            }
+        },
         system: async (n) => {
             try {
                 const value = await n.value

packages/core/src/types/prompt_template.d.ts

@@ -1568,6 +1568,7 @@ interface Parsers {
      * @param data data to render
      */
     mustache(text: string | WorkspaceFile, data: Record<string, any>): string
+
     /**
      * Renders a jinja template
      */
@@ -1589,6 +1590,13 @@ interface Parsers {
      */
     tidyData(rows: object[], options?: DataFilter): object[]
 
+    /**
+     * Applies a jq query to the data
+     * @param data data object to filter
+     * @param query jq query
+     */
+    jq(data: any, query: string): any
+
     /**
      * Computes a sha1 that can be used for hashing purpose, not cryptographic.
      * @param content content to hash
@@ -2411,6 +2419,11 @@ interface DefDataOptions
      * Output format in the prompt. Defaults to Markdown table rendering.
      */
     format?: "json" | "yaml" | "csv"
+
+    /**
+     * jq query to filter the data
+     */
+    query?: string
 }
 
 interface DefSchemaOptions {
@@ -2560,7 +2573,7 @@ interface ChatTurnGenerationContext {
     ): string
     defData(
         name: string,
-        data: object[] | object,
+        data: Awaitable<object[] | object>,
         options?: DefDataOptions
     ): string
     defDiff<T extends string | WorkspaceFile>(

packages/core/src/types/prompt_type.d.ts

@@ -267,7 +267,7 @@ declare function defImages(
  */
 declare function defData(
     name: string,
-    data: object[] | object,
+    data: Awaitable<object[] | object>,
     options?: DefDataOptions
 ): string