feat: add tasuku for task hierarchy logging

examples/hf0-demo.ts

@@ -78,8 +78,7 @@ async function main() {
     )
     .call({ topic, questions })
 
-  console.log('\n\n\n')
+  console.log(`\n\n\n${res}\n\n\n`)
-  console.log(res)
 }
 
 main()

examples/search-and-crawl.ts

@@ -15,7 +15,7 @@ async function main() {
       messages: [
         {
           role: 'system',
-          content: `You are a McKinsey analyst who is an expert at writing executive summaries. Always respond using markdown.`
+          content: `You are a McKinsey analyst who is an expert at writing executive summaries. Always cite your sources and respond using markdown.`
         },
         {
           role: 'user',
@@ -32,8 +32,7 @@ async function main() {
     )
     .call({ topic })
 
-  console.log('\n\n\n')
+  console.log(`\n\n\n${res}\n\n\n`)
-  console.log(res)
 }
 
 main()

examples/summarize-news.ts

@@ -22,7 +22,7 @@ async function main() {
       topic
     })
 
-  console.log('\n\n\n' + res)
+  console.log(`\n\n\n${res}\n\n\n`)
 }
 
 main()

package.json

@@ -61,6 +61,7 @@
     "pino": "^8.14.1",
     "pino-pretty": "^10.0.0",
     "quick-lru": "^6.1.1",
+    "tasuku": "^2.0.1",
     "ts-dedent": "^2.2.0",
     "uuid": "^9.0.0",
     "zod": "^3.21.4",

pnpm-lock.yaml

@@ -1,4 +1,4 @@
-lockfileVersion: '6.1'
+lockfileVersion: '6.0'
 
 settings:
   autoInstallPeers: true
@@ -74,6 +74,9 @@ dependencies:
   quick-lru:
     specifier: ^6.1.1
     version: 6.1.1
+  tasuku:
+    specifier: ^2.0.1
+    version: 2.0.1
   ts-dedent:
     specifier: ^2.2.0
     version: 2.2.0
@@ -1021,6 +1024,10 @@ packages:
     resolution: {integrity: sha512-kNnC1GFBLuhImSnV7w4njQkUiJi0ZXUycu1rUaouPqiKlXkh77JKgdRnTAp1x5eBwcIwbtI+3otwzuIDEuDoxQ==}
     dev: true
 
+  /@types/yoga-layout@1.9.2:
+    resolution: {integrity: sha512-S9q47ByT2pPvD65IvrWp7qppVMpk9WGMbVq9wbWZOHg6tnXSD4vyhao6nOSBwwfDdV2p3Kx9evA9vI+XWTfDvw==}
+    dev: false
+
   /@typescript-eslint/eslint-plugin@5.59.11(@typescript-eslint/parser@5.59.11)(eslint@8.43.0)(typescript@5.1.3):
     resolution: {integrity: sha512-XxuOfTkCUiOSyBWIvHlUraLw/JT/6Io1365RO6ZuI88STKMavJZPNMU0lFcUTeQXEhHiv64CbxYxBNoDVSmghg==}
     engines: {node: ^12.22.0 || ^14.17.0 || >=16.0.0}
@@ -4247,6 +4254,12 @@ packages:
     engines: {node: '>= 0.4'}
     dev: true
 
+  /tasuku@2.0.1:
+    resolution: {integrity: sha512-BXWDEJzpC1mUiOz5Csba85LS93o9a5pGKRTArLiXJZ2HGF/mXHIl+4SBF706Yxqg+GlJDQurvLxds8tC7EwyRA==}
+    dependencies:
+      yoga-layout-prebuilt: 1.10.0
+    dev: false
+
   /temp-dir@3.0.0:
     resolution: {integrity: sha512-nHc6S/bwIilKHNRgK/3jlhDoIHcp45YgyiwcAk46Tr0LfEqGBVpmiAyuiuxeVE44m3mXnEeVhaipLOEWmH+Njw==}
     engines: {node: '>=14.16'}
