@pikku/core 0.12.19 → 0.12.21

package/src/wirings/ai-agent/ai-agent-runner.ts

@@ -20,6 +20,7 @@ import {
   resolveMemoryServices,
   loadContextMessages,
   trimMessages,
+  getWorkingMemoryMiddleware,
 } from './ai-agent-memory.js'
 import {
   prepareAgentRun,
@@ -47,6 +48,21 @@ function stripNulls(obj: unknown): unknown {
 }
 import { randomUUID } from 'crypto'
 
+function extractStructuredAssistantOutput(object: unknown): {
+  text: string | null
+  uiSpec: unknown | null
+} {
+  if (!object || typeof object !== 'object') {
+    return { text: null, uiSpec: null }
+  }
+
+  const record = object as Record<string, unknown>
+  return {
+    text: typeof record.text === 'string' ? record.text : null,
+    uiSpec: record.ui ?? null,
+  }
+}
+
 export async function runAIAgent(
   agentName: string,
   input: AIAgentInput,
@@ -65,7 +81,6 @@ export async function runAIAgent(
     runnerParams,
     maxSteps,
     missingRpcs,
-    workingMemoryJsonSchema,
     workingMemorySchemaName,
   } = await prepareAgentRun(agentName, input, params, sessionMap)
 
@@ -96,7 +111,15 @@ export async function runAIAgent(
     }
   }
 
-  const aiMiddlewares: PikkuAIMiddlewareHooks[] = agent.aiMiddleware ?? []
+  const aiMiddlewares: PikkuAIMiddlewareHooks[] = [
+    ...getWorkingMemoryMiddleware(memoryConfig, storage, {
+      threadId,
+      workingMemorySchemaName,
+      logger: singletonServices.logger,
+      schemaService: singletonServices.schema,
+    }),
+    ...(agent.aiMiddleware ?? []),
+  ]
 
   let modifiedMessages = runnerParams.messages
   let modifiedInstructions = runnerParams.instructions
@@ -228,13 +251,7 @@ export async function runAIAgent(
           input.resourceId,
           memoryConfig,
           userMessage,
-          { text: '', steps: completedStepsForSave },
-          {
-            workingMemoryJsonSchema,
-            workingMemorySchemaName,
-            logger: singletonServices.logger,
-            schemaService: singletonServices.schema,
-          }
+          { text: '', steps: completedStepsForSave }
         )
 
         if (storage) {
@@ -280,30 +297,16 @@ export async function runAIAgent(
       appendStepMessages(runnerParams, stepResult)
     }
 
-    const finalText = lastStepResult?.text ?? ''
     const finalObject = lastStepResult?.object
+    const structuredOutput = extractStructuredAssistantOutput(finalObject)
+    const finalText = structuredOutput.text ?? lastStepResult?.text ?? ''
 
     const result = {
       text: finalText,
       steps: accumulatedSteps,
     }
 
-    const responseText = await saveMessages(
-      storage,
-      threadId,
-      input.resourceId,
-      memoryConfig,
-      userMessage,
-      result,
-      {
-        workingMemoryJsonSchema,
-        workingMemorySchemaName,
-        logger: singletonServices.logger,
-        schemaService: singletonServices.schema,
-      }
-    )
-
-    let outputText = responseText
+    let outputText = finalText
     let outputMessages = runnerParams.messages
     for (let i = aiMiddlewares.length - 1; i >= 0; i--) {
       const mw = aiMiddlewares[i]
@@ -318,6 +321,19 @@ export async function runAIAgent(
       }
     }
 
+    await saveMessages(
+      storage,
+      threadId,
+      input.resourceId,
+      memoryConfig,
+      userMessage,
+      {
+        ...result,
+        text: outputText,
+        uiSpec: structuredOutput.uiSpec,
+      }
+    )
+
     await aiRunState.updateRun(runId, {
       status: 'completed',
       usage: { ...totalUsage, model: agent.model },
@@ -535,7 +551,15 @@ async function continueAfterToolResultSync(
 
   const instructions = await buildInstructions(resolvedName, packageName)
 
-  const aiMiddlewares: PikkuAIMiddlewareHooks[] = agent.aiMiddleware ?? []
+  const aiMiddlewares: PikkuAIMiddlewareHooks[] = [
+    ...getWorkingMemoryMiddleware(memoryConfig, storage, {
+      threadId: run.threadId,
+      workingMemorySchemaName,
+      logger: singletonServices.logger,
+      schemaService: singletonServices.schema,
+    }),
+    ...(agent.aiMiddleware ?? []),
+  ]
   let modifiedMessages = trimmedMessages
   let modifiedInstructions = instructions
   for (const mw of aiMiddlewares) {
@@ -666,13 +690,7 @@ async function continueAfterToolResultSync(
             run.resourceId,
             memoryConfig,
             null,
-            { text: '', steps: completedSteps },
-            {
-              workingMemoryJsonSchema,
-              workingMemorySchemaName,
-              logger: singletonServices.logger,
-              schemaService: singletonServices.schema,
-            }
+            { text: '', steps: completedSteps }
           )
         }
 
@@ -727,22 +745,7 @@ async function continueAfterToolResultSync(
       steps: accumulatedSteps,
     }
 
-    const responseText = await saveMessages(
-      storage,
-      run.threadId,
-      run.resourceId,
-      memoryConfig,
-      null as any,
-      result,
-      {
-        workingMemoryJsonSchema,
-        workingMemorySchemaName,
-        logger: singletonServices.logger,
-        schemaService: singletonServices.schema,
-      }
-    )
-
-    let outputText = responseText
+    let outputText = finalText
     let outputMessages = runnerParams.messages
     for (let i = aiMiddlewares.length - 1; i >= 0; i--) {
       const mw = aiMiddlewares[i]
@@ -757,6 +760,18 @@ async function continueAfterToolResultSync(
       }
     }
 
+    await saveMessages(
+      storage,
+      run.threadId,
+      run.resourceId,
+      memoryConfig,
+      null as any,
+      {
+        ...result,
+        text: outputText,
+      }
+    )
+
     await aiRunState.updateRun(run.runId, {
       status: 'completed',
       usage: { ...totalUsage, model: agent.model },