@pikku/core 0.11.2 → 0.12.1

package/src/wirings/ai-agent/ai-agent-helpers.test.ts

@@ -0,0 +1,152 @@
+import { describe, test } from 'node:test'
+import assert from 'node:assert'
+import {
+  agent,
+  agentStream,
+  agentResume,
+  agentApprove,
+} from './ai-agent-helpers.js'
+
+describe('agent helpers', () => {
+  describe('agent', () => {
+    test('should return an object with func', () => {
+      const a = agent('myAgent')
+      assert.strictEqual(typeof a.func, 'function')
+    })
+
+    test('func should call rpc.agent.run', async () => {
+      const a = agent('myAgent')
+      let calledWith: any
+      const rpc = {
+        agent: {
+          run: async (...args: any[]) => {
+            calledWith = args
+            return {
+              runId: 'run-1',
+              result: 'done',
+              usage: { inputTokens: 10, outputTokens: 20 },
+            }
+          },
+        },
+      }
+      const input = { message: 'hello', threadId: 't1', resourceId: 'r1' }
+      const result = await a.func({}, input, { rpc })
+      assert.strictEqual(calledWith[0], 'myAgent')
+      assert.deepStrictEqual(calledWith[1], input)
+      assert.strictEqual(result.runId, 'run-1')
+    })
+  })
+
+  describe('agentStream', () => {
+    test('should return an object with func', () => {
+      const a = agentStream('myAgent')
+      assert.strictEqual(typeof a.func, 'function')
+    })
+
+    test('func should call rpc.agent.stream with bound name', async () => {
+      const a = agentStream('myAgent')
+      let calledWith: any
+      const rpc = {
+        agent: {
+          stream: async (...args: any[]) => {
+            calledWith = args
+          },
+        },
+      }
+      const input = { message: 'hello', threadId: 't1', resourceId: 'r1' }
+      await a.func({}, input, { rpc })
+      assert.strictEqual(calledWith[0], 'myAgent')
+      assert.deepStrictEqual(calledWith[1], input)
+    })
+
+    test('func should extract agentName from data when no name bound', async () => {
+      const a = agentStream()
+      let calledWith: any
+      const rpc = {
+        agent: {
+          stream: async (...args: any[]) => {
+            calledWith = args
+          },
+        },
+      }
+      const input = {
+        agentName: 'dynamicAgent',
+        message: 'hello',
+        threadId: 't1',
+        resourceId: 'r1',
+      }
+      await a.func({}, input, { rpc })
+      assert.strictEqual(calledWith[0], 'dynamicAgent')
+      assert.deepStrictEqual(calledWith[1], {
+        message: 'hello',
+        threadId: 't1',
+        resourceId: 'r1',
+      })
+    })
+
+    test('func should throw when no name bound and no agentName in data', async () => {
+      const a = agentStream()
+      const rpc = { agent: { stream: async () => {} } }
+      await assert.rejects(() => a.func({}, { message: 'hello' }, { rpc }), {
+        message:
+          'agentStream requires an agentName either as a parameter or in the input data',
+      })
+    })
+  })
+
+  describe('agentResume', () => {
+    test('should return an object with func', () => {
+      const a = agentResume()
+      assert.strictEqual(typeof a.func, 'function')
+    })
+
+    test('func should call rpc.agent.resume', async () => {
+      const a = agentResume()
+      let calledWith: any
+      const rpc = {
+        agent: {
+          resume: async (...args: any[]) => {
+            calledWith = args
+          },
+        },
+      }
+      const data = { runId: 'run-1', toolCallId: 'tc1', approved: true }
+      await a.func({}, data, { rpc })
+      assert.strictEqual(calledWith[0], 'run-1')
+      assert.deepStrictEqual(calledWith[1], {
+        toolCallId: 'tc1',
+        approved: true,
+      })
+    })
+  })
+
+  describe('agentApprove', () => {
+    test('should return an object with func', () => {
+      const a = agentApprove('myAgent')
+      assert.strictEqual(typeof a.func, 'function')
+    })
+
+    test('func should call rpc.agent.approve', async () => {
+      const a = agentApprove('myAgent')
+      let calledWith: any
+      const rpc = {
+        agent: {
+          approve: async (...args: any[]) => {
+            calledWith = args
+            return { approved: true }
+          },
+        },
+      }
+      const data = {
+        runId: 'run-1',
+        approvals: [{ toolCallId: 'tc1', approved: true }],
+      }
+      const result = await a.func({}, data, { rpc })
+      assert.strictEqual(calledWith[0], 'run-1')
+      assert.deepStrictEqual(calledWith[1], [
+        { toolCallId: 'tc1', approved: true },
+      ])
+      assert.deepStrictEqual(result, { approved: true })
+    })
+  })
+})

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

