@dxos/assistant 0.8.4-main.69d29f4 → 0.8.4-main.6fa680abb7

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 { type FunctionInvocationService, TracingService } from '@dxos/functions';
 import { log } from '@dxos/log';
-import { Message } from '@dxos/types';
+import { type ContentBlock, Message } from '@dxos/types';
 
 import { type AiAssistantError } from '../errors';
 
@@ -42,7 +41,12 @@ export type AiSessionRunRequirements =
   | TracingService
   | FunctionInvocationService;
 
-export type AiSessionOptions = {};
+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;
@@ -53,6 +57,12 @@ export type AiSessionRunProps<Tools extends Record<string, Tool.Any>> = {
   blueprints?: readonly Blueprint.Blueprint[];
   toolkit?: Toolkit.WithHandler<Tools>;
   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>;
 };
 
 /**
@@ -98,28 +108,41 @@ export class AiSession {
     blueprints = [],
     toolkit,
     observer = GenerationObserver.noop(),
+    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: Message.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 }).pipe(Effect.orDie);
-
-      // log('system', { prompt: system });
 
       // Tool call loop.
       do {
@@ -131,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(),
@@ -147,23 +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),
         );
 
-        // log('blocks', { blocks });
-
-        // Create the response message.
-        const response = yield* submitMessage(
-          Obj.make(Message.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) {
@@ -172,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,
               })),
             ),
           );
@@ -187,6 +220,7 @@ 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(Message.Message, {
             created: new Date().toISOString(),
@@ -202,23 +236,6 @@ export class AiSession {
       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: AiSessionRunProps<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

@@ -6,13 +6,17 @@ import type * as Tool from '@effect/ai/Tool';
 import * as Toolkit from '@effect/ai/Toolkit';
 import * as Effect from 'effect/Effect';
 
-import { type AiToolNotFoundError, ToolExecutionService, 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 CreateToolkitProps = {
   toolkit?: Toolkit.Any;
   blueprints?: readonly Blueprint.Blueprint[];
+  /**
+   * Self-contained with handlers toolkits.
+   */
+  genericToolkits?: readonly GenericToolkit.GenericToolkit[];
 };
 
 /**
@@ -21,6 +25,7 @@ export type CreateToolkitProps = {
 export const createToolkit = ({
   toolkit: toolkitProp,
   blueprints = [],
+  genericToolkits = [],
 }: CreateToolkitProps): Effect.Effect<
   Toolkit.WithHandler<any>,
   AiToolNotFoundError,
@@ -29,11 +34,11 @@ export const createToolkit = ({
   Effect.gen(function* () {
     const blueprintToolkit = yield* ToolResolverService.resolveToolkit(blueprints.flatMap(({ tools }) => tools));
     const blueprintToolHandler = yield* blueprintToolkit.toContext(ToolExecutionService.handlersFor(blueprintToolkit));
+    const genericToolkit = GenericToolkit.merge(...genericToolkits);
 
-    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<any>
-    >;
+    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/templates/instructions/format.tpl

@@ -52,10 +52,6 @@ Do not use bullets when listing suggestions.
 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/testing/layer.ts

@@ -2,50 +2,48 @@
 // 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 type { Type } from '@dxos/echo';
+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 { makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '../functions';
-import { GenericToolkit } from '../session';
+import { ToolExecutionServices } from '../functions';
 
 interface TestLayerOptions {
   aiServicePreset?: 'direct' | 'edge-local' | 'edge-remote';
   model?: ModelName;
   functions?: FunctionDefinition.Any[];
   toolkits?: GenericToolkit.GenericToolkit[];
-  types?: Type.Entity.Any[];
+  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-0',
+  model = '@anthropic/claude-opus-4-6',
   functions = [],
   toolkits = [],
   types = [],
   credentials = [],
   tracing = 'noop',
-}: TestLayerOptions) =>
-  Layer.mergeAll(
-    AiService.model(model),
-    makeToolResolverFromFunctions(functions, GenericToolkit.merge(...toolkits).toolkit as any),
-    makeToolExecutionServiceFromFunctions(
-      GenericToolkit.merge(...toolkits).toolkit as any,
-      GenericToolkit.merge(...toolkits).layer as any,
-    ),
-  ).pipe(
+  disableLlmMemoization = false,
+}: TestLayerOptions = {}) => {
+  const toolkit = GenericToolkit.merge(...toolkits);
+  return Layer.mergeAll(AiService.model(model), ToolExecutionServices).pipe(
     Layer.provideMerge(
       FunctionInvocationServiceLayerTest({
         functions,
@@ -53,26 +51,35 @@ export const AssistantTestLayer = ({
     ),
     Layer.provideMerge(
       Layer.mergeAll(
-        TestAiService({ preset: aiServicePreset }),
+        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()),
+          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') }),
+    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

@@ -5,7 +5,7 @@
 import type * as Effect from 'effect/Effect';
 import * as Schema from 'effect/Schema';
 
-import { type Err, Type } from '@dxos/echo';
+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';
@@ -16,16 +16,16 @@ 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 `RefFromLLM` instead.
+ * @deprecated Use `Ref.Ref(XXX)` instead.
  */
-// TODO(burdon): Rename RefFromLLM? -- yes -dm
 export const ArtifactId: Schema.Schema<string> & {
   toDXN: (reference: ArtifactId, owningSpaceId?: SpaceId) => DXN;
-  resolve: <S extends Type.Entity.Any>(
+  resolve: <S extends Type.AnyEntity>(
     schema: S,
     ref: ArtifactId,
   ) => Effect.Effect<Schema.Schema.Type<S>, Err.ObjectNotFoundError, Database.Service>;
@@ -33,11 +33,16 @@ export const ArtifactId: Schema.Schema<string> & {
   // 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.
@@ -61,12 +66,12 @@ export const ArtifactId: Schema.Schema<string> & {
   /**
    * Resolves an artifact ID to an object.
    */
-  static resolve<S extends Type.Entity.Any>(
+  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.Service.resolve(dxn, schema);
+    return Database.resolve(dxn, schema);
   }
 };
 
@@ -75,8 +80,10 @@ 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), {
+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,
 });