This commit is contained in:
Lakee Sivaraya
2026-01-16 13:52:10 -08:00
parent e4dd14df7a
commit cca1772ae1
2 changed files with 104 additions and 46 deletions

View File

@@ -5,8 +5,92 @@
* with table-specific information so LLMs can construct proper queries.
*/
import { createLogger } from '@sim/logger'
import type { TableSummary } from '../types'
const logger = createLogger('TableLLMEnrichment')
export interface TableEnrichmentContext {
workspaceId: string
workflowId: string
executeTool: (toolId: string, params: Record<string, any>) => Promise<any>
}
export interface TableEnrichmentResult {
description: string
parameters: {
properties: Record<string, any>
required: string[]
}
}
/**
* Enriches a table tool for LLM consumption by fetching its schema
* and injecting column information into the description and parameters.
*
* @param toolId - The table tool ID (e.g., 'table_query_rows')
* @param originalDescription - The tool's original description
* @param llmSchema - The original LLM schema
* @param userProvidedParams - Parameters provided by the user (must include tableId)
* @param context - Execution context with workspaceId, workflowId, and executeTool
* @returns Enriched description and parameters, or null if enrichment not applicable
*/
export async function enrichTableToolForLLM(
toolId: string,
originalDescription: string,
llmSchema: { properties?: Record<string, any>; required?: string[] },
userProvidedParams: Record<string, any>,
context: TableEnrichmentContext
): Promise<TableEnrichmentResult | null> {
const { tableId } = userProvidedParams
// Need a tableId to fetch schema
if (!tableId) {
return null
}
try {
logger.info(`Fetching schema for table ${tableId}`)
const schemaResult = await context.executeTool('table_get_schema', {
tableId,
_context: {
workspaceId: context.workspaceId,
workflowId: context.workflowId,
},
})
if (!schemaResult.success || !schemaResult.output) {
logger.warn(`Failed to fetch table schema: ${schemaResult.error}`)
return null
}
const tableSchema: TableSummary = {
name: schemaResult.output.name,
columns: schemaResult.output.columns || [],
}
// Apply enrichment using the existing utility functions
const enrichedDescription = enrichTableToolDescription(originalDescription, tableSchema, toolId)
const enrichedParams = enrichTableToolParameters(llmSchema, tableSchema, toolId)
logger.info(`Enriched ${toolId} with ${tableSchema.columns.length} columns`)
return {
description: enrichedDescription,
parameters: {
properties: enrichedParams.properties,
required:
enrichedParams.required.length > 0 ? enrichedParams.required : llmSchema.required || [],
},
}
} catch (error) {
logger.warn(`Error fetching table schema:`, error)
return null
}
}
/**
* Operations that use filters and need filter-specific enrichment.
*/

View File

@@ -3,7 +3,7 @@ import type { ChatCompletionChunk } from 'openai/resources/chat/completions'
import type { CompletionUsage } from 'openai/resources/completions'
import { env } from '@/lib/core/config/env'
import { isHosted } from '@/lib/core/config/feature-flags'
import { enrichTableToolDescription, enrichTableToolParameters } from '@/lib/table/llm'
import { enrichTableToolForLLM } from '@/lib/table/llm'
import { isCustomTool } from '@/executor/constants'
import {
getComputerUseModels,
@@ -497,60 +497,34 @@ export async function transformBlockTool(
uniqueToolId = `${toolConfig.id}_${userProvidedParams.knowledgeBaseId}`
}
// Enrich table tool descriptions with schema information
let enrichedDescription = toolConfig.description
let enrichedLlmSchema = llmSchema
if (
toolId.startsWith('table_') &&
userProvidedParams.tableId &&
workspaceId &&
workflowId &&
executeTool
) {
try {
logger.info(`[transformBlockTool] Fetching schema for table ${userProvidedParams.tableId}`)
const schemaResult = await executeTool('table_get_schema', {
tableId: userProvidedParams.tableId,
_context: { workspaceId, workflowId },
})
// Apply table tool enrichment if applicable
let finalDescription = toolConfig.description
let finalSchema = llmSchema
if (schemaResult.success && schemaResult.output) {
const tableSchema = {
name: schemaResult.output.name,
columns: schemaResult.output.columns || [],
}
// Enrich description and parameters using lib/table utilities
enrichedDescription = enrichTableToolDescription(
toolConfig.description,
tableSchema,
toolId
)
const enrichedParams = enrichTableToolParameters(llmSchema, tableSchema, toolId)
enrichedLlmSchema = {
...llmSchema,
properties: enrichedParams.properties,
required:
enrichedParams.required.length > 0 ? enrichedParams.required : llmSchema.required,
}
logger.info(
`[transformBlockTool] Enriched ${toolId} with ${tableSchema.columns.length} columns`
)
} else {
logger.warn(`[transformBlockTool] Failed to fetch table schema: ${schemaResult.error}`)
if (toolId.startsWith('table_') && workspaceId && workflowId && executeTool) {
const result = await enrichTableToolForLLM(
toolId,
toolConfig.description,
llmSchema,
userProvidedParams,
{
workspaceId,
workflowId,
executeTool,
}
} catch (error) {
logger.warn(`[transformBlockTool] Error fetching table schema:`, error)
)
if (result) {
finalDescription = result.description
finalSchema = { ...llmSchema, ...result.parameters }
}
}
return {
id: uniqueToolId,
name: toolConfig.name,
description: enrichedDescription,
description: finalDescription,
params: userProvidedParams,
parameters: enrichedLlmSchema,
parameters: finalSchema,
}
}