@@ -0,0 +1,76 @@
+export function agent<TAgentMap extends Record<string, { output: any }>>(
+  agentName: string & keyof TAgentMap
+): { func: (services: any, data: any, wire: any) => Promise<any> } {
+  return {
+    func: async (_services: any, data: any, { rpc }: any) => {
+      return rpc.agent.run(agentName, data)
+    },
+  }
+}
+
+export function agentStream<TAgentMap extends Record<string, { output: any }>>(
+  agentName?: string & keyof TAgentMap
+): { func: (services: any, data: any, wire: any) => Promise<void> } {
+  return {
+    func: async (_services: any, data: any, { rpc }: any) => {
+      if (agentName) {
+        await rpc.agent.stream(agentName, data)
+      } else {
+        const { agentName: name, ...rest } = data
+        if (!name) {
+          throw new Error(
+            'agentStream requires an agentName either as a parameter or in the input data'
+          )
+        }
+        await rpc.agent.stream(name, rest)
+      }
+    },
+  }
+}
+
+export function agentResume(): {
+  func: (
+    services: any,
+    data: { runId: string; toolCallId: string; approved: boolean },
+    wire: any
+  ) => Promise<void>
+} {
+  return {
+    func: async (
+      _services: any,
+      data: { runId: string; toolCallId: string; approved: boolean },
+      { rpc }: any
+    ) => {
+      await rpc.agent.resume(data.runId, {
+        toolCallId: data.toolCallId,
+        approved: data.approved,
+      })
+    },
+  }
+}
+
+export function agentApprove<TAgentMap extends Record<string, { output: any }>>(
+  _agentName: string & keyof TAgentMap
+): {
+  func: (
+    services: any,
+    data: {
+      runId: string
+      approvals: { toolCallId: string; approved: boolean }[]
+    },
+    wire: any
+  ) => Promise<any>
+} {
+  return {
+    func: async (
+      _services: any,
+      data: {
+        runId: string
+        approvals: { toolCallId: string; approved: boolean }[]
+      },
+      { rpc }: any
+    ) => {
+      return rpc.agent.approve(data.runId, data.approvals)
+    },
+  }
+}

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

