@dxos/assistant 0.8.4-main.ead640a → 0.8.4-main.ef1bc66f44

package/src/session/session.test.ts

@@ -9,13 +9,12 @@ import * as Effect from 'effect/Effect';
 import * as Layer from 'effect/Layer';
 import * as Schema from 'effect/Schema';
 
-import { AiService, MemoizedAiService, ToolExecutionService, ToolResolverService } from '@dxos/ai';
-import { AiServiceTestingPreset } from '@dxos/ai/testing';
 import { Obj, Type } from '@dxos/echo';
-import { TestHelpers } from '@dxos/effect';
-import { TracingService } from '@dxos/functions';
+import { TestHelpers } from '@dxos/effect/testing';
 import { log } from '@dxos/log';
 
+import { AssistantTestLayer } from '../testing';
+
 import { AiSession } from './session';
 
 // Define a calendar event artifact schema.
@@ -25,7 +24,7 @@ const CalendarEventSchema = Schema.Struct({
   endTime: Schema.String,
   description: Schema.String,
 }).pipe(
-  Type.Obj({
+  Type.object({
     typename: 'example.com/type/CalendarEvent',
     version: '0.1.0',
   }),
@@ -33,7 +32,7 @@ const CalendarEventSchema = Schema.Struct({
 
 type CalendarEvent = Schema.Schema.Type<typeof CalendarEventSchema>;
 
-class TestToolkit extends Toolkit.make(
+const TestToolkit = Toolkit.make(
   Tool.make('Calculator', {
     description: 'Basic calculator tool',
     parameters: {
@@ -46,7 +45,7 @@ class TestToolkit extends Toolkit.make(
     }),
     failure: Schema.Never,
   }),
-) {}
+);
 
 // Tool handlers.
 const toolkitLayer = TestToolkit.toLayer({
@@ -66,15 +65,10 @@ const toolkitLayer = TestToolkit.toLayer({
 });
 
 const TestLayer = Layer.mergeAll(
+  AssistantTestLayer({
+    types: [CalendarEventSchema],
+  }),
   toolkitLayer,
-  AiService.model('@anthropic/claude-3-5-sonnet-20241022'),
-  TracingService.layerNoop,
-  ToolResolverService.layerEmpty,
-  ToolExecutionService.layerEmpty,
-).pipe(
-  //
-  Layer.provideMerge(MemoizedAiService.layerTest()),
-  Layer.provide(AiServiceTestingPreset('direct')),
 );
 
 describe('AiSession', () => {
@@ -111,6 +105,25 @@ describe('AiSession', () => {
       TestHelpers.provideTestContext,
     ),
   );
+
+  it.effect(
+    'tool schema error',
+    Effect.fnUntraced(
+      function* (_) {
+        const session = new AiSession({ operationModel: 'configured' });
+        const toolkit = yield* TestToolkit;
+        const response = yield* session.run({
+          toolkit,
+          prompt:
+            'I am testing error handling in tool paramter parsing. I want you to call the calculator tool but omit the input parameter to intentionally differ from the tool schema.',
+          history: [],
+        });
+        log.info('response', { response });
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+  );
 });
 
 // Travel to rome, florence, livorno, siena, madrid for conferences

package/src/session/session.ts

@@ -20,13 +20,14 @@ import {
   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 { Message } from '@dxos/types';
 
 import { type AiAssistantError } from '../errors';
 
@@ -39,20 +40,22 @@ export type AiSessionRunRequirements =
   | LanguageModel.LanguageModel
   | ToolExecutionService
   | ToolResolverService
-  | TracingService;
+  | TracingService
+  | FunctionInvocationService;
 
-export type AiSessionRunParams<Tools extends Record<string, Tool.Any>> = {
+export type AiSessionOptions = {};
+
+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>;
 };
 
-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 +72,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,13 +99,14 @@ export class AiSession {
     blueprints = [],
     toolkit,
     observer = GenerationObserver.noop(),
-  }: AiSessionRunParams<Tools>): Effect.Effect<DataType.Message[], AiSessionRunError, AiSessionRunRequirements> =>
+  }: 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) {
+      const submitMessage = Effect.fnUntraced(function* (message: Message.Message) {
         pending.push(message);
         yield* observer.onMessage(message);
         yield* TracingService.emitConverationMessage(message);
@@ -102,7 +118,7 @@ export class AiSession {
       const promptMessage = yield* submitMessage(yield* formatUserPrompt({ prompt, history }));
 
       // Generate system and prompt messages.
-      const system = yield* formatSystemPrompt({ system: systemTemplate, blueprints, objects });
+      const system = yield* formatSystemPrompt({ system: systemTemplate, blueprints, objects }).pipe(Effect.orDie);
 
       // Tool call loop.
       do {
@@ -114,15 +130,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({
           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(),
@@ -136,7 +156,7 @@ export class AiSession {
 
         // Create the response message.
         const response = yield* submitMessage(
-          Obj.make(DataType.Message, {
+          Obj.make(Message.Message, {
             created: new Date().toISOString(),
             sender: { role: 'assistant' },
             blocks,
@@ -169,15 +189,18 @@ export class AiSession {
         // TODO(wittjosiah): Sometimes tool error results are added to the queue before the tool agent statuses.
         //   This results in a broken execution graph.
         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'));
 
@@ -187,7 +210,7 @@ export class AiSession {
   // TODO(burdon): Implement or remove.
   async runStructured<S extends Schema.Schema.AnyNoContext>(
     _schema: S,
-    _options: AiSessionRunParams<any>,
+    _options: AiSessionRunProps<any>,
   ): Promise<Schema.Schema.Type<S>> {
     return todo();
     // const parser = structuredOutputParser(schema);

package/src/session/toolkit.ts

@@ -4,45 +4,36 @@
 
 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, 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[];
 };
 
 /**
  * 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<
+}: 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 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<
+    const toolkit = Toolkit.merge(...[toolkitProp, blueprintToolkit].filter(isTruthy));
+    return yield* toolkit.pipe(Effect.provide(blueprintToolHandler)) as any as Effect.Effect<
       Toolkit.WithHandler<any>,
       never,
-      Tool.HandlersFor<Tools>
+      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,84 @@
+//
+// 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 { TestAiService } from '@dxos/ai/testing';
+import 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';
+import { GenericToolkit } from '../session';
+
+interface TestLayerOptions {
+  aiServicePreset?: 'direct' | 'edge-local' | 'edge-remote';
+  model?: ModelName;
+  functions?: FunctionDefinition.Any[];
+  toolkits?: GenericToolkit.GenericToolkit[];
+  types?: Type.Entity.Any[];
+  credentials?: ServiceCredential[];
+  /*
+   * Tracing configuration.
+   * @default 'noop'
+   */
+  tracing?: 'noop' | 'console' | 'pretty';
+
+  disableLlmMemoization?: boolean;
+}
+
+export const AssistantTestLayer = ({
+  aiServicePreset = 'direct',
+  model = '@anthropic/claude-opus-4-0',
+  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),
+        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/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, 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 `Type.Ref(XXX)` instead.
  */
-// TODO(burdon): Rename ObjectReference?
 export const ArtifactId: Schema.Schema<string> & {
   toDXN: (reference: ArtifactId, owningSpaceId?: SpaceId) => DXN;
+  resolve: <S extends Type.Entity.Any>(
+    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.Entity.Any>(
+    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, Type.Ref(Type.Obj), {
+  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,
+});