@botpress/zai 2.2.0 → 2.3.0

package/src/operations/sort.ts

@@ -0,0 +1,618 @@
+// eslint-disable consistent-type-definitions
+import { z } from '@bpinternal/zui'
+import pLimit from 'p-limit'
+import { ZaiContext } from '../context'
+import { Response } from '../response'
+import { getTokenizer } from '../tokenizer'
+import { fastHash, stringify } from '../utils'
+import { Zai } from '../zai'
+import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'
+
+export type Options = {
+  /** The maximum number of tokens per item */
+  tokensPerItem?: number
+}
+
+const _Options = z.object({
+  tokensPerItem: z
+    .number()
+    .min(1)
+    .max(100_000)
+    .optional()
+    .describe('The maximum number of tokens per item')
+    .default(250),
+})
+
+// Evaluation criteria generated by LLM
+type SortingCriteria = Record<
+  string,
+  {
+    description: string
+    labels: string[] // Ordered array of labels from FIRST to LAST, e.g., ['very_slow', 'slow', 'medium', 'fast', 'very_fast']
+  }
+>
+
+declare module '@botpress/zai' {
+  interface Zai {
+    /**
+     * Sorts an array of items based on provided instructions.
+     * Returns the sorted array directly when awaited.
+     * Use .result() to get detailed scoring information including why each item got its position.
+     *
+     * @example
+     * // Simple usage
+     * const sorted = await zai.sort(items, 'from least expensive to most expensive')
+     *
+     * @example
+     * // Get detailed results
+     * const { output: sorted, usage } = await zai.sort(items, 'by priority').result()
+     */
+    sort<T>(input: Array<T>, instructions: string, options?: Options): Response<Array<T>, Array<T>>
+  }
+}
+
+const END = '■END■'
+
+const sort = async <T>(
+  input: Array<T>,
+  instructions: string,
+  _options: Options | undefined,
+  ctx: ZaiContext
+): Promise<Array<T>> => {
+  ctx.controller.signal.throwIfAborted()
+
+  const options = _Options.parse(_options ?? {})
+  const tokenizer = await getTokenizer()
+  const model = await ctx.getModel()
+
+  const taskId = ctx.taskId
+  const taskType = 'zai.sort'
+
+  // Handle empty or single element arrays
+  if (input.length === 0) {
+    return []
+  }
+
+  if (input.length === 1) {
+    return input
+  }
+
+  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
+
+  // Phase 1: Generate sorting criteria from instructions + sample items
+  const sampleSize = Math.min(5, input.length)
+  const sampleItems = input.slice(0, sampleSize)
+  const sampleItemsText = sampleItems.map((item, idx) => `■${idx}: ${stringify(item, false)}`).join('\n')
+
+  const generateCriteriaPrompt = `Analyze this sorting instruction: "${instructions}"
+
+Sample items to be sorted:
+${sampleItemsText}
+
+Create 1-3 sorting criteria with ordered label arrays (3-10 labels each).
+
+**CRITICAL RULES**:
+1. Labels are single words, lowercase, no spaces, use underscores
+2. Labels are ordered from FIRST to LAST in sorted result
+3. If instruction says "from X to Y": first label represents X, last label represents Y
+4. If instruction says "prioritize" or "highest/lowest priority":
+   - First label = HIGHEST priority (top of todo list)
+   - Last label = LOWEST priority (bottom of todo list)
+
+Examples:
+
+"from slowest to fastest" → first=slowest, last=fastest
+■speed■
+very_slow;slow;medium;fast;very_fast
+■END■
+
+"from most dangerous to least dangerous" → first=most dangerous, last=least dangerous
+■danger■
+extremely_dangerous;very_dangerous;dangerous;moderate;slightly_dangerous;harmless
+■END■
+
+"from least urgent (spam) to most urgent (bills)" → first=spam, last=bills
+■urgency■
+spam;promotional;normal;important;urgent;critical
+■END■
+
+"prioritize: highest priority=open old tickets; lowest priority=closed" → first=high priority, last=low priority
+■status■
+open_old;open_recent;closed
+■age■
+oldest;old;recent;new
+■END■
+
+Output format:
+■criterion_name■
+label1;label2;label3;label4
+■END■
+
+Use 3-10 labels per criterion. Labels should be intuitive and match the domain.
+Keep criterion names short (1-2 words, lowercase, underscores).
+`
+
+  const { extracted: sortingCriteria } = await ctx.generateContent({
+    systemPrompt: `You are creating sorting criteria with ordered label arrays.
+
+CRITICAL: Output ordered labels from FIRST to LAST position in sorted result.
+- Labels are single words, lowercase, underscores only
+- 3-10 labels per criterion
+- Order matters: first label = appears first, last label = appears last`,
+    messages: [
+      {
+        type: 'text',
+        role: 'user',
+        content: generateCriteriaPrompt,
+      },
+    ],
+    transform: (text) => {
+      const criteria: SortingCriteria = {}
+      const criterionRegex = /■([^■]+)■\s*([^\n■]+)/g
+      let match: RegExpExecArray | null
+
+      while ((match = criterionRegex.exec(text)) !== null) {
+        const name = (match[1] ?? '').trim().toLowerCase()
+        const labelsStr = (match[2] ?? '').trim()
+
+        if (!name || name === 'end') continue
+
+        // Parse semicolon-separated labels
+        const labels = labelsStr
+          .split(';')
+          .map((l) => l.trim().toLowerCase().replace(/\s+/g, '_'))
+          .filter((l) => l.length > 0 && l.length < 50)
+
+        if (labels.length >= 3 && labels.length <= 10) {
+          criteria[name] = {
+            description: `${labels.length} ordered labels`,
+            labels,
+          }
+        }
+      }
+
+      if (Object.keys(criteria).length === 0) {
+        throw new Error(`Failed to parse sorting criteria. LLM output: ${text.slice(0, 500)}`)
+      }
+
+      return criteria
+    },
+  })
+
+  const criteriaKeys = Object.keys(sortingCriteria)
+  if (criteriaKeys.length === 0) {
+    throw new Error('No sorting criteria generated')
+  }
+
+  // Phase 2: Chunk items and score them in parallel
+  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.2)
+  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX
+
+  const MAX_ITEMS_PER_CHUNK = 50
+
+  // Prepare elements with indices
+  const elements = input.map((element, idx) => ({
+    element,
+    index: idx,
+    stringified: stringify(element, false),
+  }))
+
+  // Chunk elements
+  const chunks: Array<typeof elements> = []
+  let currentChunk: typeof elements = []
+  let currentTokens = 0
+
+  for (const elem of elements) {
+    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
+    const elemTokens = tokenizer.count(truncated)
+
+    if (
+      (currentTokens + elemTokens > TOKENS_ITEMS_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) &&
+      currentChunk.length > 0
+    ) {
+      chunks.push(currentChunk)
+      currentChunk = []
+      currentTokens = 0
+    }
+
+    currentChunk.push(elem)
+    currentTokens += elemTokens
+  }
+
+  if (currentChunk.length > 0) {
+    chunks.push(currentChunk)
+  }
+
+  // Phase 3: Score each chunk
+  type ItemScore = {
+    elementIndex: number
+    scores: Record<string, number>
+    totalScore: number
+  }
+
+  const scoreChunk = async (chunk: typeof elements): Promise<ItemScore[]> => {
+    ctx.controller.signal.throwIfAborted()
+
+    const chunkSize = chunk.length
+    const chunkInputStr = JSON.stringify(chunk.map((c) => c.element))
+
+    // Get examples from adapter for active learning
+    const examples =
+      taskId && ctx.adapter
+        ? await ctx.adapter.getExamples<string, ItemScore[]>({
+            input: chunkInputStr.slice(0, 1000),
+            taskType,
+            taskId,
+          })
+        : []
+
+    // Check for exact match (cache hit)
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: chunkInputStr,
+        instructions,
+      })
+    )
+
+    const exactMatch = examples.find((x) => x.key === key)
+    if (exactMatch && exactMatch.output) {
+      return exactMatch.output
+    }
+
+    const elementsText = chunk
+      .map((elem, i) => {
+        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
+        return `■${i}: ${truncated}■`
+      })
+      .join('\n')
+
+    const criteriaText = criteriaKeys
+      .map((key) => {
+        const criterion = sortingCriteria[key]
+        const labelsText = criterion.labels.join(';')
+        return `**${key}**: ${labelsText}`
+      })
+      .join('\n')
+
+    // Format examples for few-shot learning
+    const exampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []
+
+    for (const example of examples.slice(0, 3)) {
+      try {
+        const exampleInput = JSON.parse(example.input)
+        const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput]
+
+        exampleMessages.push({
+          type: 'text',
+          role: 'user',
+          content: `Expert Example - Items to score:\n${exampleItems.map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}■`).join('\n')}\n\nScore each item.`,
+        })
+
+        const exampleOutput = example.output
+        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
+          const formattedScores = exampleOutput
+            .map((score) => {
+              const pairs = criteriaKeys.map((key) => `${key}=${score.scores[key] ?? 0}`).join(';')
+              return `■${score.elementIndex}:${pairs}■`
+            })
+            .join('\n')
+
+          exampleMessages.push({
+            type: 'text',
+            role: 'assistant',
+            content: `${formattedScores}\n${END}`,
+          })
+
+          if (example.explanation) {
+            exampleMessages.push({
+              type: 'text',
+              role: 'assistant',
+              content: `Reasoning: ${example.explanation}`,
+            })
+          }
+        }
+      } catch {
+        // Skip malformed examples
+      }
+    }
+
+    const { extracted } = await ctx.generateContent({
+      systemPrompt: `You are ranking items for sorting using ordered label arrays.
+
+${criteriaText}
+
+Instructions: "${instructions}"
+
+SCORING RULES:
+- For each item and each criterion, assign ONE label from the ordered list
+- Labels are ordered: first label = appears FIRST in sorted result, last label = appears LAST
+- Choose the label that best describes each item
+
+Output format:
+■0:criterion1=label;criterion2=label■
+■1:criterion1=label;criterion2=label■
+${END}
+
+IMPORTANT:
+- Rank every item (■0 to ■${chunkSize - 1})
+- Use exact criterion names: ${criteriaKeys.join(', ')}
+- Use exact labels from the lists above (lowercase, underscores)
+- Use semicolons (;) between criteria
+- Use equals (=) between criterion and label`,
+      stopSequences: [END],
+      messages: [
+        ...exampleMessages,
+        {
+          type: 'text',
+          role: 'user',
+          content: `Items to rank (■0 to ■${chunkSize - 1}):\n${elementsText}\n\nRank each item using the labeled scales.\nOutput format: ■index:criterion1=label;criterion2=label■\n${END}`,
+        },
+      ],
+      transform: (text) => {
+        const results: ItemScore[] = []
+        const regex = /■(\d+):([^■]+)■/g
+        let match: RegExpExecArray | null
+
+        while ((match = regex.exec(text)) !== null) {
+          const idx = parseInt(match[1] ?? '', 10)
+          const labelsStr = match[2] ?? ''
+
+          if (isNaN(idx) || idx < 0 || idx >= chunkSize) continue
+
+          const scores: Record<string, number> = {}
+          let total = 0
+
+          const pairs = labelsStr.split(';').filter((x) => x.trim().length > 0)
+          for (const pair of pairs) {
+            const [criterion, labelStr] = pair.split('=').map((x) => x.trim().toLowerCase().replace(/\s+/g, '_'))
+            if (!criterion || !labelStr) continue
+
+            // Find the label index in the ordered array (index = score)
+            const labels = sortingCriteria[criterion]?.labels ?? []
+            const labelIndex = labels.findIndex((l) => l === labelStr)
+
+            if (labelIndex >= 0) {
+              // Use index as score (0 = first, higher = later)
+              scores[criterion] = labelIndex
+              total += labelIndex
+            } else {
+              // If label not found, use middle value
+              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5
+              scores[criterion] = middleIndex
+              total += middleIndex
+            }
+          }
+
+          results[idx] = {
+            elementIndex: chunk[idx].index,
+            scores,
+            totalScore: total,
+          }
+        }
+
+        // Fill in missing results with middle scores
+        for (let i = 0; i < chunkSize; i++) {
+          if (!results[i]) {
+            const scores: Record<string, number> = {}
+            let total = 0
+
+            for (const key of criteriaKeys) {
+              const labels = sortingCriteria[key]?.labels ?? []
+              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5
+              scores[key] = middleIndex
+              total += middleIndex
+            }
+
+            results[i] = {
+              elementIndex: chunk[i].index,
+              scores,
+              totalScore: total,
+            }
+          }
+        }
+
+        return results
+      },
+    })
+
+    return extracted
+  }
+
+  // Process all chunks in parallel
+  const limit = pLimit(10)
+  const chunkPromises = chunks.map((chunk) => limit(() => scoreChunk(chunk)))
+  const allScores = await Promise.all(chunkPromises)
+
+  // Phase 4: Merge scores from all chunks
+  // Build a map of elementIndex -> accumulated scores
+  const scoreMap = new Map<number, { scores: Record<string, number>; totalScore: number; tieBreakOrder?: number }>()
+
+  for (const chunkScores of allScores) {
+    for (const itemScore of chunkScores) {
+      const existing = scoreMap.get(itemScore.elementIndex)
+
+      if (existing) {
+        // Average the scores
+        for (const key of criteriaKeys) {
+          existing.scores[key] = (existing.scores[key] + (itemScore.scores[key] ?? 0)) / 2
+        }
+        existing.totalScore = (existing.totalScore + itemScore.totalScore) / 2
+      } else {
+        scoreMap.set(itemScore.elementIndex, {
+          scores: { ...itemScore.scores },
+          totalScore: itemScore.totalScore,
+        })
+      }
+    }
+  }
+
+  // Verify we have scores for all elements
+  if (scoreMap.size !== input.length) {
+    throw new Error(`Score map size mismatch: expected ${input.length}, got ${scoreMap.size}`)
+  }
+
+  // Phase 5: Identify ties
+  const scoreGroups = new Map<number, number[]>()
+  for (const [index, scoreData] of scoreMap.entries()) {
+    const roundedScore = Math.round(scoreData.totalScore * 100) // Round to avoid floating point issues
+    const group = scoreGroups.get(roundedScore) ?? []
+    group.push(index)
+    scoreGroups.set(roundedScore, group)
+  }
+
+  // Find groups with more than one item (ties)
+  const tiedGroups = Array.from(scoreGroups.values()).filter((group) => group.length > 1)
+
+  // Phase 6: Tie-breaking - process all tied groups IN PARALLEL
+  if (tiedGroups.length > 0) {
+    const tieBreakLimit = pLimit(10)
+
+    await Promise.all(
+      tiedGroups.map((tiedIndices) =>
+        tieBreakLimit(async () => {
+          if (tiedIndices.length <= 1) return
+
+          const tiedElements = tiedIndices.map((idx) => elements[idx])
+
+          // Re-score just these items for tie-breaking
+          const tieBreakText = tiedElements
+            .map((elem, i) => {
+              const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
+              return `■${i}: ${truncated}■`
+            })
+            .join('\n')
+
+          const { extracted: tieBreakOrder } = await ctx.generateContent({
+            systemPrompt: `You are breaking a tie between items with identical total scores.
+
+Instructions: ${instructions}
+
+Criteria:
+${criteriaKeys
+  .map((key) => {
+    const labels = sortingCriteria[key].labels.join(';')
+    return `- ${key}: ${labels}`
+  })
+  .join('\n')}
+
+Order these ${tiedElements.length} items from FIRST to LAST based on the instructions.
+Earlier labels in each criterion should come FIRST.
+
+Output format:
+■original_index■
+■original_index■
+${END}
+
+Output the indices in the order they should appear (first item at top).`,
+            stopSequences: [END],
+            messages: [
+              {
+                type: 'text',
+                role: 'user',
+                content: `Items with identical scores (need tie-breaking):\n${tieBreakText}\n\nOrder them from first to last.\nOutput format: ■index■ (one per line)\n${END}`,
+              },
+            ],
+            transform: (text) => {
+              const order: number[] = []
+              const regex = /■(\d+)■/g
+              let match: RegExpExecArray | null
+
+              while ((match = regex.exec(text)) !== null) {
+                const idx = parseInt(match[1] ?? '', 10)
+                if (!isNaN(idx) && idx >= 0 && idx < tiedElements.length) {
+                  order.push(idx)
+                }
+              }
+
+              // If not all items were ordered, append missing ones
+              for (let i = 0; i < tiedElements.length; i++) {
+                if (!order.includes(i)) {
+                  order.push(i)
+                }
+              }
+
+              return order
+            },
+          })
+
+          // Update scoreMap with tie-break order
+          for (let i = 0; i < tieBreakOrder.length; i++) {
+            const elementIndex = tiedElements[tieBreakOrder[i]].index
+            const scoreData = scoreMap.get(elementIndex)
+            if (scoreData) {
+              scoreData.tieBreakOrder = i
+            }
+          }
+        })
+      )
+    )
+  }
+
+  // Phase 7: Sort by total score, then by tie-break order
+  const sorted = Array.from(scoreMap.entries())
+    .sort((a, b) => {
+      // First sort by total score
+      const scoreDiff = a[1].totalScore - b[1].totalScore
+      if (scoreDiff !== 0) return scoreDiff
+
+      // If scores are equal, use tie-break order
+      const orderA = a[1].tieBreakOrder ?? 0
+      const orderB = b[1].tieBreakOrder ?? 0
+      return orderA - orderB
+    })
+    .map(([index]) => elements[index].element)
+
+  const result = sorted
+
+  // Save example for active learning
+  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
+    const key = fastHash(
+      stringify({
+        taskId,
+        taskType,
+        input: JSON.stringify(input),
+        instructions,
+      })
+    )
+
+    await ctx.adapter.saveExample({
+      key,
+      taskType,
+      taskId,
+      input: JSON.stringify(input),
+      output: result,
+      instructions,
+      metadata: {
+        cost: { input: 0, output: 0 },
+        latency: 0,
+        model: ctx.modelId,
+        tokens: { input: 0, output: 0 },
+      },
+    })
+  }
+
+  return result
+}
+
+Zai.prototype.sort = function <T>(
+  this: Zai,
+  input: Array<T>,
+  instructions: string,
+  _options?: Options
+): Response<Array<T>, Array<T>> {
+  const context = new ZaiContext({
+    client: this.client,
+    modelId: this.Model,
+    taskId: this.taskId,
+    taskType: 'zai.sort',
+    adapter: this.adapter,
+  })
+
+  return new Response<Array<T>, Array<T>>(
+    context,
+    sort(input, instructions, _options, context),
+    (result) => result // Simplified form is just the sorted array
+  )
+}