@dxos/assistant 0.8.4-main.fffef41 → 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

@@ -4,8 +4,8 @@
 
 import * as Context from 'effect/Context';
 
-import { type ObjectId } from '@dxos/echo/internal';
-import { type ObjectVersion } from '@dxos/echo-db';
+import { type ObjectVersion } from '@dxos/echo-client';
+import { type EntityId } from '@dxos/keys';
 
 /**
  * Resolves artifact ids to their versions.
@@ -25,12 +25,12 @@ export namespace ArtifactDiffResolver {
   export type Service = {
     resolve: (
       artifacts: {
-        id: ObjectId;
+        id: EntityId;
         lastVersion: ObjectVersion;
       }[],
     ) => Promise<
       Map<
-        ObjectId,
+        EntityId,
         {
           version: ObjectVersion;
           diff?: string;

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

@@ -2,23 +2,22 @@
 // Copyright 2025 DXOS.org
 //
 
-import * as Array from 'effect/Array';
 import * as Effect from 'effect/Effect';
 import * as Function from 'effect/Function';
 import * as Option from 'effect/Option';
 
-import { Template } from '@dxos/blueprints';
-import { Obj } from '@dxos/echo';
-import { type ObjectId } from '@dxos/echo/internal';
-import { DatabaseService, ObjectVersion } from '@dxos/echo-db';
+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, 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 AiSessionRunError, type AiSessionRunParams } from './session';
 
 /**
  * Formats the system prompt.
@@ -28,20 +27,24 @@ export const formatSystemPrompt = ({
   system,
   blueprints = [],
   objects = [],
-}: Pick<AiSessionRunParams<any>, 'system' | 'blueprints' | 'objects'>) =>
+}: Pick<AiRequest.RunProps, 'system' | 'blueprints' | 'objects'>): Effect.Effect<
+  string,
+  FunctionNotFoundError | EntityNotFoundError,
+  Database.Service | Registry.Service | Operation.Service
+> =>
   Effect.gen(function* () {
     const blueprintDefs = yield* Function.pipe(
       blueprints,
       Effect.forEach((blueprint) => Effect.succeed(blueprint.instructions)),
-      Effect.flatMap(Effect.forEach((template) => DatabaseService.loadOption(template.source))),
-      Effect.map(Array.filter(Option.isSome)),
-      Effect.map(
-        Array.map(
-          (template) => trim`
+      Effect.flatMap(
+        Effect.forEach((template) =>
+          Effect.gen(function* () {
+            return trim`
             <blueprint>
-              ${Template.process(template.value.content)}
+              ${yield* Template.processTemplate(template)}
             </blueprint>
-          `,
+          `;
+          }),
         ),
       ),
       Effect.map((blueprints) =>
@@ -54,7 +57,7 @@ export const formatSystemPrompt = ({
       Effect.forEach((object) =>
         Effect.succeed(trim`
           <object>
-            <dxn>${Obj.getDXN(object)}</dxn>
+            <dxn>${Obj.getURI(object)}</dxn>
             <typename>${Obj.getTypename(object)}</typename>
           </object>
         `),
@@ -66,7 +69,7 @@ export const formatSystemPrompt = ({
       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.
@@ -76,7 +79,7 @@ export const formatSystemPrompt = ({
 export const formatUserPrompt = ({
   prompt,
   history = [],
-}: Pick<AiSessionRunParams<any>, 'prompt' | 'history'>): Effect.Effect<Message.Message, AiSessionRunError> =>
+}: Pick<AiRequest.RunProps, 'prompt' | 'history'>): Effect.Effect<Message.Message, AiRequest.RunError> =>
   Effect.gen(function* () {
     const blocks: ContentBlock.Any[] = [];
 
@@ -110,12 +113,12 @@ export const formatUserPrompt = ({
     return Obj.make(Message.Message, {
       created: new Date().toISOString(),
       sender: { role: 'user' },
-      blocks: [...blocks, { _tag: 'text', text: prompt }],
+      blocks: typeof prompt === 'string' ? [...blocks, { _tag: 'text', text: prompt }] : [...blocks, ...prompt],
     });
-  });
+  }).pipe(Effect.withSpan('formatUserPrompt'));
 
-const gatherObjectVersions = (messages: Message.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') {
@@ -128,12 +131,12 @@ const gatherObjectVersions = (messages: Message.Message[]): Map<ObjectId, Object
 };
 
 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/{session → request}/observer.ts

@@ -2,6 +2,7 @@
 // Copyright 2025 DXOS.org
 //
 
+import type * as Tool from '@effect/ai/Tool';
 import * as Effect from 'effect/Effect';
 import * as Function from 'effect/Function';
 
@@ -11,7 +12,9 @@ import { type Message } from '@dxos/types';
 /**
  * Live observer of the generation process.
  */
