@dxos/assistant 0.8.4-main.fd6878d → 0.8.4-staging.60fe92afc8

package/src/request/AiRequest.ts

@@ -0,0 +1,387 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+// @import-as-namespace
+
+import type * as AiError from '@effect/ai/AiError';
+import * as LanguageModel from '@effect/ai/LanguageModel';
+import type * as Toolkit from '@effect/ai/Toolkit';
+import * as Array from 'effect/Array';
+import * as Chunk from 'effect/Chunk';
+import * as Effect from 'effect/Effect';
+import { pipe } from 'effect/Function';
+import * as Option from 'effect/Option';
+import * as Stream from 'effect/Stream';
+
+import {
+  AiParser,
+  AiPreprocessor,
+  AiSummarizer,
+  type AiToolNotFoundError,
+  callTool,
+  type OpaqueToolkit,
+  type PromptPreprocessingError,
+  type ToolExecutionService,
+  type ToolResolverService,
+  withoutToolCallParising,
+} from '@dxos/ai';
+import { type Blueprint, Trace, Operation } from '@dxos/compute';
+import { Database, Obj, Registry } from '@dxos/echo';
+import { log } from '@dxos/log';
+import { ContentBlock, Message } from '@dxos/types';
+
+import { getOperationFromTool } from '../functions/services';
+import { type AiAssistantError, CompleteBlock, PartialBlock } from '../util';
+import { formatSystemPrompt, formatUserPrompt } from './format';
+import { GenerationObserver } from './observer';
+
+export type RunError = AiError.AiError | PromptPreprocessingError | AiToolNotFoundError | AiAssistantError;
+
+export type RunRequirements =
+  | LanguageModel.LanguageModel
+  | ToolExecutionService
+  | ToolResolverService
+  | Database.Service
+  | Operation.Service
+  | Registry.Service
+  | Trace.TraceService;
+
+export type Options = {
+  /**
+   * Summarize before executing the prompt if the existing history exceeds this threshold.
+   */
+  summarizationThreshold?: number;
+
+  // TODO(dmaretskyi): Plan to phase out in favor of TracingService and the return type being a stream.
+  observer?: GenerationObserver;
+  /**
+   * Callback for when a message is received from the user, model, or tool.
+   * This is useful for streaming the output to a queue.
+   */
+  onOutput?: (message: Message.Message) => Effect.Effect<void, never, never>;
+};
+
+export type RunProps<R = never> = {
+  prompt: string | ContentBlock.Any[];
+  // TODO(wittjosiah): Rename to systemPrompt.
+  system?: string;
+  history?: Message.Message[];
+  objects?: Obj.Unknown[];
+  blueprints?: readonly Blueprint.Blueprint[];
+  toolkit?: OpaqueToolkit.OpaqueToolkit<R>;
+};
+
+export type BeginProps = {
+  prompt: string | ContentBlock.Any[];
+  system?: string;
+  history?: Message.Message[];
+  objects?: Obj.Unknown[];
+  blueprints?: readonly Blueprint.Blueprint[];
+};
+
+export type TurnProps<R = never> = {
+  system: string;
+  toolkit?: OpaqueToolkit.OpaqueToolkit<R>;
+};
+
+export type TurnResult = {
+  messages: Message.Message[];
+  done: boolean;
+};
+
+/**
+ * Contains message history, tools, current context.
+ * Current context means the state of the app, time of day, and other contextual information.
+ * It makes requests to the model, its a state machine.
+ * It keeps track of the current goal.
+ * It manages the context window.
+ * Tracks the success criteria of reaching the goal, exposing metrics (stretch).
+ * Could be run locally in the app or remotely.
+ * Could be personal or shared.
+ */
+export class Request {
+  /** Prevents concurrent execution of session. */
+  private readonly _semaphore = Effect.runSync(Effect.makeSemaphore(1));
+
+  private readonly _observer: GenerationObserver;
+  private readonly _onOutput: (message: Message.Message) => Effect.Effect<void, never, never>;
+
+  /** Message history from queue. */
+  // TODO(burdon): Evolve into supporting a git-like graph of messages.
+  private _history: Message.Message[] = [];
+
+  /** Pending messages for this session (incl. the current prompt). */
+  private _pending: Message.Message[] = [];
+
+  private _started = 0;
+  private _ended = 0;
+  private _toolCalls = 0;
+
+  constructor(private readonly _options: Options = {}) {
+    this._observer = _options.observer ?? GenerationObserver.noop();
+    this._onOutput = _options.onOutput ?? (() => Effect.void);
+  }
+
+  get duration(): number {
+    return this._ended - this._started;
+  }
+
+  get toolCalls(): number {
+    return this._toolCalls;
+  }
+
+  get pending(): readonly Message.Message[] {
+    return this._pending;
+  }
+
+  private _submitMessage = (message: Message.Message): Effect.Effect<Message.Message, never, Trace.TraceService> =>
+    Effect.gen(this, function* () {
+      this._pending.push(message);
+      yield* this._observer.onMessage(message);
+      for (const block of message.blocks) {
+        log('write complete block', {
+          messageId: message.id,
+          role: message.sender.role!,
+          block: JSON.stringify(block),
+        });
+        yield* Trace.write(CompleteBlock, {
+          messageId: message.id,
+          role: message.sender.role!,
+          block,
+        });
+      }
+      yield* this._onOutput(message);
+      return message;
+    });
+
+  getToolCalls = () =>
+    pipe(
+      [...this._history, ...this._pending],
+      Array.reverse,
+      Array.takeWhile((_) => _.sender.role === 'assistant'),
+      Array.flatMap((_) => _.blocks.filter(ContentBlock.is('toolCall')).map((block) => ({ block, message: _ }))),
+      Array.filter((_) => !_.block.providerExecuted),
+      Array.reverse,
+    );
+
+  /**
+   * Initialize a session: set up history, perform summarization if needed, and submit the user prompt.
+   * Must be called before `runTurn()`.
+   */
+  begin = ({
+    prompt,
+    system,
+    history = [],
+    blueprints = [],
+    objects = [],
+  }: BeginProps): Effect.Effect<void, RunError, RunRequirements> =>
+    Effect.gen(this, function* () {
+      this._started = Date.now();
+      this._history = [...history];
+      this._pending = [];
+
+      const systemPrompt = yield* formatSystemPrompt({ system, blueprints, objects }).pipe(Effect.orDie);
+
+      if (this._options.summarizationThreshold !== undefined) {
+        const tokenCount = yield* AiPreprocessor.estimateTokens(
+          yield* AiPreprocessor.preprocessPrompt([...this._history], {
+            system: systemPrompt,
+          }),
+        );
+        if (tokenCount > this._options.summarizationThreshold) {
+          const summary = yield* AiSummarizer.summarize([...this._history]);
+          yield* this._submitMessage(summary);
+        }
+      }
+
+      yield* this._submitMessage(yield* formatUserPrompt({ prompt, history }));
+    }).pipe(Effect.withSpan('AiRequest.begin'));
+
+  /**
+   * Execute a single turn: one LLM generation followed by tool execution.
+   * The toolkit and system prompt can be updated between turns to reflect context changes (e.g. dynamically enabled blueprints).
+   */
+  runAgentTurn = <const R = never>({
+    system,
+    toolkit: opaqueToolkit,
+  }: TurnProps<R>): Effect.Effect<TurnResult, RunError, RunRequirements | R> =>
+    Effect.gen(this, function* () {
+      log('request', {
+        system: { snippet: createSnippet(system), length: system.length },
+        pending: this._pending.length,
+        history: this._history.length,
+      });
+
+      const prompt = yield* AiPreprocessor.preprocessPrompt([...this._history, ...this._pending], {
+        system,
+        cacheControl: 'ephemeral',
+      });
+
+      const toolkit = opaqueToolkit ? yield* opaqueToolkit.handlers : undefined;
+
+      const observer = this._observer;
+      let currentMessageId: Obj.ID | null = null;
+
+      const messages = yield* LanguageModel.streamText({
+        prompt,
+        toolkit,
+        disableToolCallResolution: true,
+      }).pipe(
+        withoutToolCallParising,
+        AiParser.parseResponse({
+          emitPartial: true,
+          onBegin: () => observer.onBegin(),
+          onBlock: (block) => observer.onBlock(block),
+          onPart: (part) => observer.onPart(part as any),
+          onEnd: (summary) => observer.onEnd(summary),
+        }),
+        Stream.map((block) => enrichToolCallBlock(block, toolkit)),
+        Stream.mapEffect(
+          (block) =>
+            Effect.gen(this, function* () {
+              if (block.pending) {
+                currentMessageId ??= Obj.ID.random();
+                log('emit ephemeral message', { id: currentMessageId, type: block._tag });
+                yield* Trace.write(PartialBlock, {
+                  messageId: currentMessageId,
+                  role: 'assistant',
+                  block,
+                });
+                return Option.none();
+              } else {
+                currentMessageId ??= Obj.ID.random();
+                const id = currentMessageId;
+                currentMessageId = null;
+                log('emit complete message', { id, type: block._tag });
+                const message = Obj.make(Message.Message, {
+                  id,
+                  created: new Date().toISOString(),
+                  sender: { role: 'assistant' },
+                  blocks: [block],
+                });
+                return Option.some(yield* this._submitMessage(message));
+              }
+            }),
+          { concurrency: 1, unordered: false },
+        ),
+        Stream.filterMap((_) => _),
+        Stream.runCollect,
+        Effect.map(Chunk.toArray),
+      );
+      log('messages', { messages });
+
+      const toolCalls = this.getToolCalls();
+
+      if (toolCalls.length === 0) {
+        this._ended = Date.now();
+        return { messages, done: true };
+      } else if (!toolkit) {
+        throw new Error('No toolkit provided');
+      }
+
+      return { messages, done: false };
+    }).pipe(Effect.withSpan('AiRequest.runAgentTurn'));
+
+  runTools = <const R = never>({
+    toolkit: opaqueToolkit,
+  }: {
+    toolkit?: OpaqueToolkit.OpaqueToolkit<R>;
+  }): Effect.Effect<void, RunError, RunRequirements | R> =>
+    Effect.gen(this, function* () {
+      const toolkit = opaqueToolkit ? yield* opaqueToolkit.handlers : undefined;
+      const toolCalls = this.getToolCalls();
+      const toolResults = yield* Effect.forEach(toolCalls, ({ block, message }) => {
+        if (!toolkit) {
+          throw new Error('No toolkit provided');
+        }
+        return callTool(toolkit, block);
+      });
+
+      yield* this._submitMessage(
+        Obj.make(Message.Message, {
+          created: new Date().toISOString(),
+          sender: { role: 'tool' },
+          blocks: toolResults,
+        }),
+      );
+
+      this._toolCalls += toolResults.length;
+    }).pipe(Effect.withSpan('AiRequest.runTools'));
+
+  /**
+   * Run a full conversation turn loop. Equivalent to calling `begin()` then `runTurn()` in a loop.
+   */
+  run = <const R = never>({
+    prompt,
+    system: systemTemplate,
+    history = [],
+    objects = [],
+    blueprints = [],
+    toolkit,
+  }: RunProps<R>): Effect.Effect<Message.Message[], RunError, RunRequirements | R> =>
+    Effect.gen(this, function* () {
+      yield* this.begin({ prompt, system: systemTemplate, history, objects, blueprints });
+
+      const system = yield* formatSystemPrompt({ system: systemTemplate, blueprints, objects }).pipe(Effect.orDie);
+
+      do {
+        const { done } = yield* this.runAgentTurn({ system, toolkit });
+        if (done) {
+          break;
+        }
+        yield* this.runTools({ toolkit });
+      } while (true);
+
+      log('done', { pending: this._pending.length, duration: this.duration, tools: this._toolCalls });
+      return this._pending;
+    }).pipe(this._semaphore.withPermits(1), Effect.withSpan('AiRequest.run'));
+}
+
+/**
+ * Annotates `toolCall` blocks with metadata about the backing Operation, when one exists.
+ * Tool calls that resolve to a toolkit handler (no Operation) are left unchanged so callers can
+ * distinguish operation invocations from inline tool calls.
+ */
+const enrichToolCallBlock = (
+  block: ContentBlock.Any,
+  toolkit: Toolkit.WithHandler<any> | undefined,
+): ContentBlock.Any => {
+  if (block._tag !== 'toolCall' || !toolkit) {
+    return block;
+  }
+  const tool = toolkit.tools[block.name];
+  if (!tool) {
+    return block;
+  }
+  // Some tools (provider-defined, raw MCP) don't carry an Effect `Context` for annotations and
+  // `getOperationFromTool` throws. Be defensive: treat any failure as "no operation".
+  let operationOpt: Option.Option<Operation.Definition.Any>;
+  try {
+    operationOpt = getOperationFromTool(tool);
+  } catch {
+    return block;
+  }
+  if (Option.isNone(operationOpt)) {
+    return block;
+  }
+  const { meta } = operationOpt.value;
+  return {
+    ...block,
+    operationKey: meta.key,
+    operationName: meta.name,
+    operationIcon: meta.icon,
+  } satisfies ContentBlock.ToolCall;
+};
+
+const createSnippet = (text: string, len = 32) =>
+  text.length <= len * 2 ? text : [text.slice(0, len), '...', text.slice(-len)].join('');
+
+// TODO(dmaretskyi): Extract as a general util.
+const logDump = (message: string, data: unknown) => {
+  // eslint-disable-next-line @typescript-eslint/no-require-imports
+  const { writeFileSync } = require('node:fs');
+  const path = `/tmp/log-data-${Date.now()}.json`;
+  writeFileSync(path, JSON.stringify(data, null, 2));
+  log.info(message, { path });
+};

