From 8fdde66f3018b6cc862826641d2b4c91a7bb5a5b Mon Sep 17 00:00:00 2001
From: Peli de Halleux <pelikhan@users.noreply.github.com>
Date: Mon, 23 Dec 2024 06:41:09 +0000
Subject: [PATCH] add cot

---
 packages/core/src/groq.ts                    |  5 ++-
 packages/core/src/promptdom.ts               |  6 +--
 packages/core/src/types/prompt_template.d.ts |  2 +-
 packages/sample/genaisrc/groq.genai.mjs      | 47 +++++++-------------
 4 files changed, 22 insertions(+), 38 deletions(-)

diff --git a/packages/core/src/groq.ts b/packages/core/src/groq.ts
index 8d1038778..a4c7133f0 100644
--- a/packages/core/src/groq.ts
+++ b/packages/core/src/groq.ts
@@ -3,10 +3,11 @@ import { parse, evaluate } from "groq-js"
  * Loads and applies JQ transformation to the input data
  * @param input
  */
-export function GROQEvaluate(query: string, dataset: any): any {
+export async function GROQEvaluate(query: string, dataset: any): any {
     if (dataset === undefined) return dataset
 
     const tree = parse(query)
-    const res = evaluate(tree, { dataset })
+    const value = await evaluate(tree, { dataset })
+    const res = await value.get()
     return res
 }
diff --git a/packages/core/src/promptdom.ts b/packages/core/src/promptdom.ts
index 2c2c02783..7d3a0ef9a 100644
--- a/packages/core/src/promptdom.ts
+++ b/packages/core/src/promptdom.ts
@@ -326,7 +326,7 @@ function renderDefNode(def: PromptDefNode): string {
     return res
 }
 
-function renderDefDataNode(n: PromptDefDataNode): string {
+async function renderDefDataNode(n: PromptDefDataNode): string {
     const { name, headers, priority, ephemeral, query } = n
     let data = n.resolved
     let format = n.format
@@ -341,7 +341,7 @@ function renderDefDataNode(n: PromptDefDataNode): string {
     else if (!format) format = "yaml"
 
     if (Array.isArray(data)) data = tidyData(data as object[], n)
-    if (query) data = GROQEvaluate(query, data)
+    if (query) data = await GROQEvaluate(query, data)
 
     let text: string
     let lang: string
@@ -692,7 +692,7 @@ async function resolvePromptNode(
                 names.add(n.name)
                 const value = await n.value
                 n.resolved = value
-                const rendered = renderDefDataNode(n)
+                const rendered = await renderDefDataNode(n)
                 n.preview = rendered
                 n.tokens = estimateTokens(rendered, encoder)
                 n.children = [createTextNode(rendered)]
diff --git a/packages/core/src/types/prompt_template.d.ts b/packages/core/src/types/prompt_template.d.ts
index 1680a558e..7f8051ef5 100644
--- a/packages/core/src/types/prompt_template.d.ts
+++ b/packages/core/src/types/prompt_template.d.ts
@@ -1596,7 +1596,7 @@ interface Parsers {
      * @param query query
      * @see https://groq.dev/
      */
-    GROQ(query: string, data: any): any
+    GROQ(query: string, data: any): Promise<any>
 
     /**
      * Computes a sha1 that can be used for hashing purpose, not cryptographic.
diff --git a/packages/sample/genaisrc/groq.genai.mjs b/packages/sample/genaisrc/groq.genai.mjs
index ec9697865..cb54c1925 100644
--- a/packages/sample/genaisrc/groq.genai.mjs
+++ b/packages/sample/genaisrc/groq.genai.mjs
@@ -1207,41 +1207,24 @@ const json = `[
 ]`
 const query = "filter to keep completed tasks and userid 2"
 
-let { text: jqq } = await runPrompt(
+const res = await runPrompt(
     (ctx) => {
-        ctx.$`Based on the example <JSON> snippet and the desired <QUERY>, write a jq program
+        ctx.$`## Task
+Based on the example <DATASET> snippet and the desired <QUERY>, write a GROQ program that executes the query.
 
-Return only the GROQ (https://groq.dev/) query to be executed as a raw string, no string delimiters
-wrapping it, no yapping, no markdown, no fenced code blocks, what you return
-will be passed to GROQ directly.
-For example, if the user asks: extract the name of the first person
-You return only: .people[0].name`.role("system")
+- Infer the JSON Schema of <DATASET> and use valid field names in the GROQ query
+- The dataset does not specify types, do NOT use '_type' filters
+- Explain the query step by step
+- Emit the GROQ query in a groq code fence section.
+
+`.role("system")
         ctx.def("QUERY", query)
-        ctx.def("JSON", json, { maxTokens: 500 })
+        ctx.def("DATASET", json, { maxTokens: 500 })
     },
-    { system: [] }
+    { system: ["system", "system.output_markdown", "system.assistant"] }
 )
 
-try {
-    const resjq = parsers.GROQ(jqq, JSON.parse(json))
-    assistant(resjq.text)
-} catch (e) {
-    console.error(e)
-    const rescheck = await runPrompt(
-        (ctx) => {
-            ctx.$`Repair the GROQ <QUERY> using the <ERROR> message below.
-Return only the GROQ program to be executed as a raw string, no string delimiters
-wrapping it, no yapping, no markdown, no fenced code blocks, what you return
-will be passed to GROQ directly.
-`.role("system")
-            ctx.def("QUERY", jqq)
-            ctx.def("ERROR", e.message)
-        },
-        { system: [] }
-    )
-    jqq = rescheck.text
-}
-
-const resjq = parsers.GROQ(jqq, JSON.parse(json))
-
-assistant(resjq.text)
+const GROQ = res.fences.find(f => f.language === "groq").content
+console.log(GROQ)
+const resjq = await parsers.GROQ(GROQ, JSON.parse(json))
+console.log(resjq)