effect-assistant 0.0.0-stage → 0.1.0

package/src/AssistantActivity.ts

@@ -0,0 +1,150 @@
+/**
+ * This module provides functions that determine what is working in a chat right now, read from
+ * the chat's log.
+ *
+ * A {@link State} records whether the assistant is running a turn or is about to start one,
+ * and which specialist runs are working. Fold a chat's entries into it with {@link apply}, or
+ * build it from a list of entries with {@link fromEntries}, then check it with
+ * {@link isWorking}.
+ *
+ * @since 0.1.0
+ */
+import * as Match from "effect/Match"
+import * as Predicate from "effect/Predicate"
+import * as AssistantEvent from "./AssistantEvent.ts"
+import type { RunId } from "./RunId.ts"
+
+/**
+ * Represents a specialist run that is working.
+ *
+ * @category models
+ * @since 0.1.0
+ */
+export interface SpecialistRun {
+  readonly runId: RunId
+  readonly specialist: string
+  /** The goal the run is working on. */
+  readonly goal: string
+  /** When the run started, in epoch milliseconds. */
+  readonly since: number
+}
+
+/**
+ * Represents what is working in a chat.
+ *
+ * @category models
+ * @since 0.1.0
+ */
+export interface State {
+  /** When the running turn started, in epoch milliseconds, or `undefined` when no turn runs. */
+  readonly turnSince: number | undefined
+  /**
+   * Whether an input arrived that no finished turn has read. When this is true, a turn is about
+   * to start.
+   */
+  readonly queued: boolean
+  /** The `seq` of the newest input. */
+  readonly lastInputSeq: number
+  /** Specialist runs that are working, by run id. */
+  readonly specialists: ReadonlyMap<RunId, SpecialistRun>
+}
+
+/**
+ * The activity of a chat with an empty log, where nothing is working.
+ *
+ * @category constructors
+ * @since 0.1.0
+ */
+export const empty: State = {
+  turnSince: undefined,
+  queued: false,
+  lastInputSeq: 0,
+  specialists: new Map()
+}
+
+const withoutRun = (activity: State, runId: RunId): State => {
+  if (!activity.specialists.has(runId)) return activity
+
+  const specialists = new Map(activity.specialists)
+  specialists.delete(runId)
+
+  return { ...activity, specialists }
+}
+
+/**
+ * Represents the options for folding a chat's entries into its activity.
+ *
+ * @category models
+ * @since 0.1.0
+ */
+export interface Options {
+  /**
+   * The tags of the app's events that wake a turn (see `AppEvent`). These events count as inputs
+   * too. Defaults to none.
+   */
+  readonly wakes?: ReadonlySet<string> | undefined
+}
+
+/**
+ * Returns the activity after one more entry. Fold a chat's entries in order, or use
+ * {@link fromEntries}.
+ *
+ * @category combinators
+ * @since 0.1.0
+ */
+export const apply = (activity: State, entry: AssistantEvent.AnyEntry, options?: Options): State => {
+  const event = entry.event
+
+  if (!AssistantEvent.isLibraryEvent(event)) {
+    return options?.wakes?.has(event._tag) === true ? { ...activity, queued: true, lastInputSeq: entry.seq } : activity
+  }
+
+  if (AssistantEvent.isInput(event)) {
+    // A finished run's report is an input, and ends the run.
+    const current = Predicate.isTagged(event, "SpecialistFinished") ? withoutRun(activity, event.runId) : activity
+
+    return AssistantEvent.isWaking(event) ? { ...current, queued: true, lastInputSeq: entry.seq } : current
+  }
+
+  return Match.value(event).pipe(
+    Match.tag("TurnStarted", (): State => ({ ...activity, turnSince: entry.at })),
+    // Turns run one at a time, so an end closes whichever turn is open, also one a restart cut off.
+    Match.tag("TurnCompleted", "TurnFailed", (ended): State => ({
+      ...activity,
+      turnSince: undefined,
+      queued: activity.lastInputSeq > ended.throughSeq
+    })),
+    Match.tag("SpecialistStarted", (started): State => ({
+      ...activity,
+      specialists: new Map(activity.specialists).set(started.runId, {
+        runId: started.runId,
+        specialist: started.specialist,
+        goal: started.goal,
+        since: entry.at
+      })
+    })),
+    // Deleting a chat stops everything in it, and nothing starts again.
+    Match.tag("ChatDeleted", (): State => ({ ...empty, lastInputSeq: activity.lastInputSeq })),
+    Match.orElse(() => activity)
+  )
+}
+
+/**
+ * Returns what is working in a chat after the given entries, folded in order from
+ * {@link empty}.
+ *
+ * @category constructors
+ * @since 0.1.0
+ */
+export const fromEntries = (entries: ReadonlyArray<AssistantEvent.AnyEntry>, options?: Options): State =>
+  entries.reduce((activity, entry) => apply(activity, entry, options), empty)
+
+/**
+ * Returns whether anything is working: a turn that is running or about to start, or a specialist
+ * run.
+ *
+ * @category guards
+ * @since 0.1.0
+ */
+export const isWorking = (activity: State): boolean =>
+  activity.turnSince !== undefined || activity.queued || activity.specialists.size > 0