@@ -0,0 +1,300 @@
+import { randomUUID } from 'crypto'
+import type { CoreSingletonServices } from '../../types/core.types.js'
+import type {
+  CoreAIAgent,
+  AIAgentMemoryConfig,
+  AIAgentInput,
+  AIMessage,
+} from './ai-agent.types.js'
+import type { AIStorageService } from '../../services/ai-storage-service.js'
+import type { Logger } from '../../services/logger.js'
+import type { SchemaService } from '../../services/schema-service.js'
+
+export function resolveMemoryServices(
+  agent: CoreAIAgent,
+  singletonServices: CoreSingletonServices
+): {
+  storage: AIStorageService | undefined
+} {
+  const memoryConfig = agent.memory
+  const storage = memoryConfig?.storage
+    ? (singletonServices as any)[memoryConfig.storage]
+    : singletonServices.aiStorage
+  return { storage }
+}
+
+export function deepMergeWorkingMemory(
+  existing: Record<string, unknown>,
+  updates: Record<string, unknown>
+): Record<string, unknown> {
+  const result = { ...existing }
+  for (const key of Object.keys(updates)) {
+    const value = updates[key]
+    if (value === null) {
+      delete result[key]
+    } else if (
+      typeof value === 'object' &&
+      !Array.isArray(value) &&
+      typeof result[key] === 'object' &&
+      result[key] !== null &&
+      !Array.isArray(result[key])
+    ) {
+      result[key] = deepMergeWorkingMemory(
+        result[key] as Record<string, unknown>,
+        value as Record<string, unknown>
+      )
+    } else {
+      result[key] = value
+    }
+  }
+  return result
+}
+
+export function buildWorkingMemoryPrompt(
+  currentState: Record<string, unknown> | null,
+  jsonSchema?: Record<string, unknown>
+): string {
+  const parts: string[] = []
+
+  if (jsonSchema?.properties) {
+    const props = jsonSchema.properties as Record<
+      string,
+      { type?: string; description?: string }
+    >
+    const fieldLines = Object.entries(props).map(([name, def]) => {
+      const type = def.type ?? 'unknown'
+      const desc = def.description ? ` - ${def.description}` : ''
+      return `  - ${name} (${type})${desc}`
+    })
+    if (fieldLines.length > 0) {
+      parts.push(`Working memory fields:\n${fieldLines.join('\n')}`)
+    }
+  }
+
+  if (currentState && Object.keys(currentState).length > 0) {
+    parts.push(`Current working memory:\n${JSON.stringify(currentState)}`)
+  } else {
+    parts.push('Current working memory: (empty)')
+  }
+
+  parts.push(
+    'When you learn new information, output a partial JSON update in <working_memory> tags. ' +
+      'Only include changed fields. Set a field to null to delete it.'
+  )
+
+  return parts.join('\n\n')
+}
+
+export async function loadContextMessages(
+  memoryConfig: AIAgentMemoryConfig | undefined,
+  storage: AIStorageService | undefined,
+  input: AIAgentInput,
+  workingMemoryJsonSchema?: Record<string, unknown>
+): Promise<AIMessage[]> {
+  const contextMessages: AIMessage[] = []
+
+  if (memoryConfig?.workingMemory && storage) {
+    const workingMem = await storage.getWorkingMemory(input.threadId, 'thread')
+    const prompt = buildWorkingMemoryPrompt(workingMem, workingMemoryJsonSchema)
+    contextMessages.push({
+      id: `wm-${randomUUID()}`,
+      role: 'system',
+      content: prompt,
+      createdAt: new Date(),
+    })
+  }
+
+  return contextMessages
+}
+
+export async function saveMessages(
+  storage: AIStorageService | undefined,
+  threadId: string,
+  resourceId: string,
+  memoryConfig: AIAgentMemoryConfig | undefined,
+  userMessage: AIMessage,
+  result: {
+    text: string
+    steps: {
+      toolCalls?: {
+        name: string
+        args: Record<string, unknown>
+        result: string
+      }[]
+    }[]
+  },
+  options?: {
+    workingMemoryJsonSchema?: Record<string, unknown>
+    workingMemorySchemaName?: string | null
+    logger?: Logger
+    schemaService?: SchemaService
+  }
+): Promise<string> {
+  let responseText = result.text
+
+  if (storage) {
+    const newMessages: AIMessage[] = [userMessage]
+
+    for (const step of result.steps) {
+      if (step.toolCalls?.length) {
+        newMessages.push({
+          id: randomUUID(),
+          role: 'assistant',
+          toolCalls: step.toolCalls.map((tc) => ({
+            id: randomUUID(),
+            name: tc.name,
+            args: tc.args,
+          })),
+          createdAt: new Date(),
+        })
+        newMessages.push({
+          id: randomUUID(),
+          role: 'tool',
+          toolResults: step.toolCalls.map((tc) => ({
+            id: randomUUID(),
+            name: tc.name,
+            result: tc.result,
+          })),
+          createdAt: new Date(),
+        })
+      }
+    }
+
+    newMessages.push({
+      id: randomUUID(),
+      role: 'assistant',
+      content: responseText,
+      createdAt: new Date(),
+    })
+
+    await storage.saveMessages(threadId, newMessages)
+
+    if (memoryConfig?.workingMemory) {
+      const parsed = parseWorkingMemory(responseText)
+      if (parsed) {
+        const existing =
+          (await storage.getWorkingMemory(threadId, 'thread')) ?? {}
+        const merged = deepMergeWorkingMemory(existing, parsed)
+
+        if (options?.schemaService && options?.workingMemorySchemaName) {
+          try {
+            await options.schemaService.validateSchema(
+              options.workingMemorySchemaName,
+              merged
+            )
+          } catch (err) {
+            options.logger?.warn(
+              `Working memory validation failed: ${err instanceof Error ? err.message : String(err)}`
+            )
+          }
+        }
+
+        await storage.saveWorkingMemory(threadId, 'thread', merged)
+        responseText = stripWorkingMemoryTags(responseText)
+      }
+    }
+  }
+
+  return responseText
+}
+
+export function trimMessages(
+  messages: AIMessage[],
+  maxTokenBudget: number = 100000
+): AIMessage[] {
+  const sanitized = sanitizeToolMessages(messages)
+
+  let estimatedTokens = 0
+  const result: AIMessage[] = []
+
+  for (let i = sanitized.length - 1; i >= 0; i--) {
+    const msg = sanitized[i]
+    const msgTokens = estimateTokens(msg)
+    if (estimatedTokens + msgTokens > maxTokenBudget && result.length > 0) {
+      break
+    }
+    estimatedTokens += msgTokens
+    result.unshift(msg)
+  }
+
+  if (
+    result.length > 0 &&
+    result[0].role !== 'user' &&
+    result[0].role !== 'system'
+  ) {
+    const firstUserIdx = result.findIndex(
+      (m) => m.role === 'user' || m.role === 'system'
+    )
+    if (firstUserIdx > 0) {
+      return result.slice(firstUserIdx)
+    }
+  }
+
+  return result
+}
+
+function sanitizeToolMessages(messages: AIMessage[]): AIMessage[] {
+  const result: AIMessage[] = []
+
+  for (let i = 0; i < messages.length; i++) {
+    const msg = messages[i]
+
+    if (msg.role === 'assistant' && msg.toolCalls?.length) {
+      const next = messages[i + 1]
+      if (!next || next.role !== 'tool' || !next.toolResults?.length) {
+        if (msg.content) {
+          result.push({ ...msg, toolCalls: undefined })
+        }
+        continue
+      }
+
+      const toolCallIds = new Set(msg.toolCalls.map((tc) => tc.id))
+      const resultIds = new Set(next.toolResults.map((tr) => tr.id))
+      const allMatched = [...toolCallIds].every((id) => resultIds.has(id))
+
+      if (!allMatched) {
+        if (msg.content) {
+          result.push({ ...msg, toolCalls: undefined })
+        }
+        result.push({
+          ...next,
+          toolResults:
+            next.toolResults.filter((tr) => !toolCallIds.has(tr.id)) ||
+            undefined,
+        })
+        i++
+        continue
+      }
+    }
+
+    result.push(msg)
+  }
+
+  return result.filter(
+    (m) => m.role !== 'tool' || (m.toolResults && m.toolResults.length > 0)
+  )
+}
+
+function estimateTokens(msg: AIMessage): number {
+  let chars = 0
+  if (msg.content) chars += msg.content.length
+  if (msg.toolCalls) chars += JSON.stringify(msg.toolCalls).length
+  if (msg.toolResults) chars += JSON.stringify(msg.toolResults).length
+  return Math.ceil(chars / 4)
+}
+
+export function parseWorkingMemory(
+  text: string
+): Record<string, unknown> | null {
+  const match = text.match(/<working_memory>([\s\S]*?)<\/working_memory>/)
+  if (!match) return null
+  try {
+    return JSON.parse(match[1].trim())
+  } catch {
+    return null
+  }
+}
+
+export function stripWorkingMemoryTags(text: string): string {
+  return text.replace(/<working_memory>[\s\S]*?<\/working_memory>/g, '').trim()
+}

