wehub-resource-sync
972 lines
32 KiB
TypeScript
972 lines
32 KiB
TypeScript
// eslint-disable consistent-type-definitions
|
|
import { z } from '@bpinternal/zui'
|
|
import { clamp } from 'lodash-es'
|
|
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 Group<T> = {
|
|
id: string
|
|
label: string
|
|
elements: T[]
|
|
}
|
|
|
|
type InitialGroup = {
|
|
id: string
|
|
label: string
|
|
elements?: unknown[]
|
|
}
|
|
|
|
const _InitialGroup = z.object({
|
|
id: z.string().min(1).max(100),
|
|
label: z.string().min(1).max(250),
|
|
elements: z.array(z.any()).optional().default([]),
|
|
})
|
|
|
|
export type Options = {
|
|
instructions?: string
|
|
tokensPerElement?: number
|
|
chunkLength?: number
|
|
initialGroups?: Array<InitialGroup>
|
|
maxGroups?: number
|
|
minElements?: number
|
|
}
|
|
|
|
const _Options = z.object({
|
|
instructions: z.string().optional(),
|
|
tokensPerElement: z.number().min(1).max(100_000).optional().default(250),
|
|
chunkLength: z.number().min(100).max(100_000).optional().default(16_000),
|
|
initialGroups: z.array(_InitialGroup).optional().default([]),
|
|
maxGroups: z.number().min(2).optional(),
|
|
minElements: z.number().min(1).optional(),
|
|
})
|
|
|
|
declare module '@botpress/zai' {
|
|
interface Zai {
|
|
/**
|
|
* Groups array items into categories based on semantic similarity or criteria.
|
|
*
|
|
* This operation intelligently categorizes items by analyzing their content and
|
|
* creating meaningful groups. It can discover natural groupings automatically or
|
|
* use predefined categories. Perfect for clustering, classification, and organization.
|
|
*
|
|
* @param input - Array of items to group
|
|
* @param options - Configuration for grouping behavior, instructions, and initial categories
|
|
* @param options.maxGroups - Maximum number of groups allowed (minimum 2). When set, groups are merged at the end until within limit.
|
|
* @param options.minElements - Minimum elements per group (minimum 1). Groups below this threshold have their elements redistributed via AI.
|
|
* @returns Response with groups array (simplified to Record<groupLabel, items[]>)
|
|
*
|
|
* @example Automatic grouping
|
|
* ```typescript
|
|
* const messages = [
|
|
* "I can't log in to my account",
|
|
* "How do I reset my password?",
|
|
* "When will my order arrive?",
|
|
* "The app keeps crashing",
|
|
* "I haven't received my package",
|
|
* "Error 500 on checkout"
|
|
* ]
|
|
*
|
|
* const groups = await zai.group(messages, {
|
|
* instructions: 'Group by type of customer issue'
|
|
* })
|
|
* // Result (simplified):
|
|
* // {
|
|
* // "Login Issues": ["I can't log in...", "How do I reset..."],
|
|
* // "Shipping Questions": ["When will my order...", "I haven't received..."],
|
|
* // "Technical Errors": ["The app keeps crashing", "Error 500..."]
|
|
* // }
|
|
*
|
|
* // Full result:
|
|
* const { output } = await zai.group(messages, { instructions: '...' }).result()
|
|
* // output: [
|
|
* // { id: 'login_issues', label: 'Login Issues', elements: [...] },
|
|
* // { id: 'shipping', label: 'Shipping Questions', elements: [...] },
|
|
* // { id: 'errors', label: 'Technical Errors', elements: [...] }
|
|
* // ]
|
|
* ```
|
|
*
|
|
* @example With predefined categories
|
|
* ```typescript
|
|
* const articles = [
|
|
* "How to build a React app",
|
|
* "Python machine learning tutorial",
|
|
* "Understanding Docker containers",
|
|
* "Vue.js best practices",
|
|
* "Deep learning with TensorFlow"
|
|
* ]
|
|
*
|
|
* const groups = await zai.group(articles, {
|
|
* instructions: 'Categorize by technology',
|
|
* initialGroups: [
|
|
* { id: 'frontend', label: 'Frontend Development' },
|
|
* { id: 'backend', label: 'Backend Development' },
|
|
* { id: 'ml', label: 'Machine Learning' },
|
|
* { id: 'devops', label: 'DevOps & Infrastructure' }
|
|
* ]
|
|
* })
|
|
* // Groups articles into predefined categories
|
|
* ```
|
|
*
|
|
* @example Grouping products
|
|
* ```typescript
|
|
* const products = [
|
|
* { name: 'Laptop', price: 999, category: 'Electronics' },
|
|
* { name: 'Desk', price: 299, category: 'Furniture' },
|
|
* { name: 'Mouse', price: 29, category: 'Electronics' },
|
|
* { name: 'Chair', price: 199, category: 'Furniture' }
|
|
* ]
|
|
*
|
|
* const grouped = await zai.group(products, {
|
|
* instructions: 'Group by price range: budget (< $100), mid-range ($100-$500), premium (> $500)'
|
|
* })
|
|
* // Result:
|
|
* // {
|
|
* // "Budget": [Mouse],
|
|
* // "Mid-range": [Chair, Desk],
|
|
* // "Premium": [Laptop]
|
|
* // }
|
|
* ```
|
|
*
|
|
* @example Content categorization
|
|
* ```typescript
|
|
* const emails = [
|
|
* { subject: 'Meeting tomorrow', body: '...', from: 'boss@company.com' },
|
|
* { subject: 'Invoice #1234', body: '...', from: 'billing@vendor.com' },
|
|
* { subject: 'Weekly report', body: '...', from: 'team@company.com' },
|
|
* { subject: 'Payment received', body: '...', from: 'accounting@company.com' }
|
|
* ]
|
|
*
|
|
* const categorized = await zai.group(emails, {
|
|
* instructions: 'Categorize by email type: work communication, financial, reports',
|
|
* tokensPerElement: 300 // Allow more context per email
|
|
* })
|
|
* ```
|
|
*
|
|
* @example Customer feedback grouping
|
|
* ```typescript
|
|
* const feedback = [
|
|
* "Love the new UI!",
|
|
* "App is too slow",
|
|
* "Great customer service",
|
|
* "Confusing navigation",
|
|
* "Fast shipping!",
|
|
* "Hard to find features"
|
|
* ]
|
|
*
|
|
* const grouped = await zai.group(feedback, {
|
|
* instructions: 'Group by aspect: UI/UX, Performance, Customer Service, Shipping'
|
|
* })
|
|
* ```
|
|
*
|
|
* @example Topic clustering for research
|
|
* ```typescript
|
|
* const papers = [
|
|
* { title: 'Transformer Networks for NLP', abstract: '...' },
|
|
* { title: 'CNN Image Classification', abstract: '...' },
|
|
* { title: 'BERT Language Understanding', abstract: '...' },
|
|
* { title: 'Object Detection with YOLO', abstract: '...' }
|
|
* ]
|
|
*
|
|
* const clusters = await zai.group(papers, {
|
|
* instructions: 'Group by research area',
|
|
* chunkLength: 10000 // Allow more tokens for detailed abstracts
|
|
* })
|
|
* // Result: Groups papers by topic (NLP, Computer Vision, etc.)
|
|
* ```
|
|
*
|
|
* @example With initial seed groups
|
|
* ```typescript
|
|
* const tasks = [
|
|
* "Fix login bug",
|
|
* "Update documentation",
|
|
* "Add dark mode",
|
|
* "Optimize database queries"
|
|
* ]
|
|
*
|
|
* const grouped = await zai.group(tasks, {
|
|
* instructions: 'Categorize development tasks',
|
|
* initialGroups: [
|
|
* { id: 'bugs', label: 'Bug Fixes', elements: [] },
|
|
* { id: 'features', label: 'New Features', elements: [] },
|
|
* { id: 'docs', label: 'Documentation', elements: [] },
|
|
* { id: 'performance', label: 'Performance', elements: [] }
|
|
* ]
|
|
* })
|
|
* ```
|
|
*/
|
|
/**
|
|
* @example Limiting number of groups
|
|
* ```typescript
|
|
* const items = ['apple', 'banana', 'carrot', 'chicken', 'rice', 'bread', 'salmon', 'milk']
|
|
*
|
|
* const groups = await zai.group(items, {
|
|
* instructions: 'Group by food type',
|
|
* maxGroups: 3 // At most 3 groups — smallest groups get merged if exceeded
|
|
* })
|
|
* // Guarantees no more than 3 groups in the result
|
|
* ```
|
|
*/
|
|
group<T>(input: Array<T>, options?: Options): Response<Array<Group<T>>, Record<string, T[]>>
|
|
}
|
|
}
|
|
|
|
const END = '■END■'
|
|
|
|
// Simplified data structures
|
|
type GroupInfo = {
|
|
id: string
|
|
label: string
|
|
normalizedLabel: string
|
|
}
|
|
|
|
const normalizeLabel = (label: string): string => {
|
|
return label
|
|
.trim()
|
|
.toLowerCase()
|
|
.replace(/^(group|new group|new)\s*[-:]\s*/i, '')
|
|
.replace(/^(group|new group|new)\s+/i, '')
|
|
.trim()
|
|
}
|
|
|
|
const group = async <T>(input: Array<T>, _options: Options | undefined, ctx: ZaiContext): Promise<Array<Group<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.group'
|
|
|
|
if (input.length === 0) {
|
|
return []
|
|
}
|
|
|
|
// Simple data structures
|
|
const groups = new Map<string, GroupInfo>() // groupId -> GroupInfo
|
|
const groupElements = new Map<string, Set<number>>() // groupId -> Set of element indices
|
|
const elementGroups = new Map<number, Set<string>>() // elementIndex -> Set of groupIds seen/assigned
|
|
const labelToGroupId = new Map<string, string>() // normalized label -> groupId
|
|
let groupIdCounter = 0
|
|
|
|
// Initialize with provided groups
|
|
options.initialGroups.forEach((ig) => {
|
|
const normalized = normalizeLabel(ig.label)
|
|
groups.set(ig.id, { id: ig.id, label: ig.label, normalizedLabel: normalized })
|
|
groupElements.set(ig.id, new Set())
|
|
labelToGroupId.set(normalized, ig.id)
|
|
})
|
|
|
|
// Prepare elements
|
|
const elements = input.map((element, idx) => ({
|
|
element,
|
|
index: idx,
|
|
stringified: stringify(element, false),
|
|
}))
|
|
|
|
// Token budget
|
|
const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER
|
|
const TOKENS_INSTRUCTIONS_MAX = options.instructions
|
|
? clamp(tokenizer.count(options.instructions), 100, TOKENS_TOTAL_MAX * 0.2)
|
|
: 0
|
|
const TOKENS_AVAILABLE = TOKENS_TOTAL_MAX - TOKENS_INSTRUCTIONS_MAX
|
|
const TOKENS_FOR_GROUPS_MAX = Math.floor(TOKENS_AVAILABLE * 0.4)
|
|
const TOKENS_FOR_ELEMENTS_MAX = Math.floor(TOKENS_AVAILABLE * 0.6)
|
|
|
|
// Chunk elements by token budget
|
|
const MAX_ELEMENTS_PER_CHUNK = 50
|
|
const elementChunks: number[][] = [] // Array of element indices
|
|
let currentChunk: number[] = []
|
|
let currentTokens = 0
|
|
|
|
for (const elem of elements) {
|
|
const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
|
|
const elemTokens = tokenizer.count(truncated)
|
|
|
|
if (
|
|
(currentTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || currentChunk.length >= MAX_ELEMENTS_PER_CHUNK) &&
|
|
currentChunk.length > 0
|
|
) {
|
|
elementChunks.push(currentChunk)
|
|
currentChunk = []
|
|
currentTokens = 0
|
|
}
|
|
|
|
currentChunk.push(elem.index)
|
|
currentTokens += elemTokens
|
|
}
|
|
|
|
if (currentChunk.length > 0) {
|
|
elementChunks.push(currentChunk)
|
|
}
|
|
|
|
// Helper to chunk groups
|
|
const getGroupChunks = (): string[][] => {
|
|
const allGroupIds = Array.from(groups.keys())
|
|
if (allGroupIds.length === 0) return [[]]
|
|
|
|
const chunks: string[][] = []
|
|
let currentChunk: string[] = []
|
|
let currentTokens = 0
|
|
|
|
for (const groupId of allGroupIds) {
|
|
const group = groups.get(groupId)!
|
|
const groupTokens = tokenizer.count(`${group.label}`) + 10
|
|
|
|
if (currentTokens + groupTokens > TOKENS_FOR_GROUPS_MAX && currentChunk.length > 0) {
|
|
chunks.push(currentChunk)
|
|
currentChunk = []
|
|
currentTokens = 0
|
|
}
|
|
|
|
currentChunk.push(groupId)
|
|
currentTokens += groupTokens
|
|
}
|
|
|
|
if (currentChunk.length > 0) {
|
|
chunks.push(currentChunk)
|
|
}
|
|
|
|
return chunks.length > 0 ? chunks : [[]]
|
|
}
|
|
|
|
// Process elements against groups and get assignments
|
|
const processChunk = async (
|
|
elementIndices: number[],
|
|
groupIds: string[]
|
|
): Promise<Array<{ elementIndex: number; label: string }>> => {
|
|
// Get examples from adapter for active learning
|
|
const chunkElements = elementIndices.map((idx) => elements[idx].element)
|
|
const chunkInputStr = JSON.stringify(chunkElements)
|
|
|
|
const examples =
|
|
taskId && ctx.adapter
|
|
? await ctx.adapter.getExamples<string, Array<{ elementIndex: number; label: string }>>({
|
|
input: chunkInputStr.slice(0, 1000), // Limit search string length
|
|
taskType,
|
|
taskId,
|
|
})
|
|
: []
|
|
|
|
// Check for exact match (cache hit)
|
|
|
|
const key = fastHash(
|
|
stringify({
|
|
taskId,
|
|
taskType,
|
|
input: chunkInputStr,
|
|
instructions: options.instructions ?? '',
|
|
groupIds: groupIds.join(','),
|
|
})
|
|
)
|
|
|
|
const exactMatch = examples.find((x) => x.key === key)
|
|
if (exactMatch && exactMatch.output) {
|
|
return exactMatch.output
|
|
}
|
|
|
|
const elementsText = elementIndices
|
|
.map((idx, i) => {
|
|
const elem = elements[idx]
|
|
const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
|
|
return `■${i}: ${truncated}■`
|
|
})
|
|
.join('\n')
|
|
|
|
const groupsList = groupIds.map((gid) => groups.get(gid)!.label)
|
|
const groupsText =
|
|
groupsList.length > 0
|
|
? `**Existing Groups (prefer reusing these):**\n${groupsList.map((l) => `- ${l}`).join('\n')}\n\n`
|
|
: ''
|
|
|
|
// Format examples for few-shot learning
|
|
const exampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []
|
|
|
|
for (const example of examples.slice(0, 5)) {
|
|
try {
|
|
const exampleInput = JSON.parse(example.input)
|
|
const exampleElements = Array.isArray(exampleInput) ? exampleInput : [exampleInput]
|
|
|
|
// User message
|
|
const exampleElementsText = exampleElements
|
|
.map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}■`)
|
|
.join('\n')
|
|
|
|
exampleMessages.push({
|
|
type: 'text',
|
|
role: 'user',
|
|
content: `Expert Example - Elements to group:
|
|
${exampleElementsText}
|
|
|
|
Group each element.`,
|
|
})
|
|
|
|
// Assistant message
|
|
const exampleOutput = example.output
|
|
if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
|
|
const formattedAssignments = exampleOutput
|
|
.map((assignment) => `■${assignment.elementIndex}:${assignment.label}■`)
|
|
.join('\n')
|
|
|
|
exampleMessages.push({
|
|
type: 'text',
|
|
role: 'assistant',
|
|
content: `${formattedAssignments}\n${END}`,
|
|
})
|
|
|
|
if (example.explanation) {
|
|
exampleMessages.push({
|
|
type: 'text',
|
|
role: 'assistant',
|
|
content: `Reasoning: ${example.explanation}`,
|
|
})
|
|
}
|
|
}
|
|
} catch {
|
|
// Skip malformed examples
|
|
}
|
|
}
|
|
|
|
const systemPrompt = `You are grouping elements into cohesive groups.
|
|
|
|
${options.instructions ? `**Instructions:** ${options.instructions}\n` : '**Instructions:** Group similar elements together.'}
|
|
|
|
**Important:**
|
|
- Each element gets exactly ONE group label
|
|
- Use EXACT SAME label for similar items (case-sensitive)
|
|
- Create new descriptive labels when needed
|
|
|
|
**Output Format:**
|
|
One line per element:
|
|
■0:Group Label■
|
|
■1:Group Label■
|
|
${END}`.trim()
|
|
|
|
const userPrompt = `${groupsText}**Elements (■0 to ■${elementIndices.length - 1}):**
|
|
${elementsText}
|
|
|
|
**Task:** For each element, output one line with its group label.
|
|
${END}`.trim()
|
|
|
|
const { extracted } = await ctx.generateContent({
|
|
systemPrompt,
|
|
stopSequences: [END],
|
|
messages: [...exampleMessages, { type: 'text', role: 'user', content: userPrompt }],
|
|
transform: (text) => {
|
|
const assignments: Array<{ elementIndex: number; label: string }> = []
|
|
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 >= elementIndices.length) continue
|
|
|
|
const label = (match[2] ?? '').trim()
|
|
if (!label) continue
|
|
|
|
assignments.push({
|
|
elementIndex: elementIndices[idx],
|
|
label: label.slice(0, 250),
|
|
})
|
|
}
|
|
|
|
return assignments
|
|
},
|
|
})
|
|
|
|
return extracted
|
|
}
|
|
|
|
// Phase 1: Process all element chunks against current groups IN PARALLEL
|
|
const elementLimit = pLimit(10) // Separate limiter for element chunks
|
|
const groupLimit = pLimit(10) // Separate limiter for group chunks
|
|
|
|
// Collect all assignments from parallel processing
|
|
const allChunkResults = await Promise.all(
|
|
elementChunks.map((elementChunk) =>
|
|
elementLimit(async () => {
|
|
const groupChunks = getGroupChunks()
|
|
|
|
const allAssignments = await Promise.all(
|
|
groupChunks.map((groupChunk) => groupLimit(() => processChunk(elementChunk, groupChunk)))
|
|
)
|
|
|
|
return allAssignments.flat()
|
|
})
|
|
)
|
|
)
|
|
|
|
// Process all assignments sequentially to avoid race conditions
|
|
for (const assignments of allChunkResults) {
|
|
for (const { elementIndex, label } of assignments) {
|
|
const normalized = normalizeLabel(label)
|
|
let groupId = labelToGroupId.get(normalized)
|
|
|
|
if (!groupId) {
|
|
// Create new group
|
|
groupId = `group_${groupIdCounter++}`
|
|
groups.set(groupId, { id: groupId, label, normalizedLabel: normalized })
|
|
groupElements.set(groupId, new Set())
|
|
labelToGroupId.set(normalized, groupId)
|
|
}
|
|
|
|
// Add element to group
|
|
groupElements.get(groupId)!.add(elementIndex)
|
|
|
|
// Track that element saw this group
|
|
if (!elementGroups.has(elementIndex)) {
|
|
elementGroups.set(elementIndex, new Set())
|
|
}
|
|
elementGroups.get(elementIndex)!.add(groupId)
|
|
}
|
|
}
|
|
|
|
// Phase 2: Ensure all elements saw all groups (coverage guarantee)
|
|
const allGroupIds = Array.from(groups.keys())
|
|
|
|
if (allGroupIds.length > 0) {
|
|
const elementsNeedingReview: number[] = []
|
|
|
|
for (const elem of elements) {
|
|
const seenGroups = elementGroups.get(elem.index) ?? new Set()
|
|
const unseenCount = allGroupIds.filter((gid) => !seenGroups.has(gid)).length
|
|
|
|
if (unseenCount > 0) {
|
|
elementsNeedingReview.push(elem.index)
|
|
}
|
|
}
|
|
|
|
if (elementsNeedingReview.length > 0) {
|
|
// Chunk elements needing review
|
|
const reviewChunks: number[][] = []
|
|
let reviewChunk: number[] = []
|
|
let reviewTokens = 0
|
|
|
|
for (const elemIdx of elementsNeedingReview) {
|
|
const elem = elements[elemIdx]
|
|
const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
|
|
const elemTokens = tokenizer.count(truncated)
|
|
|
|
const shouldStartNewChunk =
|
|
(reviewTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX || reviewChunk.length >= MAX_ELEMENTS_PER_CHUNK) &&
|
|
reviewChunk.length > 0
|
|
|
|
if (shouldStartNewChunk) {
|
|
reviewChunks.push(reviewChunk)
|
|
reviewChunk = []
|
|
reviewTokens = 0
|
|
}
|
|
|
|
reviewChunk.push(elemIdx)
|
|
reviewTokens += elemTokens
|
|
}
|
|
|
|
if (reviewChunk.length > 0) {
|
|
reviewChunks.push(reviewChunk)
|
|
}
|
|
|
|
// Process review chunks IN PARALLEL
|
|
const reviewResults = await Promise.all(
|
|
reviewChunks.map((chunk) =>
|
|
elementLimit(async () => {
|
|
const groupChunks = getGroupChunks()
|
|
|
|
const allAssignments = await Promise.all(
|
|
groupChunks.map((groupChunk) => groupLimit(() => processChunk(chunk, groupChunk)))
|
|
)
|
|
|
|
return allAssignments.flat()
|
|
})
|
|
)
|
|
)
|
|
|
|
// Mark groups as seen and update assignments (sequential to avoid races)
|
|
const updateElementGroupAssignment = (elementIndex: number, label: string) => {
|
|
const normalized = normalizeLabel(label)
|
|
const groupId = labelToGroupId.get(normalized)
|
|
if (!groupId) return
|
|
|
|
// Add to group and mark as seen
|
|
groupElements.get(groupId)!.add(elementIndex)
|
|
|
|
// Initialize element groups if needed
|
|
const elemGroups = elementGroups.get(elementIndex) ?? new Set()
|
|
if (!elementGroups.has(elementIndex)) {
|
|
elementGroups.set(elementIndex, elemGroups)
|
|
}
|
|
elemGroups.add(groupId)
|
|
}
|
|
|
|
for (const assignments of reviewResults) {
|
|
for (const { elementIndex, label } of assignments) {
|
|
updateElementGroupAssignment(elementIndex, label)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Phase 3: Resolve conflicts (elements in multiple groups)
|
|
for (const [elementIndex, groupSet] of elementGroups.entries()) {
|
|
if (groupSet.size > 1) {
|
|
// Element is in multiple groups, keep only the most common assignment
|
|
const groupIds = Array.from(groupSet)
|
|
|
|
// Remove from all groups
|
|
for (const gid of groupIds) {
|
|
groupElements.get(gid)?.delete(elementIndex)
|
|
}
|
|
|
|
// Re-assign to first group (or could use LLM to decide)
|
|
const finalGroupId = groupIds[0]
|
|
groupElements.get(finalGroupId)!.add(elementIndex)
|
|
}
|
|
}
|
|
|
|
// Phase 4: Merge groups if maxGroups is set (AI-driven)
|
|
if (options.maxGroups !== undefined) {
|
|
const nonEmptyGroupIds = () =>
|
|
Array.from(groupElements.entries())
|
|
.filter(([, s]) => s.size > 0)
|
|
.map(([id]) => id)
|
|
|
|
let currentIds = nonEmptyGroupIds()
|
|
|
|
if (currentIds.length > options.maxGroups) {
|
|
// Build a summary of each group: label + element count + sample elements
|
|
const groupSummaries = currentIds.map((gid, idx) => {
|
|
const info = groups.get(gid)!
|
|
const elemIndices = Array.from(groupElements.get(gid)!)
|
|
const sampleElements = elemIndices
|
|
.slice(0, 3)
|
|
.map((i) => tokenizer.truncate(elements[i].stringified, 60))
|
|
.join(', ')
|
|
return `■${idx}:${info.label} (${elemIndices.length} elements, e.g. ${sampleElements})■`
|
|
})
|
|
|
|
const mergeSystemPrompt = `You are consolidating groups into fewer, broader categories.
|
|
|
|
${options.instructions ? `**Original instructions:** ${options.instructions}\n` : ''}
|
|
**Task:** Merge ${currentIds.length} groups down to at most ${options.maxGroups} groups.
|
|
Combine the most semantically related groups together. Give each merged group a new descriptive label.
|
|
|
|
**Output Format:**
|
|
For each input group (■0 to ■${currentIds.length - 1}), output which target label it maps to:
|
|
■0:Merged Label■
|
|
■1:Merged Label■
|
|
${END}
|
|
|
|
Use the EXACT SAME label for groups that should be merged together.`.trim()
|
|
|
|
const mergeUserPrompt = `**Current groups:**
|
|
${groupSummaries.join('\n')}
|
|
|
|
Merge into at most ${options.maxGroups} groups.
|
|
${END}`.trim()
|
|
|
|
const { extracted: mergeAssignments } = await ctx.generateContent({
|
|
systemPrompt: mergeSystemPrompt,
|
|
stopSequences: [END],
|
|
messages: [{ type: 'text', role: 'user', content: mergeUserPrompt }],
|
|
transform: (text) => {
|
|
const assignments: Array<{ sourceIdx: number; label: string }> = []
|
|
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 >= currentIds.length) continue
|
|
|
|
const label = (match[2] ?? '').trim()
|
|
if (!label) continue
|
|
|
|
assignments.push({ sourceIdx: idx, label: label.slice(0, 250) })
|
|
}
|
|
|
|
return assignments
|
|
},
|
|
})
|
|
|
|
// Build merge map: normalized merge label → list of source group IDs
|
|
const mergeMap = new Map<string, { label: string; sourceGroupIds: string[] }>()
|
|
|
|
for (const { sourceIdx, label } of mergeAssignments) {
|
|
const sourceGid = currentIds[sourceIdx]
|
|
if (!sourceGid) continue
|
|
|
|
const normalized = normalizeLabel(label)
|
|
if (!mergeMap.has(normalized)) {
|
|
mergeMap.set(normalized, { label, sourceGroupIds: [] })
|
|
}
|
|
mergeMap.get(normalized)!.sourceGroupIds.push(sourceGid)
|
|
}
|
|
|
|
// Apply merges: for each merge target, pick the first source group as the target
|
|
// and move all elements from other source groups into it
|
|
for (const [, { label, sourceGroupIds }] of mergeMap) {
|
|
if (sourceGroupIds.length <= 1) continue
|
|
|
|
const targetGid = sourceGroupIds[0]
|
|
const targetSet = groupElements.get(targetGid)!
|
|
|
|
// Update label on the target group
|
|
const targetInfo = groups.get(targetGid)!
|
|
targetInfo.label = label
|
|
targetInfo.normalizedLabel = normalizeLabel(label)
|
|
|
|
for (let i = 1; i < sourceGroupIds.length; i++) {
|
|
const sourceGid = sourceGroupIds[i]
|
|
const sourceSet = groupElements.get(sourceGid)!
|
|
sourceSet.forEach((elemIdx) => targetSet.add(elemIdx))
|
|
sourceSet.clear()
|
|
}
|
|
}
|
|
|
|
// Safety: if LLM still produced too many groups, fall back to merging smallest pairs
|
|
currentIds = nonEmptyGroupIds()
|
|
while (currentIds.length > options.maxGroups) {
|
|
currentIds.sort((a, b) => groupElements.get(a)!.size - groupElements.get(b)!.size)
|
|
|
|
const sourceSet = groupElements.get(currentIds[0])!
|
|
const targetSet = groupElements.get(currentIds[1])!
|
|
for (const elemIdx of sourceSet) {
|
|
targetSet.add(elemIdx)
|
|
}
|
|
sourceSet.clear()
|
|
|
|
currentIds = nonEmptyGroupIds()
|
|
}
|
|
}
|
|
}
|
|
|
|
// Phase 5: Redistribute undersized groups if minElements is set
|
|
// Reuses processChunk so orphans see the valid groups as available buckets
|
|
if (options.minElements !== undefined && options.minElements > 1) {
|
|
const getNonEmptyGroupIds = () =>
|
|
Array.from(groupElements.entries())
|
|
.filter(([, s]) => s.size > 0)
|
|
.map(([id]) => id)
|
|
|
|
// Collect orphan elements from all undersized groups
|
|
const orphanIndices: number[] = []
|
|
|
|
for (const gid of getNonEmptyGroupIds()) {
|
|
const elemSet = groupElements.get(gid)!
|
|
if (elemSet.size > 0 && elemSet.size < options.minElements) {
|
|
for (const idx of elemSet) {
|
|
orphanIndices.push(idx)
|
|
}
|
|
elemSet.clear()
|
|
}
|
|
}
|
|
|
|
if (orphanIndices.length > 0) {
|
|
// Valid groups = everything that's still non-empty (i.e. above minElements)
|
|
const validGroupIds = getNonEmptyGroupIds()
|
|
|
|
// Chunk orphans and run them through processChunk with only valid groups visible
|
|
const orphanChunks: number[][] = []
|
|
let currentOrphanChunk: number[] = []
|
|
let currentOrphanTokens = 0
|
|
|
|
for (const elemIdx of orphanIndices) {
|
|
const elem = elements[elemIdx]
|
|
const truncated = tokenizer.truncate(elem.stringified, options.tokensPerElement)
|
|
const elemTokens = tokenizer.count(truncated)
|
|
|
|
if (
|
|
(currentOrphanTokens + elemTokens > TOKENS_FOR_ELEMENTS_MAX ||
|
|
currentOrphanChunk.length >= MAX_ELEMENTS_PER_CHUNK) &&
|
|
currentOrphanChunk.length > 0
|
|
) {
|
|
orphanChunks.push(currentOrphanChunk)
|
|
currentOrphanChunk = []
|
|
currentOrphanTokens = 0
|
|
}
|
|
|
|
currentOrphanChunk.push(elemIdx)
|
|
currentOrphanTokens += elemTokens
|
|
}
|
|
|
|
if (currentOrphanChunk.length > 0) {
|
|
orphanChunks.push(currentOrphanChunk)
|
|
}
|
|
|
|
// Process orphan chunks against valid groups (reuses the same processChunk as Phase 1)
|
|
const orphanResults = await Promise.all(
|
|
orphanChunks.map((chunk) =>
|
|
elementLimit(async () => {
|
|
// If there are valid groups, chunk them; otherwise pass empty so LLM creates new groups
|
|
const groupChunksForOrphans = validGroupIds.length > 0 ? getGroupChunks() : [[]]
|
|
|
|
const allAssignments = await Promise.all(
|
|
groupChunksForOrphans
|
|
.filter((gc) => gc.length === 0 || gc.some((gid) => validGroupIds.includes(gid)))
|
|
.map((groupChunk) => {
|
|
// Only show valid groups (exclude the orphaned/undersized ones)
|
|
const filteredGroupChunk = groupChunk.filter((gid) => validGroupIds.includes(gid))
|
|
return groupLimit(() => processChunk(chunk, filteredGroupChunk))
|
|
})
|
|
)
|
|
|
|
return allAssignments.flat()
|
|
})
|
|
)
|
|
)
|
|
|
|
// Apply assignments
|
|
const flatAssignments = orphanResults.flat()
|
|
for (const { elementIndex, label } of flatAssignments) {
|
|
const normalized = normalizeLabel(label)
|
|
let groupId = labelToGroupId.get(normalized)
|
|
|
|
if (!groupId) {
|
|
groupId = `group_${groupIdCounter++}`
|
|
groups.set(groupId, { id: groupId, label, normalizedLabel: normalized })
|
|
groupElements.set(groupId, new Set())
|
|
labelToGroupId.set(normalized, groupId)
|
|
}
|
|
groupElements.get(groupId)!.add(elementIndex)
|
|
}
|
|
|
|
// Safety: any orphans the LLM missed get placed into the largest group
|
|
const isAssigned = (idx: number) => {
|
|
for (const [, elemSet] of groupElements) {
|
|
if (elemSet.has(idx)) return true
|
|
}
|
|
return false
|
|
}
|
|
const unassigned = orphanIndices.filter((idx) => !isAssigned(idx))
|
|
const placeIntoLargest = (indices: number[]) => {
|
|
const allNonEmpty = getNonEmptyGroupIds()
|
|
if (allNonEmpty.length === 0) return
|
|
const largestGid = allNonEmpty.reduce((a, b) =>
|
|
groupElements.get(a)!.size >= groupElements.get(b)!.size ? a : b
|
|
)
|
|
for (const idx of indices) {
|
|
groupElements.get(largestGid)!.add(idx)
|
|
}
|
|
}
|
|
|
|
if (unassigned.length > 0) {
|
|
placeIntoLargest(unassigned)
|
|
}
|
|
|
|
// Second pass: if any groups are still undersized after redistribution,
|
|
// merge their elements into the largest group
|
|
const mergeUndersizedGroups = () => {
|
|
const allNonEmpty = getNonEmptyGroupIds()
|
|
if (allNonEmpty.length <= 1) return false
|
|
|
|
const largestGid = allNonEmpty.reduce((a, b) =>
|
|
groupElements.get(a)!.size >= groupElements.get(b)!.size ? a : b
|
|
)
|
|
const targetSet = groupElements.get(largestGid)!
|
|
let merged = false
|
|
|
|
for (const gid of allNonEmpty) {
|
|
if (gid === largestGid) continue
|
|
const elemSet = groupElements.get(gid)!
|
|
if (elemSet.size > 0 && elemSet.size < options.minElements) {
|
|
elemSet.forEach((idx) => targetSet.add(idx))
|
|
elemSet.clear()
|
|
merged = true
|
|
}
|
|
}
|
|
return merged
|
|
}
|
|
|
|
while (mergeUndersizedGroups()) {
|
|
// keep merging until no undersized groups remain
|
|
}
|
|
}
|
|
}
|
|
|
|
// Build final result
|
|
const result: Array<Group<T>> = []
|
|
|
|
for (const [groupId, elementIndices] of groupElements.entries()) {
|
|
if (elementIndices.size > 0) {
|
|
const groupInfo = groups.get(groupId)!
|
|
result.push({
|
|
id: groupInfo.id,
|
|
label: groupInfo.label,
|
|
elements: Array.from(elementIndices).map((idx) => elements[idx].element),
|
|
})
|
|
}
|
|
}
|
|
|
|
// Save example for active learning
|
|
if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
|
|
const key = fastHash(
|
|
stringify({
|
|
taskId,
|
|
taskType,
|
|
input: JSON.stringify(input),
|
|
instructions: options.instructions ?? '',
|
|
})
|
|
)
|
|
|
|
// Build output format for saving
|
|
const outputAssignments: Array<{ elementIndex: number; label: string }> = []
|
|
for (const [groupId, elementIndices] of groupElements.entries()) {
|
|
const groupInfo = groups.get(groupId)!
|
|
for (const idx of elementIndices) {
|
|
outputAssignments.push({
|
|
elementIndex: idx,
|
|
label: groupInfo.label,
|
|
})
|
|
}
|
|
}
|
|
|
|
// Note: We don't have direct access to usage metadata here since it's distributed
|
|
// across many parallel operations. We'll use default values.
|
|
await ctx.adapter.saveExample({
|
|
key,
|
|
taskType,
|
|
taskId,
|
|
input: JSON.stringify(input),
|
|
output: result,
|
|
instructions: options.instructions ?? '',
|
|
metadata: {
|
|
cost: { input: 0, output: 0 },
|
|
latency: 0,
|
|
model: ctx.modelId,
|
|
tokens: { input: 0, output: 0 },
|
|
},
|
|
})
|
|
}
|
|
|
|
return result
|
|
}
|
|
|
|
Zai.prototype.group = function <T>(
|
|
this: Zai,
|
|
input: Array<T>,
|
|
_options?: Options
|
|
): Response<Array<Group<T>>, Record<string, T[]>> {
|
|
const context = new ZaiContext({
|
|
client: this.client,
|
|
modelId: this.Model,
|
|
taskId: this.taskId,
|
|
taskType: 'zai.group',
|
|
adapter: this.adapter,
|
|
memoizer: this._resolveMemoizer(),
|
|
})
|
|
|
|
return new Response<Array<Group<T>>, Record<string, T[]>>(context, group(input, _options, context), (result) => {
|
|
const merged: Record<string, T[]> = {}
|
|
result.forEach((group) => {
|
|
if (!merged[group.label]) {
|
|
merged[group.label] = []
|
|
}
|
|
merged[group.label].push(...group.elements)
|
|
})
|
|
return merged
|
|
})
|
|
}
|