package/src/AssistantEngine.ts

@@ -0,0 +1,343 @@
+/**
+ * This module provides the `AssistantEngine` service, which runs assistants' chats, and its two
+ * implementations.
+ *
+ * An assistant's `toLayer` registers the assistant with the engine, and the assistant's `submit`
+ * and `subscribe` go through it. Pick one engine for the app, as with a workflow engine.
+ * {@link layerMemory} runs chats in the process, for tests and single-process apps.
+ * {@link layerCluster} runs each chat as a cluster entity, so a turn cut off by a crash resumes.
+ *
+ * @since 0.1.0
+ */
+import * as ClusterSchema from "effect/cluster/ClusterSchema"
+import * as Entity from "effect/cluster/Entity"
+import * as Sharding from "effect/cluster/Sharding"
+import * as Context from "effect/Context"
+import * as Data from "effect/Data"
+import * as Effect from "effect/Effect"
+import * as Exit from "effect/Exit"
+import * as Layer from "effect/Layer"
+import * as MutableHashMap from "effect/MutableHashMap"
+import * as Option from "effect/Option"
+import * as Predicate from "effect/Predicate"
+import * as RcMap from "effect/RcMap"
+import * as Rpc from "effect/rpc/Rpc"
+import * as Schema from "effect/Schema"
+import * as Scope from "effect/Scope"
+import * as Stream from "effect/Stream"
+import * as AssistantEvent from "./AssistantEvent.ts"
+import { ChatId } from "./ChatId.ts"
+import { ChatLog } from "./ChatLog.ts"
+import * as Chat from "./internal/chat.ts"
+import * as History from "./internal/history.ts"
+import type * as Turns from "./internal/turns.ts"
+
+/**
+ * Service that runs assistants' chats. An assistant's `toLayer` registers the assistant with it,
+ * and the assistant's `submit` and `subscribe` use it.
+ *
+ * @category services
+ * @since 0.1.0
+ */
+export class AssistantEngine extends Context.Service<
+  AssistantEngine,
+  {
+    /** Registers an assistant for as long as the scope is open. */
+    readonly register: (registration: Turns.Registration) => Effect.Effect<void, never, Scope.Scope>
+    /**
+     * Stores an input in one of the assistant's chats. The assistant works on it in the
+     * background.
+     */
+    readonly submit: (
+      assistant: string,
+      chatId: ChatId,
+      input: AssistantEvent.Input | AssistantEvent.StoredAppEvent
+    ) => Effect.Effect<void>
+    /** Streams the chat's stored entries after the `after` seq, then live ones. */
+    readonly subscribe: (assistant: string, chatId: ChatId, after: number) => Stream.Stream<AssistantEvent.Entry>
+    /**
+     * Deletes one of the assistant's chats. The chat's log records the deletion, and everything
+     * that works in the chat stops.
+     */
+    readonly delete: (assistant: string, chatId: ChatId) => Effect.Effect<void>
+  }
+>()("effect-assistant/AssistantEngine") {}
+
+class ChatKey extends Data.Class<{ readonly assistant: string; readonly chatId: ChatId }> {}
+
+const notRegistered = (assistant: string) =>
+  Effect.die(`The ${assistant} assistant isn't registered: provide its toLayer.`)
+
+/**
+ * Layer that runs chats in the current process. Each chat's turns run as fibers, one at a time.
+ *
+ * **Details**
+ *
+ * With `ChatLog.layerSql`, the engine resumes chats at startup. When an assistant registers, the
+ * engine finds the chats that assistant was working on when the process stopped (a turn cut off
+ * partway, or an input that no turn read) and runs a turn for each. Run this engine in one process
+ * only, because two processes would run the same chat and write its log at the same time.
+ *
+ * **Example** (Running chats in memory)
+ *
+ * ```ts
+ * import { Assistant, AssistantEngine, ChatLog } from "effect-assistant"
+ * import type * as LanguageModel from "effect/ai/LanguageModel"
+ * import * as Layer from "effect/Layer"
+ *
+ * declare const ModelLive: Layer.Layer<LanguageModel.LanguageModel>
+ *
+ * const MyAssistant = Assistant.make("my-assistant")
+ *
+ * const AppLive = MyAssistant.toLayer((turn) => turn.loop({ instructions: "Be brief." })).pipe(
+ *   Layer.provide(ModelLive),
+ *   Layer.provideMerge(AssistantEngine.layerMemory),
+ *   Layer.provideMerge(ChatLog.layerMemory)
+ * )
+ * ```
+ *
+ * @see {@link layerCluster} for running each chat as a cluster entity
+ *
+ * @category layers
+ * @since 0.1.0
+ */
+export const layerMemory: Layer.Layer<AssistantEngine, never, ChatLog> = Layer.effect(
+  AssistantEngine,
+  Effect.gen(function* () {
+    const context = yield* Effect.context<ChatLog>()
+    const log = yield* ChatLog
+    const registered = new Map<string, Turns.Registration>()
+    // A chat's keep-alive: a hold on its entry, so it stays open while runs or turns work in the background.
+    const held = MutableHashMap.empty<ChatKey, Scope.Closeable>()
+
+    // One entry per open chat, like a cluster entity: opened when used, closed five minutes after the last use.
+    const chats: RcMap.RcMap<ChatKey, Chat.Chat> = yield* RcMap.make({
+      lookup: Effect.fnUntraced(function* (key: ChatKey) {
+        const registration = registered.get(key.assistant)
+
+        if (registration === undefined) return yield* notRegistered(key.assistant)
+
+        return yield* Chat.make({
+          registration,
+          chatId: key.chatId,
+          keepAlive: (enabled) => keepAlive(key, enabled)
+        }).pipe(Effect.provide(context))
+      }),
+      idleTimeToLive: "5 minutes"
+    })
+
+    // On: a hold on the chat's entry, in a scope of its own. Off: closing that scope lets the entry go idle.
+    const keepAlive = Effect.fnUntraced(function* (key: ChatKey, enabled: boolean) {
+      if (enabled) {
+        const scope = yield* Scope.make()
+
+        yield* Effect.provideService(RcMap.get(chats, key), Scope.Scope, scope)
+
+        return MutableHashMap.set(held, key, scope)
+      }
+
+      const scope = MutableHashMap.get(held, key)
+
+      MutableHashMap.remove(held, key)
+
+      if (Option.isSome(scope)) yield* Scope.close(scope.value, Exit.void)
+    })
+
+    const withChat = Effect.fnUntraced(function* <A>(
+      assistant: string,
+      chatId: ChatId,
+      f: (chat: Chat.Chat) => Effect.Effect<A>
+    ) {
+      const chat = yield* RcMap.get(chats, new ChatKey({ assistant, chatId }))
+
+      return yield* f(chat)
+    }, Effect.scoped)
+
+    // A chat owes a turn when its last turn was cut off, or an input came after.
+    const resumeAll = Effect.fnUntraced(function* (registration: Turns.Registration) {
+      const { name, app } = registration.definition
+
+      for (const chatId of yield* log.unfinished) {
+        const entries = yield* log.read(chatId)
+
+        const lastTurn = entries.findLast(({ event }) => Predicate.isTagged(event, "TurnStarted"))?.event
+        const ours = Predicate.isTagged(lastTurn, "TurnStarted") && lastTurn.assistant === name
+
+        if (ours && History.needsTurn(entries, app.wakes)) yield* withChat(name, chatId, (chat) => chat.resume)
+      }
+    })
+
+    return AssistantEngine.of({
+      register: Effect.fn("AssistantEngine.register")(function* (registration: Turns.Registration) {
+        registered.set(registration.definition.name, registration)
+        // Chats this assistant was working on when the process stopped: pick them up again.
+        yield* Effect.forkScoped(resumeAll(registration))
+      }),
+      submit: (assistant, chatId, input) => withChat(assistant, chatId, (chat) => chat.submitInBackground(input)),
+      // The chat deletes itself, so a chat that's open stops; the entry goes, and the next open finds it deleted.
+      delete: Effect.fn("AssistantEngine.delete")(function* (assistant: string, chatId: ChatId) {
+        yield* withChat(assistant, chatId, (chat) => chat.delete)
+        yield* RcMap.invalidate(chats, new ChatKey({ assistant, chatId }))
+      }),
+      // The subscriber holds the chat open while it streams.
+      subscribe: (assistant, chatId, after) =>
+        Stream.unwrap(
+          Effect.gen(function* () {
+            const chat = yield* RcMap.get(chats, new ChatKey({ assistant, chatId }))
+
+            return chat.stream(after)
+          })
+        )
+    })
+  })
+)
+
+interface EntityClient {
+  readonly submit: (chatId: ChatId, input: AssistantEvent.Input | AssistantEvent.StoredAppEvent) => Effect.Effect<void>
+  readonly subscribe: (chatId: ChatId, after: number) => Stream.Stream<AssistantEvent.Entry>
+  readonly delete: (chatId: ChatId) => Effect.Effect<void>
+}
+
+const Submit = Rpc.make("Submit", {
+  payload: { input: Schema.Union([AssistantEvent.Input, AssistantEvent.StoredAppEvent]) }
+}).annotate(ClusterSchema.Persisted, true)
+
+const Subscribe = Rpc.make("Subscribe", {
+  payload: { after: Schema.Int },
+  success: AssistantEvent.Entry,
+  stream: true
+})
+
+const Delete = Rpc.make("Delete")
+
+const chatEntity = (assistant: string) => Entity.make(`Assistant/${assistant}`, [Submit, Subscribe, Delete])
+
+/**
+ * Layer that runs chats as cluster entities, one entity per chat, keyed by chat id.
+ *
+ * **Details**
+ *
+ * The layer requires `Sharding`. Use a `SingleRunner` for one process with SQL storage, a runner
+ * of the app's cluster, or `TestRunner` in tests.
+ *
+ * **Example** (Running chats on a single runner with Postgres)
+ *
+ * ```ts
+ * import * as NodeCrypto from "@effect/platform-node/NodeCrypto"
+ * import * as PgClient from "@effect/sql-pg/PgClient"
+ * import { Assistant, AssistantEngine, ChatLog } from "effect-assistant"
+ * import type * as LanguageModel from "effect/ai/LanguageModel"
+ * import * as SingleRunner from "effect/cluster/SingleRunner"
+ * import * as Config from "effect/Config"
+ * import * as Layer from "effect/Layer"
+ *
+ * declare const ModelLive: Layer.Layer<LanguageModel.LanguageModel>
+ *
+ * const MyAssistant = Assistant.make("my-assistant")
+ *
+ * const AppLive = MyAssistant.toLayer((turn) => turn.loop({ instructions: "Be brief." })).pipe(
+ *   Layer.provide(ModelLive),
+ *   Layer.provideMerge(AssistantEngine.layerCluster),
+ *   Layer.provideMerge(
+ *     Layer.merge(
+ *       ChatLog.layerSql,
+ *       SingleRunner.layer({ runnerStorage: "sql" }).pipe(Layer.provide(NodeCrypto.layer))
+ *     )
+ *   ),
+ *   Layer.provide(PgClient.layerConfig({ url: Config.Redacted("DATABASE_URL") }))
+ * )
+ * ```
+ *
+ * @see {@link layerMemory} for running chats in one process
+ *
+ * @category layers
+ * @since 0.1.0
+ */
+export const layerCluster: Layer.Layer<AssistantEngine, never, ChatLog | Sharding.Sharding> = Layer.effect(
+  AssistantEngine,
+  Effect.gen(function* () {
+    const context = yield* Effect.context<ChatLog | Sharding.Sharding>()
+    const sharding = yield* Sharding.Sharding
+    const clients = new Map<string, EntityClient>()
+
+    const register = Effect.fnUntraced(function* (registration: Turns.Registration) {
+      const ChatEntity = chatEntity(registration.definition.name)
+
+      const handlers = Effect.gen(function* () {
+        const address = yield* Entity.CurrentAddress
+        const chatId = yield* Schema.decodeUnknownEffect(ChatId)(address.entityId).pipe(Effect.orDie)
+        // The entity's own: `Entity.keepAlive` needs its address and keep-alive latch, which is in there too.
+        const entityContext = yield* Effect.context<Entity.CurrentAddress | Sharding.Sharding>()
+
+        const keepAlive = (enabled: boolean) =>
+          Entity.keepAlive(enabled).pipe(
+            Effect.provide(entityContext),
+            Effect.catchCause((cause) => Effect.logWarning("assistant chat keep-alive failed", cause))
+          )
+
+        // A keep-alive left on from before a restart (by a run cut off) has done its job: it started the entity.
+        yield* keepAlive(false)
+
+        const chat = yield* Chat.make({ registration, chatId, keepAlive })
+        yield* chat.resume
+
+        return ChatEntity.of({
+          // The turns run inside the request, so a crash leaves it unprocessed and the cluster redelivers it.
+          Submit: (request) => chat.submit(request.payload.input),
+          Subscribe: (request) => chat.stream(request.payload.after),
+          Delete: () => chat.delete
+        })
+      })
+
+      yield* sharding.registerEntity(ChatEntity, handlers.pipe(Effect.provide(context)), {
+        maxIdleTime: "5 minutes",
+        concurrency: "unbounded"
+      })
+    })
+
+    // Made on first use, so a process that registers no assistant can still reach its chats.
+    const clientOf = Effect.fnUntraced(function* (assistant: string) {
+      const known = clients.get(assistant)
+
+      if (known !== undefined) return known
+
+      const makeClient = yield* chatEntity(assistant).client.pipe(Effect.provide(context))
+
+      const client: EntityClient = {
+        submit: (chatId, input) => makeClient(chatId).Submit({ input }, { discard: true }).pipe(Effect.orDie),
+        subscribe: (chatId, after) => makeClient(chatId).Subscribe({ after }).pipe(Stream.orDie),
+        delete: (chatId) => makeClient(chatId).Delete().pipe(Effect.orDie)
+      }
+
+      clients.set(assistant, client)
+
+      return client
+    })
+
+    return AssistantEngine.of({
+      register: (registration) => register(registration).pipe(Effect.provide(context)),
+      submit: Effect.fn("AssistantEngine.submit")(function* (
+        assistant: string,
+        chatId: ChatId,
+        input: AssistantEvent.Input | AssistantEvent.StoredAppEvent
+      ) {
+        const client = yield* clientOf(assistant)
+
+        yield* client.submit(chatId, input)
+      }),
+      subscribe: (assistant, chatId, after) =>
+        Stream.unwrap(
+          Effect.gen(function* () {
+            const client = yield* clientOf(assistant)
+
+            return client.subscribe(chatId, after)
+          })
+        ),
+      delete: Effect.fn("AssistantEngine.delete")(function* (assistant: string, chatId: ChatId) {
+        const client = yield* clientOf(assistant)
+
+        yield* client.delete(chatId)
+      })
+    })
+  })
+)