package/src/{session → request}/artifact-diff.ts

@@ -2,10 +2,10 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Context } from 'effect';
+import * as Context from 'effect/Context';
 
-import { type ObjectVersion } from '@dxos/echo-db';
-import { type ObjectId } from '@dxos/echo-schema';
+import { type ObjectVersion } from '@dxos/echo-client';
+import { type EntityId } from '@dxos/keys';
 
 /**
  * Resolves artifact ids to their versions.
@@ -15,7 +15,6 @@ import { type ObjectId } from '@dxos/echo-schema';
  *
  * Can be optionally provided to the session run call.
  */
-// TODO(burdon): Move.
 // TODO(dmaretskyi): Convert to Context.Reference
 export class ArtifactDiffResolver extends Context.Tag('@dxos/assistant/ArtifactDiffResolver')<
   ArtifactDiffResolver,
@@ -24,9 +23,14 @@ export class ArtifactDiffResolver extends Context.Tag('@dxos/assistant/ArtifactD
 
 export namespace ArtifactDiffResolver {
   export type Service = {
-    resolve: (artifacts: { id: ObjectId; lastVersion: ObjectVersion }[]) => Promise<
+    resolve: (
+      artifacts: {
+        id: EntityId;
+        lastVersion: ObjectVersion;
+      }[],
+    ) => Promise<
       Map<
-        ObjectId,
+        EntityId,
         {
           version: ObjectVersion;
           diff?: string;

package/src/{session → request}/format.ts

@@ -2,21 +2,22 @@
 // Copyright 2025 DXOS.org
 //
 
-import { Array, Effect, Option, pipe } from 'effect';
-
-import { Template } from '@dxos/blueprints';
-import { Obj } from '@dxos/echo';
-import { ObjectVersion } from '@dxos/echo-db';
-import { type ObjectId } from '@dxos/echo-schema';
-import { DatabaseService } from '@dxos/functions';
+import * as Effect from 'effect/Effect';
+import * as Function from 'effect/Function';
+import * as Option from 'effect/Option';
+
+import { type FunctionNotFoundError, type Operation, Template } from '@dxos/compute';
+import { type Database, Obj, type Registry } from '@dxos/echo';
+import { ObjectVersion } from '@dxos/echo-client';
+import { type EntityNotFoundError } from '@dxos/echo/Err';
+import { type EntityId } from '@dxos/keys';
 import { log } from '@dxos/log';
-import { type ContentBlock, DataType } from '@dxos/schema';
+import { type ContentBlock, Message } from '@dxos/types';
 import { trim } from '@dxos/util';
 
-import { AiAssistantError } from '../errors';
-
+import { AiAssistantError } from '../util';
+import type * as AiRequest from './AiRequest';
 import { ArtifactDiffResolver } from './artifact-diff';
-import { type SessionRunParams } from './session';
 
 /**
  * Formats the system prompt.
@@ -26,20 +27,24 @@ export const formatSystemPrompt = ({
   system,
   blueprints = [],
   objects = [],
-}: Pick<SessionRunParams<any>, 'system' | 'blueprints' | 'objects'>) =>
+}: Pick<AiRequest.RunProps, 'system' | 'blueprints' | 'objects'>): Effect.Effect<
+  string,
+  FunctionNotFoundError | EntityNotFoundError,
+  Database.Service | Registry.Service | Operation.Service
+> =>
   Effect.gen(function* () {
-    // TOOD(burdon): Should process templates.
-    const blueprintDefs = yield* pipe(
+    const blueprintDefs = yield* Function.pipe(
       blueprints,
       Effect.forEach((blueprint) => Effect.succeed(blueprint.instructions)),
-      Effect.flatMap(Effect.forEach((template) => DatabaseService.load(template.source))),
-      Effect.map(
-        Array.map(
-          (template) => trim`
+      Effect.flatMap(
+        Effect.forEach((template) =>
+          Effect.gen(function* () {
+            return trim`
             <blueprint>
-              ${Template.process(template.content)}
+              ${yield* Template.processTemplate(template)}
             </blueprint>
-          `,
+          `;
+          }),
         ),
       ),
       Effect.map((blueprints) =>
@@ -47,12 +52,12 @@ export const formatSystemPrompt = ({
       ),
     );
 
-    const objectDefs = yield* pipe(
+    const objectDefs = yield* Function.pipe(
       objects,
       Effect.forEach((object) =>
         Effect.succeed(trim`
           <object>
-            <dxn>${Obj.getDXN(object)}</dxn>
+            <dxn>${Obj.getURI(object)}</dxn>
             <typename>${Obj.getTypename(object)}</typename>
           </object>
         `),
@@ -60,20 +65,23 @@ export const formatSystemPrompt = ({
       Effect.map((objects) => (objects.length > 0 ? ['## Context Objects', ...objects].join('\n\n') : undefined)),
     );
 
-    return yield* pipe(
+    return yield* Function.pipe(
       Effect.succeed([system, blueprintDefs, objectDefs].filter((def): def is string => def !== undefined)),
       Effect.map((parts) => parts.join('\n\n')),
     );
-  });
+  }).pipe(Effect.withSpan('formatSystemPrompt'));
 
 /**
  * Formats the user prompt.
  */
 // TODO(burdon): Move to AiPreprocessor.
 // TODO(burdon): Convert util below to `Effect.fn` (to preserve stack info)
-export const formatUserPrompt = ({ prompt, history = [] }: Pick<SessionRunParams<any>, 'prompt' | 'history'>) =>
+export const formatUserPrompt = ({
+  prompt,
+  history = [],
+}: Pick<AiRequest.RunProps, 'prompt' | 'history'>): Effect.Effect<Message.Message, AiRequest.RunError> =>
   Effect.gen(function* () {
-    const prelude: ContentBlock.Any[] = [];
+    const blocks: ContentBlock.Any[] = [];
 
     // TODO(dmaretskyi): Evaluate other approaches as `serviceOption` isn't represented in the type system.
     const artifactDiffResolver = yield* Effect.serviceOption(ArtifactDiffResolver);
@@ -84,33 +92,33 @@ export const formatUserPrompt = ({ prompt, history = [] }: Pick<SessionRunParams
           artifactDiffResolver.value.resolve(
             [...versions.entries()].map(([id, version]) => ({ id, lastVersion: version })),
           ),
-        catch: AiAssistantError.wrap('Artifact diff resolution error'),
+        catch: AiAssistantError.wrap({ message: 'Artifact diff resolution error' }),
       });
 
-      log.info('version', { artifactDiff, versions });
+      log('version', { artifactDiff, versions });
       for (const [id, { version }] of [...artifactDiff.entries()]) {
         if (ObjectVersion.equals(version, versions.get(id)!)) {
           artifactDiff.delete(id);
           continue;
         }
 
-        prelude.push({ _tag: 'anchor', objectId: id, version });
+        blocks.push({ _tag: 'anchor', objectId: id, version });
       }
 
       if (artifactDiff.size > 0) {
-        prelude.push(createArtifactUpdateBlock(artifactDiff));
+        blocks.push(createArtifactUpdateBlock(artifactDiff));
       }
     }
 
-    return Obj.make(DataType.Message, {
+    return Obj.make(Message.Message, {
       created: new Date().toISOString(),
       sender: { role: 'user' },
-      blocks: [...prelude, { _tag: 'text', text: prompt }],
+      blocks: typeof prompt === 'string' ? [...blocks, { _tag: 'text', text: prompt }] : [...blocks, ...prompt],
     });
-  });
+  }).pipe(Effect.withSpan('formatUserPrompt'));
 
-const gatherObjectVersions = (messages: DataType.Message[]): Map<ObjectId, ObjectVersion> => {
-  const artifactIds = new Map<ObjectId, ObjectVersion>();
+const gatherObjectVersions = (messages: Message.Message[]): Map<EntityId, ObjectVersion> => {
+  const artifactIds = new Map<EntityId, ObjectVersion>();
   for (const message of messages) {
     for (const block of message.blocks) {
       if (block._tag === 'anchor') {
@@ -123,12 +131,12 @@ const gatherObjectVersions = (messages: DataType.Message[]): Map<ObjectId, Objec
 };
 
 const createArtifactUpdateBlock = (
-  artifactDiff: Map<ObjectId, { version: ObjectVersion; diff?: string }>,
+  artifactDiff: Map<EntityId, { version: ObjectVersion; diff?: string }>,
 ): ContentBlock.Any => {
   return {
     _tag: 'text',
     // TODO(dmaretskyi): Does this need to be a special content-block?
-    disposition: 'artifact-update',
+    disposition: 'synthetic',
     text: trim`
       The following artifacts have been updated since the last message:
       ${[...artifactDiff.entries()]

package/src/request/index.ts

@@ -0,0 +1,8 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './format';
+export * from './observer';
+
+export * as AiRequest from './AiRequest';

package/src/request/observer.ts

@@ -0,0 +1,62 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import type * as Tool from '@effect/ai/Tool';
+import * as Effect from 'effect/Effect';
+import * as Function from 'effect/Function';
+
+import { type AiParser, type ConsolePrinter } from '@dxos/ai';
+import { type Message } from '@dxos/types';
+
+/**
+ * Live observer of the generation process.
+ */
+export interface GenerationObserver<
+  Tools extends Record<string, Tool.Any> = any,
+> extends AiParser.ParseResponseCallbacks<Tools> {
+  /**
+   * Complete messages fired during the session, both from the model and from the user.
+   * This message will contain all message blocks emitted through the `onBlock` callback.
+   */
+  onMessage: (message: Message.Message) => Effect.Effect<void>;
+}
+
+export const GenerationObserver = Object.freeze({
+  make: ({
+    onBegin = Function.constant(Effect.void),
+    onPart = Function.constant(Effect.void),
+    onBlock = Function.constant(Effect.void),
+    onEnd = Function.constant(Effect.void),
+    onMessage = Function.constant(Effect.void),
+  }: Partial<GenerationObserver> = {}): GenerationObserver => ({
+    onBegin,
+    onPart,
+    onBlock,
+    onEnd,
+    onMessage,
+  }),
+
+  noop: () => GenerationObserver.make(),
+
+  /**
+   * Debug printer to be used in unit-tests and browser devtools.
+   */
+  fromPrinter: (printer: ConsolePrinter) =>
+    GenerationObserver.make({
+      onBlock: (block) =>
+        Effect.sync(() => {
+          if (block.pending) {
+            return; // Only prints full blocks (better for unit-tests and browser devtools).
+          }
+          printer.printContentBlock(block);
+        }),
+      onMessage: (message) =>
+        Effect.sync(() => {
+          if (message.sender.role === 'assistant') {
+            return; // Skip assistant messages since they are printed in the `onBlock` callback.
+          }
+          printer.printMessage(message);
+        }),
+    }),
+});