package/src/wirings/ai-agent/ai-agent-model-config.test.ts

@@ -0,0 +1,115 @@
+import { describe, test, beforeEach } from 'node:test'
+import assert from 'node:assert'
+import { resolveModelConfig } from './ai-agent-model-config.js'
+import { resetPikkuState, pikkuState } from '../../pikku-state.js'
+
+beforeEach(() => {
+  resetPikkuState()
+})
+
+describe('resolveModelConfig', () => {
+  test('should return agent model when no config', () => {
+    const result = resolveModelConfig('testAgent', {
+      model: 'anthropic/claude-3',
+      temperature: 0.7,
+    })
+    assert.strictEqual(result.model, 'anthropic/claude-3')
+    assert.strictEqual(result.temperature, 0.7)
+  })
+
+  test('should resolve model alias to string', () => {
+    pikkuState(null, 'models', 'config', {
+      models: { fast: 'anthropic/claude-haiku' },
+    } as any)
+
+    const result = resolveModelConfig('testAgent', { model: 'fast' })
+    assert.strictEqual(result.model, 'anthropic/claude-haiku')
+  })
+
+  test('should resolve model alias with temperature and maxSteps', () => {
+    pikkuState(null, 'models', 'config', {
+      models: {
+        smart: {
+          model: 'anthropic/claude-opus',
+          temperature: 0.3,
+          maxSteps: 10,
+        },
+      },
+    } as any)
+
+    const result = resolveModelConfig('testAgent', { model: 'smart' })
+    assert.strictEqual(result.model, 'anthropic/claude-opus')
+    assert.strictEqual(result.temperature, 0.3)
+    assert.strictEqual(result.maxSteps, 10)
+  })
+
+  test('should throw for unknown model alias', () => {
+    pikkuState(null, 'models', 'config', {
+      models: {},
+    } as any)
+
+    assert.throws(() => resolveModelConfig('testAgent', { model: 'unknown' }), {
+      message: "Unknown model alias 'unknown'.",
+    })
+  })
+
+  test('should pass through model with slash (not an alias)', () => {
+    pikkuState(null, 'models', 'config', {
+      models: {},
+    } as any)
+
+    const result = resolveModelConfig('testAgent', { model: 'openai/gpt-4' })
+    assert.strictEqual(result.model, 'openai/gpt-4')
+  })
+
+  test('should use agent override model', () => {
+    pikkuState(null, 'models', 'config', {
+      models: { fast: 'anthropic/haiku' },
+      agentOverrides: { testAgent: { model: 'fast' } },
+    } as any)
+
+    const result = resolveModelConfig('testAgent', {
+      model: 'anthropic/default',
+    })
+    assert.strictEqual(result.model, 'anthropic/haiku')
+  })
+
+  test('should use agent override temperature over alias', () => {
+    pikkuState(null, 'models', 'config', {
+      models: { smart: { model: 'anthropic/opus', temperature: 0.5 } },
+      agentOverrides: { testAgent: { model: 'smart', temperature: 0.1 } },
+    } as any)
+
+    const result = resolveModelConfig('testAgent', {
+      model: 'anthropic/default',
+    })
+    assert.strictEqual(result.temperature, 0.1)
+  })
+
+  test('should use agentDefaults as fallback', () => {
+    pikkuState(null, 'models', 'config', {
+      models: {},
+      agentDefaults: { temperature: 0.5, maxSteps: 20 },
+    } as any)
+
+    const result = resolveModelConfig('testAgent', {
+      model: 'anthropic/claude-3',
+    })
+    assert.strictEqual(result.temperature, 0.5)
+    assert.strictEqual(result.maxSteps, 20)
+  })
+
+  test('should prioritize: override > alias > defaults > agent', () => {
+    pikkuState(null, 'models', 'config', {
+      models: { fast: { model: 'anthropic/haiku', temperature: 0.3 } },
+      agentDefaults: { temperature: 0.5 },
+      agentOverrides: { testAgent: { temperature: 0.1 } },
+    } as any)
+
+    const result = resolveModelConfig('testAgent', {
+      model: 'fast',
+      temperature: 0.9,
+    })
+    assert.strictEqual(result.temperature, 0.1)
+  })
+})