@@ -4635,6 +4648,13 @@ packages:
     engines: {node: '>=12.20'}
     dev: true
 
+  /yoga-layout-prebuilt@1.10.0:
+    resolution: {integrity: sha512-YnOmtSbv4MTf7RGJMK0FvZ+KD8OEe/J5BNnR0GHhD8J/XcG/Qvxgszm0Un6FTHWW4uHlTgP0IztiXQnGyIR45g==}
+    engines: {node: '>=8'}
+    dependencies:
+      '@types/yoga-layout': 1.9.2
+    dev: false
+
   /zod-to-json-schema@3.21.2(zod@3.21.4):
     resolution: {integrity: sha512-02yfKymfmIf2rM/5LYGlyw0daEel/f3MsSGMNJZWWf44ato+Y+diFugOpDtgvEUn3cYM5oDAGWW2NHeSD4mByw==}
     peerDependencies:

src/agentic.ts

@@ -1,4 +1,5 @@
 import defaultKy from 'ky'
+import taskTracker from 'tasuku'
 import { SetOptional } from 'type-fest'
 
 import * as types from './types'
@@ -12,6 +13,7 @@ import { defaultIDGeneratorFn, isFunction, isString } from './utils'
 export class Agentic {
   protected _ky: types.KyInstance
   protected _logger: types.Logger
+  protected _taskTracker: types.TaskTracker
 
   protected _openai?: types.openai.OpenAIClient
   protected _anthropic?: types.anthropic.Client
@@ -35,6 +37,7 @@ export class Agentic {
     idGeneratorFn?: types.IDGeneratorFunction
     logger?: types.Logger
     ky?: types.KyInstance
+    taskTracker?: types.TaskTracker
   }) {
     // TODO: This is a bit hacky, but we're doing it to have a slightly nicer API
     // for the end developer when creating subclasses of `BaseTask` to use as
@@ -48,6 +51,7 @@ export class Agentic {
 
     this._ky = opts.ky ?? defaultKy
     this._logger = opts.logger ?? defaultLogger
+    this._taskTracker = opts.taskTracker ?? taskTracker
 
     this._openaiModelDefaults = {
       provider: 'openai',
@@ -94,6 +98,10 @@ export class Agentic {
     return this._logger
   }
 
+  public get taskTracker(): types.TaskTracker {
+    return this._taskTracker
+  }
+
   public get humanFeedbackDefaults() {
     return this._humanFeedbackDefaults
   }

src/llms/chat.ts

@@ -318,7 +318,10 @@ export abstract class BaseChatCompletion<
         )
 
         // TODO: handle sub-task errors gracefully
-        const toolCallResponse = await tool.callWithMetadata(functionArguments)
+        const toolCallResponse = await tool.callWithMetadata(
+          functionArguments,
+          ctx
+        )
 
         this._logger.info(
           {

src/task.ts

@@ -9,7 +9,11 @@ import {
   HumanFeedbackOptions,
   HumanFeedbackType
 } from './human-feedback'
-import { defaultIDGeneratorFn, isValidTaskIdentifier } from './utils'
+import {
+  defaultIDGeneratorFn,
+  isValidTaskIdentifier,
+  stringifyForDebugging
+} from './utils'
 
 /**
  * A `Task` is an async function call that may be non-deterministic. It has
@@ -174,6 +178,31 @@ export abstract class BaseTask<
   ): Promise<types.TaskResponse<TOutput>> {
     this.validate()
 
+    let _resolve: (value: unknown) => void | undefined
+    let _reject: (err: Error) => void | undefined
+    let _taskInnerAPI: types.TaskTrackerInnerAPI | undefined
+
+    const taskP = new Promise((resolve, reject) => {
+      _resolve = resolve
+      _reject = reject
+    })
+
+    const title = `${this.nameForModel}(${stringifyForDebugging(input, {
+      maxLength: 120
+    })})`
+
+    if (parentCtx?.tracker) {
+      parentCtx.tracker.task(title, (taskInnerAPI) => {
+        _taskInnerAPI = taskInnerAPI
+        return taskP
+      })
+    } else {
+      this._agentic!.taskTracker(title, (taskInnerAPI) => {
+        _taskInnerAPI = taskInnerAPI
+        return taskP
+      })
+    }
+
     this._logger.info({ input }, `Task call "${this.nameForHuman}"`)
 
     if (this.inputSchema) {
@@ -189,6 +218,7 @@ export abstract class BaseTask<
     const ctx: types.TaskCallContext<TInput> = {
       input,
       attemptNumber: 0,
+      tracker: _taskInnerAPI!,
       metadata: {
         taskName: this.nameForModel,
         taskId: this.id,
@@ -198,74 +228,91 @@ export abstract class BaseTask<
       }
     }
 
-    for (const preHook of this._preHooks) {
+    try {
-      await preHook(ctx)
+      for (const preHook of this._preHooks) {
-    }
+        await preHook(ctx)
-
-    const result = await pRetry(
-      async () => {
-        const result = await this._call(ctx)
-
-        for (const postHook of this._postHooks) {
-          await postHook(result, ctx)
-        }
-
-        return result
-      },
-      {
-        ...this._retryConfig,
-        onFailedAttempt: async (err: FailedAttemptError) => {
-          this._logger.warn(
-            err,
-            `Task error "${this.nameForHuman}" failed attempt ${
-              err.attemptNumber
-            }${input ? ': ' + JSON.stringify(input) : ''}`
-          )
-
-          if (this._retryConfig.onFailedAttempt) {
-            await Promise.resolve(this._retryConfig.onFailedAttempt(err))
-          }
-
-          // TODO: log this task error
-          ctx.attemptNumber = err.attemptNumber + 1
-          ctx.metadata.error = err
-
-          if (err instanceof errors.ZodOutputValidationError) {
-            ctx.retryMessage = err.message
-            return
-          } else if (err instanceof errors.OutputValidationError) {
-            ctx.retryMessage = err.message
-            return
-          } else if (err instanceof errors.HumanFeedbackDeclineError) {
-            ctx.retryMessage = err.message
-            return
-          } else if (
-            err instanceof errors.KyTimeoutError ||
-            err instanceof errors.TimeoutError ||
-            (err as any).name === 'TimeoutError'
-          ) {
-            // TODO
-            return
-          } else if ((err.cause as any)?.code === 'UND_ERR_HEADERS_TIMEOUT') {
-            // TODO: This is a pretty common OpenAI error, and I think it either has
-            // to do with OpenAI's servers being flaky or the combination of Node.js
-            // `undici` and OpenAI's HTTP requests. Either way, let's just retry the
-            // task for now.
-            return
-          } else {
-            throw err
-          }
-        }
       }
-    )
 
-    ctx.metadata.success = true
+      const result = await pRetry(
-    ctx.metadata.numRetries = ctx.attemptNumber
+        async () => {
-    ctx.metadata.error = undefined
+          const result = await this._call(ctx)
 
-    return {
+          for (const postHook of this._postHooks) {
-      result,
+            await postHook(result, ctx)
-      metadata: ctx.metadata
+          }
+
+          return result
+        },
+        {
+          ...this._retryConfig,
+          onFailedAttempt: async (err: FailedAttemptError) => {
+            this._logger.warn(
+              err,
+              `Task error "${this.nameForHuman}" failed attempt ${
+                err.attemptNumber
+              }${input ? ': ' + JSON.stringify(input) : ''}`
+            )
+
+            // const error = `error ${err.attemptNumber}: ${err.message}`
+            // const errorAsString =
+            //   error.length > 80 ? error.slice(0, 80 - 3) + '...' : error
+            // ctx.tracker.setStatus(errorAsString)
+
+            if (this._retryConfig.onFailedAttempt) {
+              await Promise.resolve(this._retryConfig.onFailedAttempt(err))
+            }
+
+            // TODO: log this task error
+            ctx.attemptNumber = err.attemptNumber + 1
+            ctx.metadata.error = err
+
+            if (err instanceof errors.ZodOutputValidationError) {
+              ctx.retryMessage = err.message
+              return
+            } else if (err instanceof errors.OutputValidationError) {
+              ctx.retryMessage = err.message
+              return
+            } else if (err instanceof errors.HumanFeedbackDeclineError) {
+              ctx.retryMessage = err.message
+              return
+            } else if (
+              err instanceof errors.KyTimeoutError ||
+              err instanceof errors.TimeoutError ||
+              (err as any).name === 'TimeoutError'
+            ) {
+              // TODO
+              return
+            } else if ((err.cause as any)?.code === 'UND_ERR_HEADERS_TIMEOUT') {
+              // TODO: This is a pretty common OpenAI error, and I think it either has
+              // to do with OpenAI's servers being flaky or the combination of Node.js
+              // `undici` and OpenAI's HTTP requests. Either way, let's just retry the
+              // task for now.
+              return
+            } else {
+              throw err
+            }
+          }
+        }
+      )
+
+      ctx.metadata.success = true
+      ctx.metadata.numRetries = ctx.attemptNumber
+      ctx.metadata.error = undefined
+
+      ctx.tracker.setOutput(stringifyForDebugging(result, { maxLength: 100 }))
+
+      // @ts-expect-error "_resolve" should be defined above
+      _resolve(result)
+
+      return {
+        result,
+        metadata: ctx.metadata
+      }
+    } catch (err: any) {
+      // @ts-expect-error "_reject" should be defined above
+      _reject(err)
+
+      throw err
     }
   }
 

src/types.ts

@@ -2,6 +2,10 @@ import * as anthropic from '@anthropic-ai/sdk'
 import * as openai from 'openai-fetch'
 import ky from 'ky'
 import type { Options as RetryOptions } from 'p-retry'
+import type {
+  Task as TaskTracker,
+  TaskInnerAPI as TaskTrackerInnerAPI
+} from 'tasuku'
 import type { JsonObject, Jsonifiable } from 'type-fest'
 import { SafeParseReturnType, ZodType, ZodTypeAny, output, z } from 'zod'
 
@@ -15,6 +19,7 @@ import type { BaseTask } from './task'
 
 export { anthropic, openai }
 
+export type { TaskTracker, TaskTrackerInnerAPI }
 export type { Jsonifiable, Logger }
 export type KyInstance = typeof ky
 
@@ -40,6 +45,8 @@ export interface BaseTaskOptions {
   retryConfig?: RetryConfig
   id?: string
 
+  taskTracker?: TaskTracker
+
   // TODO
   // caching config
   // logging config
@@ -156,6 +163,7 @@ export interface TaskCallContext<
 > {
   input: TInput
   retryMessage?: string
+  tracker: TaskTrackerInnerAPI
 
   attemptNumber: number
   metadata: TMetadata

src/utils.ts

@@ -107,6 +107,29 @@ export function chunkMultipleStrings(
   return textSections.map((section) => chunkString(section, maxLength)).flat()
 }
 
+export function stringifyForDebugging(
+  json?: types.Jsonifiable | void,
+  {
+    maxLength
+  }: {
+    maxLength?: number
+  } = {}
+): string {
+  if (json === undefined) {
+    return ''
+  }
+
+  const out = stringifyForModel(json)
+
+  if (maxLength) {
+    return out.length > maxLength
+      ? out.substring(0, Math.max(0, maxLength - 1)) + '…'
+      : out
+  } else {
+    return out
+  }
+}
+
 /**
  * Stringifies a JSON value for use in an LLM prompt.
  *
@@ -114,9 +137,13 @@ export function chunkMultipleStrings(
  * @returns stringified value with all double quotes around object keys removed
  */
 export function stringifyForModel(
-  json: types.Jsonifiable,
+  json: types.Jsonifiable | void,
   omit: string[] = []
 ): string {
+  if (json === undefined) {
+    return ''
+  }
+
   const UNIQUE_PREFIX = defaultIDGeneratorFn()
   return (
     JSON.stringify(json, replacer)