@dxos/assistant 0.8.4-main.bc674ce → 0.8.4-main.bcb3aa67d6

package/src/session/session.test.ts

@@ -4,7 +4,7 @@
 
 import * as Tool from '@effect/ai/Tool';
 import * as Toolkit from '@effect/ai/Toolkit';
-import { describe, it } from '@effect/vitest';
+import { describe, expect, it } from '@effect/vitest';
 import * as Effect from 'effect/Effect';
 import * as Layer from 'effect/Layer';
 import * as Schema from 'effect/Schema';
@@ -12,6 +12,7 @@ import * as Schema from 'effect/Schema';
 import { Obj, Type } from '@dxos/echo';
 import { TestHelpers } from '@dxos/effect/testing';
 import { log } from '@dxos/log';
+import { Message } from '@dxos/types';
 
 import { AssistantTestLayer } from '../testing';
 
@@ -25,7 +26,7 @@ const CalendarEventSchema = Schema.Struct({
   description: Schema.String,
 }).pipe(
   Type.object({
-    typename: 'example.com/type/CalendarEvent',
+    typename: 'com.example.type.calendar-event',
     version: '0.1.0',
   }),
 );
@@ -67,6 +68,7 @@ const toolkitLayer = TestToolkit.toLayer({
 const TestLayer = Layer.mergeAll(
   AssistantTestLayer({
     types: [CalendarEventSchema],
+    tracing: 'pretty',
   }),
   toolkitLayer,
 );
@@ -76,7 +78,7 @@ describe('AiSession', () => {
     'no tools',
     Effect.fnUntraced(
       function* (_) {
-        const session = new AiSession({ operationModel: 'configured' });
+        const session = new AiSession();
         const response = yield* session.run({
           prompt: 'Hello world!',
           history: [],
@@ -92,7 +94,7 @@ describe('AiSession', () => {
     'calculator',
     Effect.fnUntraced(
       function* (_) {
-        const session = new AiSession({ operationModel: 'configured' });
+        const session = new AiSession();
         const toolkit = yield* TestToolkit;
         const response = yield* session.run({
           toolkit,
@@ -105,6 +107,52 @@ describe('AiSession', () => {
       TestHelpers.provideTestContext,
     ),
   );
+
+  it.effect(
+    'tool schema error',
+    Effect.fnUntraced(
+      function* (_) {
+        const session = new AiSession();
+        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,
+    ),
+  );
+
+  it.effect(
+    'summarization',
+    Effect.fnUntraced(
+      function* (_) {
+        const session = new AiSession({ summarizationThreshold: 0 }); // Force summarization.
+        const response = yield* session.run({
+          prompt: 'What did we talk about?',
+          history: [
+            Obj.make(Message.Message, {
+              created: '2024-01-01T10:00:00Z',
+              sender: { role: 'user' },
+              blocks: [{ _tag: 'text', text: 'How many apples are in the basket?' }],
+            }),
+            Obj.make(Message.Message, {
+              created: '2024-01-01T11:00:00Z',
+              sender: { role: 'assistant' },
+              blocks: [{ _tag: 'text', text: 'There are 10 apples in the basket.' }],
+            }),
+          ],
+        });
+        expect(response).toBeDefined();
+      },
+      Effect.provide(TestLayer),
+      TestHelpers.provideTestContext,
+    ),
+  );
 });
 
 // Travel to rome, florence, livorno, siena, madrid for conferences

package/src/session/session.ts

@@ -8,30 +8,32 @@ 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 * as Array from 'effect/Array';
+import * as Option from 'effect/Option';
 import {
   AiParser,
   AiPreprocessor,
+  AiSummarizer,
   type AiToolNotFoundError,
+  callTool,
   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 { type FunctionInvocationService, Trace, TracingService } from '@dxos/functions';
 import { log } from '@dxos/log';
-import { Message } from '@dxos/types';
+import { ContentBlock, Message } from '@dxos/types';
 
 import { type AiAssistantError } from '../errors';
 
 import { formatSystemPrompt, formatUserPrompt } from './format';
 import { GenerationObserver } from './observer';
+import { pipe } from 'effect/Function';
+import { CompleteBlock, PartialBlock } from '../tracing';
 
 export type AiSessionRunError = AiError.AiError | PromptPreprocessingError | AiToolNotFoundError | AiAssistantError;
 
@@ -39,10 +41,28 @@ export type AiSessionRunRequirements =
   | LanguageModel.LanguageModel
   | ToolExecutionService
   | ToolResolverService
-  | TracingService
-  | FunctionInvocationService;
+  | FunctionInvocationService
+  | Trace.TraceService
+  /**
+   * @deprecated Retained for backward compatibility with tool handlers that use TracingService.emitStatus().
+   *   New code should use Trace.TraceService instead.
+   */
+  | TracingService;
 
-export type AiSessionOptions = {};
+export type AiSessionOptions = {
+  /**
+   * 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 AiSessionRunProps<Tools extends Record<string, Tool.Any>> = {
   prompt: string;
@@ -52,7 +72,24 @@ export type AiSessionRunProps<Tools extends Record<string, Tool.Any>> = {
   objects?: Obj.Unknown[];
   blueprints?: readonly Blueprint.Blueprint[];
   toolkit?: Toolkit.WithHandler<Tools>;
-  observer?: GenerationObserver<Tools>;
+};
+
+export type AiSessionBeginProps = {
+  prompt: string;
+  system?: string;
+  history?: Message.Message[];
+  objects?: Obj.Unknown[];
+  blueprints?: readonly Blueprint.Blueprint[];
+};
+
+export type AiSessionTurnProps<Tools extends Record<string, Tool.Any>> = {
+  system: string;
+  toolkit?: Toolkit.WithHandler<Tools>;
+};
+
+export type AiSessionTurnResult = {
+  messages: Message.Message[];
+  done: boolean;
 };
 
 /**
@@ -65,10 +102,14 @@ export type AiSessionRunProps<Tools extends Record<string, Tool.Any>> = {
  * Could be run locally in the app or remotely.
  * Could be personal or shared.
  */
+// TODO(dmaretskyi): Rename AiSessionRequest
 export class AiSession {
   /** 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[] = [];
@@ -80,7 +121,10 @@ export class AiSession {
   private _ended = 0;
   private _toolCalls = 0;
 
-  constructor(private readonly _options: AiSessionOptions = {}) {}
+  constructor(private readonly _options: AiSessionOptions = {}) {
+    this._observer = _options.observer ?? GenerationObserver.noop();
+    this._onOutput = _options.onOutput ?? (() => Effect.void);
+  }
 
   get duration(): number {
     return this._ended - this._started;
@@ -90,135 +134,205 @@ export class AiSession {
     return this._toolCalls;
   }
 
-  run = <Tools extends Record<string, Tool.Any>>({
+  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) {
+        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),
+    );
+
+  /**
+   * Initialize a session: set up history, perform summarization if needed, and submit the user prompt.
+   * Must be called before `runTurn()`.
+   */
+  begin = ({
     prompt,
-    system: systemTemplate,
+    system,
     history = [],
-    objects = [],
     blueprints = [],
-    toolkit,
-    observer = GenerationObserver.noop(),
-  }: AiSessionRunProps<Tools>): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements> =>
+    objects = [],
+  }: AiSessionBeginProps): Effect.Effect<void, 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 systemPrompt = yield* formatSystemPrompt({ system, blueprints, objects }).pipe(Effect.orDie);
 
-      // Submit the prompt.
-      // TODO(burdon): Remove if cancelled?
-      const promptMessage = yield* submitMessage(yield* formatUserPrompt({ prompt, history }));
+      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);
+        }
+      }
 
-      // Generate system and prompt messages.
-      const system = yield* formatSystemPrompt({ system: systemTemplate, blueprints, objects }).pipe(Effect.orDie);
+      yield* this._submitMessage(yield* formatUserPrompt({ prompt, history }));
+    });
 
-      // log('system', { prompt: system });
+  /**
+   * 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 = <Tools extends Record<string, Tool.Any>>({
+    system,
+    toolkit,
+  }: AiSessionTurnProps<Tools>): Effect.Effect<AiSessionTurnResult, AiSessionRunError, AiSessionRunRequirements> =>
+    Effect.gen(this, function* () {
+      log('request', {
+        system: { snippet: createSnippet(system), length: system.length },
+        pending: this._pending.length,
+        history: this._history.length,
+      });
 
-      // Tool call loop.
-      do {
-        log('request', {
-          prompt: promptMessage,
-          system: { snippet: createSnippet(system), length: system.length },
-          pending: this._pending.length,
-          history: this._history.length,
-          objects: objects?.length ?? 0,
-        });
+      const prompt = yield* AiPreprocessor.preprocessPrompt([...this._history, ...this._pending], {
+        system,
+        cacheControl: 'ephemeral',
+      });
 
-        const prompt = yield* AiPreprocessor.preprocessPrompt([...this._history, ...this._pending], { system });
-
-        // Execute the stream request.
-        // logDump('prompt', Prompt.Prompt.pipe(Schema.encodeSync)(prompt));
-        const blocks = yield* LanguageModel.streamText({
-          prompt,
-          toolkit,
-          disableToolCallResolution: true,
-        }).pipe(
-          // TOOD(dmaretskyi): Error mapping.
-          AiParser.parseResponse({
-            onBegin: () => observer.onBegin(),
-            onBlock: (block) => observer.onBlock(block),
-            onPart: (part) => observer.onPart(part),
-            onEnd: (summary) => observer.onEnd(summary),
-          }),
-          Stream.runCollect,
-          Effect.map(Chunk.toArray),
+      const observer = this._observer;
+      let currentMessageId: Obj.ID | null = null; // Consistent IDs for pending blocks preceding the complete one.
+      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.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),
+      );
+
+      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 };
+    });
+
+  runTools = <Tools extends Record<string, Tool.Any>>({ toolkit }: { toolkit?: Toolkit.WithHandler<Tools> }) =>
+    Effect.gen(this, function* () {
+      const toolCalls = this.getToolCalls();
+      // TODO(burdon): Retry errors? Write result when each completes individually?
+      const toolResults = yield* Effect.forEach(toolCalls, ({ block, message }) => {
+        if (!toolkit) {
+          throw new Error('No toolkit provided');
+        }
+        return callTool(toolkit, block).pipe(
+          Effect.provide(
+            TracingService.layerSubframe((context) => ({
+              ...context,
+              parentMessage: message.id,
+              toolCallId: block.toolCallId,
+            })),
+          ),
         );
+      });
 
-        // log('blocks', { blocks });
+      // TODO(wittjosiah): Sometimes tool error results are added to the queue before the tool agent statuses.
+      // TODO(dmaretskyi): Stream tool results one by one.
+      yield* this._submitMessage(
+        Obj.make(Message.Message, {
+          created: new Date().toISOString(),
+          sender: { role: 'tool' },
+          blocks: toolResults,
+        }),
+      );
+
+      this._toolCalls += toolResults.length;
+    });
+  /**
+   * Run a full conversation turn loop. Equivalent to calling `begin()` then `runTurn()` in a loop.
+   */
+  run = <Tools extends Record<string, Tool.Any>>({
+    prompt,
+    system: systemTemplate,
+    history = [],
+    objects = [],
+    blueprints = [],
+    toolkit,
+  }: AiSessionRunProps<Tools>): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements> =>
+    Effect.gen(this, function* () {
+      yield* this.begin({ prompt, system: systemTemplate, history, objects, blueprints });
 
-        // Create the response message.
-        const response = yield* submitMessage(
-          Obj.make(Message.Message, {
-            created: new Date().toISOString(),
-            sender: { role: 'assistant' },
-            blocks,
-          }),
-        );
+      const system = yield* formatSystemPrompt({ system: systemTemplate, blueprints, objects }).pipe(Effect.orDie);
 
-        // Parse the response for tool calls.
-        const toolCalls = getToolCalls(response);
-        if (toolCalls.length === 0) {
+      do {
+        const { done } = yield* this.runAgentTurn({ system, toolkit });
+        if (done) {
           break;
-        } else if (!toolkit) {
-          throw new Error('No toolkit provided'); // TODO(burdon): Throw user error?
         }
-
-        // 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(
-            Effect.provide(
-              TracingService.layerSubframe((context) => ({
-                ...context,
-                parentMessage: response.id,
-                toolCallId: toolCall.toolCallId,
-              })),
-            ),
-          );
-        });
-
-        // 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.
-        yield* submitMessage(
-          Obj.make(Message.Message, {
-            created: new Date().toISOString(),
-            sender: { role: 'tool' },
-            blocks: toolResults,
-          }),
-        );
-
-        this._toolCalls++;
+        yield* this.runTools({ toolkit });
       } while (true);
 
-      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: 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/templates/instructions/system.tpl

@@ -43,6 +43,15 @@
 - Suggested actions could include actions that create artifacts.
 - After creating an artifact use a tool to add it to the chat context.
 
-## Context
+## Navigation
+
+- When the user asks to open, go to, navigate to, or show a specific document, page, or object, use the navigation tools instead of loading content inline.
+- First, call the Resolve navigation targets tool to find the navigation path for the requested item.
+  - Without arguments, it returns available pages that can be navigated to.
+  - With a DXN in the query, it resolves that specific object to a navigation path.
+- If you know the object's DXN (e.g. from a database query or context objects), pass it in the query to get the exact navigation path.
+- Then, call the Open tool with the resolved navigation path(s) in the subject array to navigate the user to that item.
+- Do not respond with the document's content when the user asks to open it — open it in the main content area instead.
+- If multiple targets match, present the options to the user and let them choose which to open.
 
-Today is {{DATETIME}}.
+## Context

package/src/templates/system.ts

@@ -29,13 +29,14 @@ export const createSystemPrompt = (options: Partial<SystemPromptOptions> = {}) =
     }),
   );
 
+/**
+ * NOTE: We are forbidden from accessing randomness at model level in cloudflare, so we have to have pre-generated ids.
+ */
+const SYSTEM_ID = '01K2J2FZWH39F55CS7B26GP1Q3';
+
 export const templates = {
-  /** Editable system prompt.*/
   system: Template.make({
-    /**
-     * NOTE: We are forbidden from accessing randomness at model level in cloudflare, so we have to have pre-generated ids.
-     */
-    id: '01K2J2FZWH39F55CS7B26GP1Q3',
+    id: SYSTEM_ID,
     source: SYSTEM,
   }),
 };