package/src/wirings/ai-agent/ai-agent-model-config.ts

@@ -0,0 +1,43 @@
+import { pikkuState } from '../../pikku-state.js'
+
+export function resolveModelConfig(
+  agentName: string,
+  agent: { model: string; temperature?: number; maxSteps?: number }
+): { model: string; temperature?: number; maxSteps?: number } {
+  let config: any = {}
+  try {
+    config = pikkuState(null, 'models', 'config') ?? {}
+  } catch {}
+  const models = config.models ?? {}
+  const defaults = config.agentDefaults
+  const agentOverride = config.agentOverrides?.[agentName]
+
+  let rawModel = agentOverride?.model ?? agent.model
+
+  let aliasTemperature: number | undefined
+  let aliasMaxSteps: number | undefined
+  if (!rawModel.includes('/')) {
+    const entry = models[rawModel]
+    if (!entry) throw new Error(`Unknown model alias '${rawModel}'.`)
+    if (typeof entry === 'string') {
+      rawModel = entry
+    } else {
+      rawModel = entry.model
+      aliasTemperature = entry.temperature
+      aliasMaxSteps = entry.maxSteps
+    }
+  }
+
+  const temperature =
+    agentOverride?.temperature ??
+    aliasTemperature ??
+    defaults?.temperature ??
+    agent.temperature
+  const maxSteps =
+    agentOverride?.maxSteps ??
+    aliasMaxSteps ??
+    defaults?.maxSteps ??
+    agent.maxSteps
+
+  return { model: rawModel, temperature, maxSteps }
+}