updates

Author: lakeesiv

apps/sim/app/api/wand/route.ts

@@ -11,6 +11,7 @@ import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
 import { env } from '@/lib/core/config/env'
 import { getCostMultiplier, isBillingEnabled } from '@/lib/core/config/feature-flags'
 import { generateRequestId } from '@/lib/core/utils/request'
+import { enrichTableSchema } from '@/lib/table/llm/wand'
 import { verifyWorkspaceMembership } from '@/app/api/workflows/utils'
 import { getModelPricing } from '@/providers/utils'
 
@@ -60,6 +61,7 @@ interface RequestBody {
   history?: ChatMessage[]
   workflowId?: string
   generationType?: string
+  wandContext?: Record<string, unknown>
 }
 
 function safeStringify(value: unknown): string {
@@ -70,6 +72,38 @@ function safeStringify(value: unknown): string {
   }
 }
 
+/**
+ * Wand enricher function type.
+ * Enrichers add context to the system prompt based on generationType.
+ */
+type WandEnricher = (
+  workspaceId: string | null,
+  context: Record<string, unknown>
+) => Promise<string | null>
+
+/**
+ * Registry of wand enrichers by generationType.
+ * Each enricher returns additional context to append to the system prompt.
+ */
+const wandEnrichers: Partial<Record<string, WandEnricher>> = {
+  timestamp: async () => {
+    const now = new Date()
+    return `Current date and time context for reference:
+- Current UTC timestamp: ${now.toISOString()}
+- Current Unix timestamp (seconds): ${Math.floor(now.getTime() / 1000)}
+- Current Unix timestamp (milliseconds): ${now.getTime()}
+- Current date (UTC): ${now.toISOString().split('T')[0]}
+- Current year: ${now.getUTCFullYear()}
+- Current month: ${now.getUTCMonth() + 1}
+- Current day of month: ${now.getUTCDate()}
+- Current day of week: ${['Sunday', 'Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday'][now.getUTCDay()]}
+
+Use this context to calculate relative dates like "yesterday", "last week", "beginning of this month", etc.`
+  },
+
+  'table-schema': enrichTableSchema,
+}
+
 async function updateUserStatsForWand(
   userId: string,
   usage: {
@@ -159,7 +193,15 @@ export async function POST(req: NextRequest) {
   try {
     const body = (await req.json()) as RequestBody
 
-    const { prompt, systemPrompt, stream = false, history = [], workflowId, generationType } = body
+    const {
+      prompt,
+      systemPrompt,
+      stream = false,
+      history = [],
+      workflowId,
+      generationType,
+      wandContext = {},
+    } = body
 
     if (!prompt) {
       logger.warn(`[${requestId}] Invalid request: Missing prompt.`)
@@ -227,20 +269,15 @@ export async function POST(req: NextRequest) {
       systemPrompt ||
       'You are a helpful AI assistant. Generate content exactly as requested by the user.'
 
-    if (generationType === 'timestamp') {
-      const now = new Date()
-      const currentTimeContext = `\n\nCurrent date and time context for reference:
-- Current UTC timestamp: ${now.toISOString()}
-- Current Unix timestamp (seconds): ${Math.floor(now.getTime() / 1000)}
-- Current Unix timestamp (milliseconds): ${now.getTime()}
-- Current date (UTC): ${now.toISOString().split('T')[0]}
-- Current year: ${now.getUTCFullYear()}
-- Current month: ${now.getUTCMonth() + 1}
-- Current day of month: ${now.getUTCDate()}
-- Current day of week: ${['Sunday', 'Monday', 'Tuesday', 'Wednesday', 'Thursday', 'Friday', 'Saturday'][now.getUTCDay()]}
-
-Use this context to calculate relative dates like "yesterday", "last week", "beginning of this month", etc.`
-      finalSystemPrompt += currentTimeContext
+    // Apply enricher if one exists for this generationType
+    if (generationType) {
+      const enricher = wandEnrichers[generationType]
+      if (enricher) {
+        const enrichment = await enricher(workspaceId, wandContext)
+        if (enrichment) {
+          finalSystemPrompt += `\n\n${enrichment}`
+        }
+      }
     }
 
     if (generationType === 'json-object') {

apps/sim/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand.ts

@@ -1,7 +1,6 @@
 import { useCallback, useRef, useState } from 'react'
 import { createLogger } from '@sim/logger'
 import { useQueryClient } from '@tanstack/react-query'
-import { fetchTableSchemaContext } from '@/lib/table/wand'
 import type { GenerationType } from '@/blocks/types'
 import { subscriptionKeys } from '@/hooks/queries/subscription'
 import { useWorkflowRegistry } from '@/stores/workflows/registry/store'
@@ -16,19 +15,16 @@ interface ChatMessage {
 interface BuildWandContextInfoOptions {
   currentValue?: string
   generationType?: string
-  tableId?: string | null
-  workspaceId?: string | null
 }
 
 /**
  * Builds rich context information based on current content and generation type.
+ * Note: Table schema context is now fetched server-side in /api/wand for simplicity.
  */
-async function buildWandContextInfo({
+function buildWandContextInfo({
   currentValue,
   generationType,
-  tableId,
-  workspaceId,
-}: BuildWandContextInfoOptions): Promise<string> {
+}: BuildWandContextInfoOptions): string {
   const hasContent = Boolean(currentValue && currentValue.trim() !== '')
   const contentLength = currentValue?.length ?? 0
   const lineCount = currentValue ? currentValue.split('\n').length : 0
@@ -62,13 +58,6 @@ async function buildWandContextInfo({
     }
   }
 
-  if (generationType === 'table-schema') {
-    const tableSchemaContext = await fetchTableSchemaContext({ tableId, workspaceId })
-    if (tableSchemaContext) {
-      contextInfo += `\n\n${tableSchemaContext}`
-    }
-  }
-
   return contextInfo
 }
 
@@ -102,7 +91,7 @@ export function useWand({
   onGenerationComplete,
 }: UseWandProps) {
   const queryClient = useQueryClient()
-  const workspaceId = useWorkflowRegistry((state) => state.hydration.workspaceId)
+  const workflowId = useWorkflowRegistry((state) => state.hydration.workflowId)
   const [isLoading, setIsLoading] = useState(false)
   const [isPromptVisible, setIsPromptVisible] = useState(false)
   const [promptInputValue, setPromptInputValue] = useState('')
@@ -173,11 +162,9 @@ export function useWand({
       }
 
       try {
-        const contextInfo = await buildWandContextInfo({
+        const contextInfo = buildWandContextInfo({
           currentValue,
           generationType: wandConfig?.generationType,
-          tableId: contextParams?.tableId,
-          workspaceId,
         })
 
         let systemPrompt = wandConfig?.prompt || ''
@@ -201,6 +188,8 @@ export function useWand({
             stream: true,
             history: wandConfig?.maintainHistory ? conversationHistory : [],
             generationType: wandConfig?.generationType,
+            workflowId,
+            wandContext: contextParams?.tableId ? { tableId: contextParams.tableId } : undefined,
           }),
           signal: abortControllerRef.current.signal,
           cache: 'no-store',
@@ -308,7 +297,7 @@ export function useWand({
       onGenerationComplete,
       queryClient,
       contextParams?.tableId,
-      workspaceId,
+      workflowId,
     ]
   )
 

apps/sim/lib/table/llm/index.ts

@@ -1,5 +1 @@
-/**
- * LLM enrichment utilities for table operations.
- */
-
 export * from './enrichment'