-export interface GenerationObserver extends AiParser.ParseResponseCallbacks {
+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.

package/src/{session → request}/version-pin.ts

@@ -4,8 +4,9 @@
 
 import * as Schema from 'effect/Schema';
 
-import { type AnyEchoObject, ObjectId } from '@dxos/echo/internal';
-import { type ObjectVersion, getVersion } from '@dxos/echo-db';
+import { type Obj } from '@dxos/echo';
+import { type ObjectVersion, getVersion } from '@dxos/echo-client';
+import { EntityId } from '@dxos/keys';
 import { type ContentBlock } from '@dxos/types';
 
 // TODO(dmaretskyi): Extract.
@@ -13,7 +14,7 @@ const ObjectVersionSchema = Schema.Unknown as Schema.Schema<ObjectVersion>;
 
 const VersionPinSchema = Schema.Struct({
   // TODO(dmaretskyi): Use Ref when those support encoding.
-  objectId: ObjectId,
+  objectId: EntityId,
   // TODO(dmaretskyi): Could be opaque.
   version: ObjectVersionSchema,
 });
@@ -26,11 +27,11 @@ export interface VersionPin extends Schema.Schema.Type<typeof VersionPinSchema>
 
 export const VersionPin: typeof VersionPinSchema & {
   DISPOSITION: 'version-pin';
-  fromObject: (object: AnyEchoObject) => VersionPin;
+  fromObject: (object: Obj.Unknown) => VersionPin;
   createBlock: (pin: VersionPin) => ContentBlock.Any;
 } = class extends VersionPinSchema {
   static readonly DISPOSITION = 'version-pin';
-  static fromObject(object: AnyEchoObject): VersionPin {
+  static fromObject(object: Obj.Unknown): VersionPin {
     return VersionPin.make({
       objectId: object.id,
       version: getVersion(object),

package/src/session/AiContext.test.ts

@@ -0,0 +1,49 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+import * as Context from 'effect/Context';
+import * as Effect from 'effect/Effect';
+import * as Layer from 'effect/Layer';
+import * as Runtime from 'effect/Runtime';
+import * as Schema from 'effect/Schema';
+import { describe, it } from 'vitest';
+
+import { DXN, Feed, Obj, Ref, Type } from '@dxos/echo';
+
+import * as AiContext from './AiContext';
+
+const createMockFeedRuntime = (): Runtime.Runtime<Feed.FeedService> => {
+  const mockFeedService: Context.Tag.Service<Feed.FeedService> = {
+    append: async () => {},
+    remove: async () => {},
+    query: () =>
+      ({
+        subscribe: () => () => {},
+        results: [],
+        run: async () => [],
+      }) as any,
+    sync: async () => {},
+    getSyncState: async () => ({ blocksToPull: 0, blocksToPush: 0, totalBlocks: 0 }),
+  };
+  const layer = Layer.succeed(Feed.FeedService, mockFeedService);
+  return Effect.runSync(Effect.runtime<Feed.FeedService>().pipe(Effect.provide(layer)));
+};
+
+describe('AiContext.Binder', () => {
+  it('should handle bind with Ref', async () => {
+    const feed = Feed.make();
+    const runtime = createMockFeedRuntime();
+    const binder = new AiContext.Binder({ feed, runtime });
+
+    const TestSchema = Schema.Struct({}).pipe(Type.makeObject(DXN.make('org.dxos.type.example', '0.1.0')));
+
+    const obj = Obj.make(TestSchema, {});
+    const ref = Ref.make(obj);
+
+    await binder.bind({
+      blueprints: [],
+      objects: [ref],
+    });
+  });
+});