@dxos/assistant 0.8.4-main.c4373fc → 0.8.4-main.c85a9c8dae

package/src/session/session.ts

@@ -8,25 +8,24 @@ import type * as Tool from '@effect/ai/Tool';
 import type * as Toolkit from '@effect/ai/Toolkit';
 import * as Chunk from 'effect/Chunk';
 import * as Effect from 'effect/Effect';
-import type * as Schema from 'effect/Schema';
 import * as Stream from 'effect/Stream';
 
 import {
   AiParser,
   AiPreprocessor,
+  AiSummarizer,
   type AiToolNotFoundError,
   type PromptPreprocessingError,
   type ToolExecutionService,
   type ToolResolverService,
   callTool,
-  getToolCalls,
+  withoutToolCallParising,
 } from '@dxos/ai';
 import { type Blueprint } from '@dxos/blueprints';
-import { todo } from '@dxos/debug';
 import { Obj } from '@dxos/echo';
-import { TracingService } from '@dxos/functions';
+import { type FunctionInvocationService, TracingService } from '@dxos/functions';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { type ContentBlock, Message } from '@dxos/types';
 
 import { type AiAssistantError } from '../errors';
 
@@ -39,20 +38,33 @@ export type AiSessionRunRequirements =
   | LanguageModel.LanguageModel
   | ToolExecutionService
   | ToolResolverService
-  | TracingService;
+  | TracingService
+  | FunctionInvocationService;
 
-export type AiSessionRunParams<Tools extends Record<string, Tool.Any>> = {
+export type AiSessionOptions = {
+  /**
+   * Summarize before executing the prompt if the existing history exceeds this threshold.
+   */
+  summarizationThreshold?: number;
+};
+
+export type AiSessionRunProps<Tools extends Record<string, Tool.Any>> = {
   prompt: string;
+  // TODO(wittjosiah): Rename to systemPrompt.
   system?: string;
-  history?: DataType.Message[];
-  objects?: Obj.Any[];
-  blueprints?: Blueprint.Blueprint[];
+  history?: Message.Message[];
+  objects?: Obj.Unknown[];
+  blueprints?: readonly Blueprint.Blueprint[];
   toolkit?: Toolkit.WithHandler<Tools>;
-  observer?: GenerationObserver;
+  observer?: GenerationObserver<Tools>;
+  /**
+   * Callback for when a message is received from the user, model, or tool.
+   * This is useful for streaming the output to a queue.
+   */
+  // TODO(dmaretskyi): This is duplicated by generation observer. Better solution would be to convert the return type from effect to stream.
+  onOutput?: (message: Message.Message) => Effect.Effect<void, never, never>;
 };
 
-export type AiSessionOptions = {};
-
 /**
  * Contains message history, tools, current context.
  * Current context means the state of the app, time of day, and other contextual information.
@@ -69,13 +81,25 @@ export class AiSession {
 
   /** Message history from queue. */
   // TODO(burdon): Evolve into supporting a git-like graph of messages.
-  private _history: DataType.Message[] = [];
+  private _history: Message.Message[] = [];
 
   /** Pending messages for this session (incl. the current prompt). */
-  private _pending: DataType.Message[] = [];
+  private _pending: Message.Message[] = [];
+
+  private _started = 0;
+  private _ended = 0;
+  private _toolCalls = 0;
 
   constructor(private readonly _options: AiSessionOptions = {}) {}
 
+  get duration(): number {
+    return this._ended - this._started;
+  }
+
+  get toolCalls(): number {
+    return this._toolCalls;
+  }
+
   run = <Tools extends Record<string, Tool.Any>>({
     prompt,
     system: systemTemplate,
@@ -84,25 +108,41 @@ export class AiSession {
     blueprints = [],
     toolkit,
     observer = GenerationObserver.noop(),
-  }: AiSessionRunParams<Tools>): Effect.Effect<DataType.Message[], AiSessionRunError, AiSessionRunRequirements> =>
+    onOutput = () => Effect.void,
+  }: AiSessionRunProps<Tools>): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements> =>
     Effect.gen(this, function* () {
+      this._started = Date.now();
       this._history = [...history];
       this._pending = [];
-      const pending = this._pending;
 
-      const submitMessage = Effect.fnUntraced(function* (message: DataType.Message) {
-        pending.push(message);
-        yield* observer.onMessage(message);
-        yield* TracingService.emitConverationMessage(message);
-        return message;
-      });
+      const submitMessage = (message: Message.Message) =>
+        Effect.gen(this, function* () {
+          this._pending.push(message);
+          yield* observer.onMessage(message);
+          yield* TracingService.emitConverationMessage(message);
+          yield* onOutput(message);
+          return message;
+        });
+
+      const system = yield* formatSystemPrompt({ system: systemTemplate, blueprints, objects }).pipe(Effect.orDie);
+
+      if (this._options.summarizationThreshold !== undefined) {
+        const tokenCount = yield* AiPreprocessor.estimateTokens(
+          yield* AiPreprocessor.preprocessPrompt([...this._history], {
+            system,
+          }),
+        );
+        if (tokenCount > this._options.summarizationThreshold) {
+          const summary = yield* AiSummarizer.summarize([...this._history]);
+          yield* submitMessage(summary);
+        }
+      }
 
       // Submit the prompt.
       // TODO(burdon): Remove if cancelled?
       const promptMessage = yield* submitMessage(yield* formatUserPrompt({ prompt, history }));
 
       // Generate system and prompt messages.
-      const system = yield* formatSystemPrompt({ system: systemTemplate, blueprints, objects });
 
       // Tool call loop.
       do {
@@ -114,15 +154,19 @@ export class AiSession {
           objects: objects?.length ?? 0,
         });
 
-        const prompt = yield* AiPreprocessor.preprocessPrompt([...this._history, ...this._pending], { system });
+        const prompt = yield* AiPreprocessor.preprocessPrompt([...this._history, ...this._pending], {
+          system,
+          cacheControl: 'ephemeral',
+        });
 
         // Execute the stream request.
-        // logDump('prompt', Prompt.Prompt.pipe(Schema.encodeSync)(prompt));
-        const blocks = yield* LanguageModel.streamText({
+        const messages = yield* LanguageModel.streamText({
           prompt,
           toolkit,
           disableToolCallResolution: true,
         }).pipe(
+          // Disable tool call parsing so that we can handle parameter schema errors on a per-tool call basis.
+          withoutToolCallParising,
           // TOOD(dmaretskyi): Error mapping.
           AiParser.parseResponse({
             onBegin: () => observer.onBegin(),
@@ -130,21 +174,29 @@ export class AiSession {
             onPart: (part) => observer.onPart(part),
             onEnd: (summary) => observer.onEnd(summary),
           }),
+          Stream.mapEffect((block) =>
+            Effect.gen(function* () {
+              return yield* submitMessage(
+                Obj.make(Message.Message, {
+                  created: new Date().toISOString(),
+                  sender: { role: 'assistant' },
+                  blocks: [block],
+                }),
+              );
+            }),
+          ),
           Stream.runCollect,
           Effect.map(Chunk.toArray),
         );
 
-        // Create the response message.
-        const response = yield* submitMessage(
-          Obj.make(DataType.Message, {
-            created: new Date().toISOString(),
-            sender: { role: 'assistant' },
-            blocks,
-          }),
-        );
-
         // Parse the response for tool calls.
-        const toolCalls = getToolCalls(response);
+        const toolCalls = messages.flatMap((message) =>
+          message.blocks
+            .filter(
+              (block): block is ContentBlock.ToolCall => block._tag === 'toolCall' && block.providerExecuted === false,
+            )
+            .map((block) => ({ block, message })),
+        );
         if (toolCalls.length === 0) {
           break;
         } else if (!toolkit) {
@@ -153,13 +205,13 @@ export class AiSession {
 
         // Execute the tool calls.
         // TODO(burdon): Retry errors? Write result when each completes individually?
-        const toolResults = yield* Effect.forEach(toolCalls, (toolCall) => {
-          return callTool(toolkit, toolCall).pipe(
+        const toolResults = yield* Effect.forEach(toolCalls, ({ block, message }) => {
+          return callTool(toolkit, block).pipe(
             Effect.provide(
               TracingService.layerSubframe((context) => ({
                 ...context,
-                parentMessage: response.id,
-                toolCallId: toolCall.toolCallId,
+                parentMessage: message.id,
+                toolCallId: block.toolCallId,
               })),
             ),
           );
@@ -168,35 +220,22 @@ export class AiSession {
         // Add to queue and continue loop.
         // TODO(wittjosiah): Sometimes tool error results are added to the queue before the tool agent statuses.
         //   This results in a broken execution graph.
+        // TODO(dmaretskyi): Stream tool results one by one.
         yield* submitMessage(
-          Obj.make(DataType.Message, {
+          Obj.make(Message.Message, {
             created: new Date().toISOString(),
             sender: { role: 'tool' },
             blocks: toolResults,
           }),
         );
+
+        this._toolCalls++;
       } while (true);
 
-      log('done', { pending: this._pending.length });
+      this._ended = Date.now();
+      log('done', { pending: this._pending.length, duration: this.duration, tools: this._toolCalls });
       return this._pending;
     }).pipe(this._semaphore.withPermits(1), Effect.withSpan('AiSession.run'));
-
-  /**
-   * @deprecated
-   */
-  // TODO(burdon): Implement or remove.
-  async runStructured<S extends Schema.Schema.AnyNoContext>(
-    _schema: S,
-    _options: AiSessionRunParams<any>,
-  ): Promise<Schema.Schema.Type<S>> {
-    return todo();
-    // const parser = structuredOutputParser(schema);
-    // const result = await this.run({
-    //   ...options,
-    //   executableTools: [...(options.executableTools ?? []), parser.tool],
-    // });
-    // return parser.getResult(result);
-  }
 }
 
 const createSnippet = (text: string, len = 32) =>

package/src/session/toolkit.ts

@@ -4,45 +4,41 @@
 
 import type * as Tool from '@effect/ai/Tool';
 import * as Toolkit from '@effect/ai/Toolkit';
-import type * as Context from 'effect/Context';
 import * as Effect from 'effect/Effect';
 
-import { type AiToolNotFoundError, ToolExecutionService, type ToolId, ToolResolverService } from '@dxos/ai';
+import { type AiToolNotFoundError, GenericToolkit, ToolExecutionService, ToolResolverService } from '@dxos/ai';
 import { type Blueprint } from '@dxos/blueprints';
 import { isTruthy } from '@dxos/util';
 
-export type ToolkitParams<Tools extends Record<string, Tool.Any>> = {
-  toolkit?: Toolkit.Toolkit<Tools>;
-  toolIds?: ToolId[];
-  blueprints?: Blueprint.Blueprint[];
+export type CreateToolkitProps = {
+  toolkit?: Toolkit.Any;
+  blueprints?: readonly Blueprint.Blueprint[];
+  /**
+   * Self-contained with handlers toolkits.
+   */
+  genericToolkits?: readonly GenericToolkit.GenericToolkit[];
 };
 
 /**
  * Build a combined toolkit from the blueprint tools and the provided toolkit.
  */
-export const createToolkit = <Tools extends Record<string, Tool.Any> = {}>({
-  toolkit: toolkitParam,
-  toolIds = [],
+export const createToolkit = ({
+  toolkit: toolkitProp,
   blueprints = [],
-}: ToolkitParams<Tools>): Effect.Effect<
+  genericToolkits = [],
+}: CreateToolkitProps): Effect.Effect<
   Toolkit.WithHandler<any>,
   AiToolNotFoundError,
-  ToolResolverService | ToolExecutionService | Tool.HandlersFor<Tools>
+  ToolResolverService | ToolExecutionService | Tool.HandlersFor<any>
 > =>
   Effect.gen(function* () {
-    const blueprintToolkit = yield* ToolResolverService.resolveToolkit([
-      ...blueprints.flatMap(({ tools }) => tools),
-      ...toolIds,
-    ]);
+    const blueprintToolkit = yield* ToolResolverService.resolveToolkit(blueprints.flatMap(({ tools }) => tools));
+    const blueprintToolHandler = yield* blueprintToolkit.toContext(ToolExecutionService.handlersFor(blueprintToolkit));
+    const genericToolkit = GenericToolkit.merge(...genericToolkits);
 
-    const blueprintToolkitHandler: Context.Context<Tool.HandlersFor<any>> = yield* blueprintToolkit.toContext(
-      ToolExecutionService.handlersFor(blueprintToolkit),
-    );
-
-    const toolkit = Toolkit.merge(...[toolkitParam, blueprintToolkit].filter(isTruthy));
-    return yield* toolkit.pipe(Effect.provide(blueprintToolkitHandler)) as any as Effect.Effect<
-      Toolkit.WithHandler<any>,
-      never,
-      Tool.HandlersFor<Tools>
-    >;
+    const toolkit = Toolkit.merge(...[toolkitProp, blueprintToolkit, genericToolkit.toolkit].filter(isTruthy));
+    return yield* toolkit.pipe(
+      Effect.provide(blueprintToolHandler),
+      Effect.provide(genericToolkit.layer),
+    ) as any as Effect.Effect<Toolkit.WithHandler<any>, never, Tool.HandlersFor<any>>;
   });

package/src/session/version-pin.ts

@@ -4,10 +4,10 @@
 
 import * as Schema from 'effect/Schema';
 
-import { type AnyEchoObject, ObjectId } from '@dxos/echo/internal';
-import { type ObjectVersion } from '@dxos/echo-db';
-import { getVersion } from '@dxos/echo-db';
-import { type ContentBlock } from '@dxos/schema';
+import { type Obj } from '@dxos/echo';
+import { type ObjectVersion, getVersion } from '@dxos/echo-db';
+import { ObjectId } from '@dxos/keys';
+import { type ContentBlock } from '@dxos/types';
 
 // TODO(dmaretskyi): Extract.
 const ObjectVersionSchema = Schema.Unknown as Schema.Schema<ObjectVersion>;
@@ -27,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/templates/{format.tpl → instructions/format.tpl}

@@ -45,16 +45,13 @@ You can reply with a block reference for objects if you have the DXN. For exampl
 
 Each suggestion must in the form of a proposed action enclosed in a <suggestion> tag.
 Example: <suggestion>Show on a map</suggestion>
+Do not use bullets when listing suggestions.
 
 ### Multiple choice questions
 
 You may ask multiple choice questions by wrapping consise options inside <option> tags inside a well formed <select> tag.
 Example: <select><option>Yes</option><option>No</option></select>
 
-## Toolkit
-
-If the user asks for a list of tools, then just emit a single self-closing <toolkit/> tag instead of listing the tools.
-
 ## Context
 
 The current date and time is {{DATETIME}}.

package/src/templates/system.ts

@@ -7,8 +7,8 @@ import defaultsDeep from 'lodash.defaultsdeep';
 import { type ModelCapabilities } from '@dxos/ai';
 import { Template } from '@dxos/blueprints';
 
-import FORMAT from './format.tpl?raw';
-import SYSTEM from './system.tpl?raw';
+import FORMAT from './instructions/format.tpl?raw';
+import SYSTEM from './instructions/system.tpl?raw';
 
 /**
  * System template variables.

package/src/testing/index.ts

@@ -0,0 +1,5 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+export * from './layer';

package/src/testing/layer.ts

@@ -0,0 +1,85 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import { Registry } from '@effect-atom/atom';
+import * as Layer from 'effect/Layer';
+import * as Match from 'effect/Match';
+
+import { AiService, type ModelName } from '@dxos/ai';
+import { GenericToolkit } from '@dxos/ai';
+import { TestAiService } from '@dxos/ai/testing';
+import { Feed, type Type } from '@dxos/echo';
+import { CredentialsService, type FunctionDefinition, type ServiceCredential, TracingService } from '@dxos/functions';
+import { TracingServiceExt, TriggerDispatcher, TriggerStateStore } from '@dxos/functions-runtime';
+import { FunctionInvocationServiceLayerTest, TestDatabaseLayer } from '@dxos/functions-runtime/testing';
+
+import { ToolExecutionServices } from '../functions';
+
+interface TestLayerOptions {
+  aiServicePreset?: 'direct' | 'edge-local' | 'edge-remote';
+  model?: ModelName;
+  functions?: FunctionDefinition.Any[];
+  toolkits?: GenericToolkit.GenericToolkit[];
+  types?: Type.AnyEntity[];
+  credentials?: ServiceCredential[];
+  /*
+   * Tracing configuration.
+   * @default 'noop'
+   */
+  tracing?: 'noop' | 'console' | 'pretty';
+
+  disableLlmMemoization?: boolean;
+}
+
+export const AssistantTestLayer = ({
+  aiServicePreset = 'direct',
+  model = '@anthropic/claude-opus-4-6',
+  functions = [],
+  toolkits = [],
+  types = [],
+  credentials = [],
+  tracing = 'noop',
+  disableLlmMemoization = false,
+}: TestLayerOptions = {}) => {
+  const toolkit = GenericToolkit.merge(...toolkits);
+  return Layer.mergeAll(AiService.model(model), ToolExecutionServices).pipe(
+    Layer.provideMerge(
+      FunctionInvocationServiceLayerTest({
+        functions,
+      }),
+    ),
+    Layer.provideMerge(
+      Layer.mergeAll(
+        TestAiService({ preset: aiServicePreset, disableMemoization: disableLlmMemoization }),
+        TestDatabaseLayer({
+          spaceKey: 'fixed',
+          types,
+        }),
+        CredentialsService.configuredLayer(credentials),
+        Feed.notAvailable,
+        Match.value(tracing).pipe(
+          Match.when('noop', () => TracingService.layerNoop),
+          Match.when('console', () => TracingServiceExt.layerLogInfo()),
+          Match.when('pretty', () =>
+            TracingServiceExt.layerConsolePrettyPrint({
+              toolkit: (toolkits.length > 0 ? GenericToolkit.merge(...toolkits) : GenericToolkit.empty).toolkit as any,
+            }),
+          ),
+          Match.exhaustive,
+        ),
+      ),
+    ),
+    Layer.provideMerge(GenericToolkit.providerLayer(toolkit)),
+  );
+};
+
+interface TestLayerWithTriggersOptions extends TestLayerOptions {}
+
+export const AssistantTestLayerWithTriggers = (options: TestLayerWithTriggersOptions) =>
+  Layer.mergeAll(
+    AssistantTestLayer(options),
+    TriggerDispatcher.layer({ timeControl: 'manual', startingTime: new Date('2025-09-05T15:01:00.000Z') }).pipe(
+      Layer.provide(Registry.layer),
+    ),
+  ).pipe(Layer.provideMerge(TriggerStateStore.layerMemory));

package/src/{types.ts → types/artifact.ts}

@@ -2,10 +2,13 @@
 // Copyright 2025 DXOS.org
 //
 
+import type * as Effect from 'effect/Effect';
 import * as Schema from 'effect/Schema';
 
-import { type ObjectId } from '@dxos/echo/internal';
-import { DXN, LOCAL_SPACE_TAG, type SpaceId } from '@dxos/keys';
+import { type Err, Obj, Ref, type Type } from '@dxos/echo';
+import { Database } from '@dxos/echo';
+import { EncodedReference } from '@dxos/echo-protocol';
+import { DXN, LOCAL_SPACE_TAG, type ObjectId, type SpaceId } from '@dxos/keys';
 import { trim } from '@dxos/util';
 
 /**
@@ -13,21 +16,33 @@ import { trim } from '@dxos/util';
  */
 export const createArtifactElement = (id: ObjectId) => `<artifact id=${id} />`;
 
+// TODO(burdon): Rename RefFromLLM? -- yes -dm
 /**
  * A model-friendly way to reference an object.
  * Supports vairous formats that will be normalized to a DXN.
+ *
+ * @deprecated Use `Ref.Ref(XXX)` instead.
  */
-// TODO(burdon): Rename ObjectReference?
 export const ArtifactId: Schema.Schema<string> & {
   toDXN: (reference: ArtifactId, owningSpaceId?: SpaceId) => DXN;
+  resolve: <S extends Type.AnyEntity>(
+    schema: S,
+    ref: ArtifactId,
+  ) => Effect.Effect<Schema.Schema.Type<S>, Err.ObjectNotFoundError, Database.Service>;
 } = class extends Schema.String.annotations({
+  // TODO(dmaretskyi): This section gets overriden.
   description: trim`
     The ID of the referenced object. Formats accepted:
-    - DXN (dxn:echo:@:XXXXX). DXNs can be prepended with an @ symbol for compatibility with in-text references.
+    - DXN (dxn:echo:@:01KG7R1ZXWFMWQ4DA1Q6TN1DG4). DXNs can be prepended with an @ symbol for compatibility with in-text references.
     - space ID, object ID tuple (spaceID:objectID)
-    - Only object ID that is assumed to be in the current space (XXXXX)
+    - Only object ID that is assumed to be in the current space (01KG7R1ZXWFMWQ4DA1Q6TN1DG4)
   `,
-  examples: ['dxn:echo:@:XXXXX', '@dxn:echo:@:XXXXX', 'spaceID:objectID', 'XXXXX'],
+  examples: [
+    'dxn:echo:@:01KG7R1ZXWFMWQ4DA1Q6TN1DG4',
+    '@dxn:echo:@:01KG7R1ZXWFMWQ4DA1Q6TN1DG4',
+    'BM3FSHFOMJCHCG5QW7JTVKGYABD2GAA7G:01KG7R1ZXWFMWQ4DA1Q6TN1DG4',
+    '01KG7R1ZXWFMWQ4DA1Q6TN1DG4',
+  ],
 }) {
   static toDXN(reference: ArtifactId, owningSpaceId?: SpaceId): DXN {
     // Allow @dxn: prefix for compatibility with in-text references.
@@ -47,6 +62,28 @@ export const ArtifactId: Schema.Schema<string> & {
       throw new Error(`Unable to parse object reference: ${reference}`);
     }
   }
+
+  /**
+   * Resolves an artifact ID to an object.
+   */
+  static resolve<S extends Type.AnyEntity>(
+    schema: S,
+    ref: ArtifactId,
+  ): Effect.Effect<Schema.Schema.Type<S>, Err.ObjectNotFoundError, Database.Service> {
+    const dxn = ArtifactId.toDXN(ref);
+    return Database.resolve(dxn, schema);
+  }
 };
 
 export type ArtifactId = Schema.Schema.Type<typeof ArtifactId>;
+
+/**
+ * Schema that decodes ECHO reference object from an LLM-friendly input.
+ */
+export const RefFromLLM = Schema.transform(ArtifactId, Ref.Ref(Obj.Unknown), {
+  decode: (fromA, fromI) => EncodedReference.fromDXN(ArtifactId.toDXN(fromA)),
+  encode: (toI, toA) => EncodedReference.toDXN(toI).toString(),
+  strict: false,
+}).annotations({
+  description: ArtifactId.ast.annotations.description as string,
+});

package/src/types/index.ts

@@ -0,0 +1,5 @@
+//
+// Copyright 2025 DXOS.org
+//
+
+export * from './artifact';