@dxos/assistant 0.8.4-main.7ace549 → 0.8.4-main.937b3ca

package/src/conversation/conversation.ts

@@ -3,14 +3,13 @@
 //
 
 import type * as Toolkit from '@effect/ai/Toolkit';
-import * as Array from 'effect/Array';
+import type { Registry } from '@effect-atom/atom-react';
 import * as Effect from 'effect/Effect';
-import * as Option from 'effect/Option';
 
 import { Resource } from '@dxos/context';
 import { Obj } from '@dxos/echo';
 import { type Queue } from '@dxos/echo-db';
-import { DatabaseService } from '@dxos/echo-db';
+import { invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
 import { Message } from '@dxos/types';
 
@@ -24,49 +23,48 @@ import {
 
 import { AiContextBinder, AiContextService, type ContextBinding } from './context';
 
-export interface AiConversationRunParams {
+export interface AiConversationRunProps {
   prompt: string;
   system?: string;
   observer?: GenerationObserver;
 }
 
+export type AiConversationOptions = {
+  queue: Queue<Message.Message | ContextBinding>;
+  toolkit?: Toolkit.Any;
+  registry?: Registry.Registry;
+};
+
 /**
  * Durable conversation state (initiated by users and agents) backed by a Queue.
  * Executes tools based on AI responses and supports cancellation of in-progress requests.
  */
 export class AiConversation extends Resource {
   /**
-   * Message and binding queue.
+   * Blueprints and objects bound to the conversation.
    */
+  private readonly _binder: AiContextBinder;
   private readonly _queue: Queue<Message.Message | ContextBinding>;
+  private readonly _toolkit?: Toolkit.Any;
 
-  /**
-   * Blueprints bound to the conversation.
-   */
-  private readonly _context: AiContextBinder;
-
-  /**
-   * Toolkit from the current session request.
-   */
-  private _toolkit: Toolkit.WithHandler<any> | undefined;
-
-  public constructor(queue: Queue<Message.Message | ContextBinding>) {
+  public constructor(options: AiConversationOptions) {
     super();
-    this._queue = queue;
-    this._context = new AiContextBinder(this._queue);
+    this._queue = options.queue;
+    this._toolkit = options.toolkit;
+    invariant(this._queue);
+    this._binder = new AiContextBinder({ queue: this._queue, registry: options.registry });
   }
 
   protected override async _open(): Promise<void> {
-    // TODO(wittjosiah): Pass in parent context?
-    await this._context.open();
+    await this._binder.open(this._ctx);
   }
 
-  protected override async _close(): Promise<void> {
-    await this._context.close();
+  public get queue() {
+    return this._queue;
   }
 
   public get context() {
-    return this._context;
+    return this._binder;
   }
 
   public get toolkit() {
@@ -74,55 +72,54 @@ export class AiConversation extends Resource {
   }
 
   public async getHistory(): Promise<Message.Message[]> {
-    const queueItems = await this._queue.queryObjects();
+    const queueItems = await this._queue.queryObjects(); // TODO(burdon): Update.
     return queueItems.filter(Obj.instanceOf(Message.Message));
   }
 
   /**
    * Creates a new cancelable request effect.
    */
-  createRequest(
-    params: AiConversationRunParams,
+  public createRequest(
+    params: AiConversationRunProps,
   ): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements> {
-    return Effect.gen(this, function* () {
-      const session = new AiSession();
-      const history = yield* Effect.promise(() => this.getHistory());
-
-      // Get context objects.
-      const context = yield* Effect.promise(() => this.context.query());
-      const blueprints = yield* Effect.forEach(context.blueprints.values(), DatabaseService.loadOption).pipe(
-        Effect.map(Array.filter(Option.isSome)),
-        Effect.map(Array.map((option) => option.value)),
-      );
-      const objects = yield* Effect.forEach(context.objects.values(), DatabaseService.loadOption).pipe(
-        Effect.map(Array.filter(Option.isSome)),
-        Effect.map(Array.map((option) => option.value)),
-      );
+    const self = this;
+    return Effect.gen(function* () {
+      const history = yield* Effect.promise(() => self.getHistory());
 
       // Create toolkit.
-      const toolkit = yield* createToolkit({ blueprints });
-      this._toolkit = toolkit;
+      const blueprints = self.context.getBlueprints();
+      const toolkit = yield* createToolkit({
+        toolkit: self._toolkit,
+        blueprints,
+      });
+
+      // Context objects.
+      const objects = self.context.getObjects();
 
-      const start = Date.now();
       log('run', {
         history: history.length,
         blueprints: blueprints.length,
+        tools: Object.keys(toolkit.tools).length,
         objects: objects.length,
-        tools: this._toolkit?.tools.length ?? 0,
       });
 
       // Process request.
-      const messages = yield* session.run({ history, blueprints, objects, toolkit, ...params }).pipe(
+      const session = new AiSession();
+      const messages = yield* session.run({ history, blueprints, toolkit, objects, ...params }).pipe(
         Effect.provideService(AiContextService, {
-          binder: this.context,
+          binder: self.context,
         }),
       );
 
-      log('result', { messages: messages.length, duration: Date.now() - start });
+      log('result', {
+        messages: messages.length,
+        duration: session.duration,
+        toolCalls: session.toolCalls,
+      });
 
       // Append to queue.
-      yield* Effect.promise(() => this._queue.append(messages));
+      yield* Effect.promise(() => self._queue.append(messages));
       return messages;
-    }).pipe(Effect.withSpan('AiConversation.request'));
+    });
   }
 }

package/src/diff.test.ts

@@ -15,7 +15,7 @@ describe('diff', () => {
   it('should apply diffs', async () => {
     const builder = new EchoTestBuilder();
     const { db, graph } = await builder.createDatabase();
-    graph.schemaRegistry.addSchema([Text.Text]);
+    await graph.schemaRegistry.register([Text.Text]);
 
     const document = trim`
       # Hello World

package/src/errors.ts

@@ -7,9 +7,9 @@ import { BaseError } from '@dxos/errors';
 /**
  * Generic error from AI model.
  */
-export class AiModelError extends BaseError.extend('AI_MODEL_ERROR', 'AI model error') {}
+export class AiModelError extends BaseError.extend('AiModelError', 'AI model error') {}
 
 /**
  * Generic error for AI agent execution.
  */
-export class AiAssistantError extends BaseError.extend('AI_ASSISTANT_ERROR', 'AI assistant error') {}
+export class AiAssistantError extends BaseError.extend('AiAssistantError', 'AI assistant error') {}

package/src/execution-graph.ts

@@ -496,7 +496,7 @@ const getParentId = (message: Message.Message) => {
   }
 };
 
-const stringifyRef = (ref: Ref.Any) => {
+const stringifyRef = (ref: Ref.Unknown) => {
   if (ref.target) {
     return stringifyObject(ref.target);
   }
@@ -504,6 +504,6 @@ const stringifyRef = (ref: Ref.Any) => {
   return ref.dxn.asEchoDXN()?.echoId ?? ref.dxn.asQueueDXN()?.objectId ?? '';
 };
 
-const stringifyObject = (obj: Obj.Any) => {
+const stringifyObject = (obj: Obj.Unknown) => {
   return Obj.getLabel(obj) ?? Obj.getTypename(obj) ?? obj.id;
 };

package/src/extraction/extraction-ner-function/extraction-ner-function.test.ts

@@ -4,7 +4,7 @@
 
 import { beforeAll, describe, test } from 'vitest';
 
-import { type EchoDatabase } from '@dxos/echo-db';
+import { type Database } from '@dxos/echo';
 import { EchoTestBuilder } from '@dxos/echo-db/testing';
 import { FunctionExecutor, ServiceContainer } from '@dxos/functions-runtime';
 import { log } from '@dxos/log';
@@ -19,7 +19,7 @@ import { extractionNerFunction } from './extraction-ner-function';
 
 describe.skip('NER EntityExtraction', () => {
   let builder: EchoTestBuilder;
-  let db: EchoDatabase;
+  let db: Database.Database;
   let executor: FunctionExecutor;
   let testData: {
     transcriptJosiah: Message.Message[];
@@ -37,7 +37,7 @@ describe.skip('NER EntityExtraction', () => {
     builder = await new EchoTestBuilder().open();
     const { db: db1 } = await builder.createDatabase({ indexing: { vector: true } });
     db = db1;
-    db.graph.schemaRegistry.addSchema(TYPES);
+    await db.graph.schemaRegistry.register(TYPES);
     const data = createTestData();
     testData = {
       transcriptJosiah: data.transcriptJosiah,

package/src/extraction/extraction-ner-function/extraction-ner-function.ts

@@ -20,7 +20,7 @@ export const extractionNerFunction = defineFunction({
     throw new Error('not implemented');
     // log.info('input', { message, options });
     // const startTime = performance.now();
-    // const { db } = context.getService(DatabaseService);
+    // const { db } = context.getService(Database.Service);
     // const { extractFullEntities } = await import('./named-entity-recognition');
 
     // const entitiesPromise = Promise.all(

package/src/extraction/extraction.ts

@@ -13,7 +13,7 @@ import { Message } from '@dxos/types';
 
 export const ExtractionInput = Schema.Struct({
   message: Message.Message,
-  objects: Schema.optional(Schema.Array(Type.Expando)),
+  objects: Schema.optional(Schema.Array(Type.Obj)),
   options: Schema.optional(
     Schema.Struct({
       timeout: Schema.optional(Schema.Number),
@@ -29,7 +29,7 @@ export const ExtractionOutput = Schema.Struct({
 });
 export interface ExtractionOutput extends Schema.Schema.Type<typeof ExtractionOutput> {}
 
-export type ProcessTranscriptMessageParams = {
+export type ProcessTranscriptMessageProps = {
   input: ExtractionInput;
   function: FunctionDefinition<ExtractionInput, ExtractionOutput>;
   executor: FunctionExecutor;
@@ -55,7 +55,7 @@ export type ExtractionFunction = FunctionDefinition<ExtractionInput, ExtractionO
  * Extract entities from the transcript message and add them to the message.
  */
 // TODO(dmaretskyi): Move context to a vector search index.
-export const processTranscriptMessage = async (params: ProcessTranscriptMessageParams): Promise<ExtractionOutput> => {
+export const processTranscriptMessage = async (params: ProcessTranscriptMessageProps): Promise<ExtractionOutput> => {
   try {
     if (params.options?.timeout && params.options.timeout > 0) {
       return await asyncTimeout(params.executor.invoke(params.function, params.input), params.options.timeout);

package/src/extraction/quotes.ts

@@ -4,8 +4,7 @@
 
 import * as Schema from 'effect/Schema';
 
-import { Filter, Query } from '@dxos/echo';
-import { type EchoDatabase } from '@dxos/echo-db';
+import { type Database, Filter, Query } from '@dxos/echo';
 import { ObjectId } from '@dxos/keys';
 import { DXN } from '@dxos/keys';
 import { log } from '@dxos/log';
@@ -28,10 +27,10 @@ export const ReferencedQuotes = Schema.Struct({
 });
 export type ReferencedQuotes = Schema.Schema.Type<typeof ReferencedQuotes>;
 
-export const findQuotes = async (quotes: string[], db: EchoDatabase): Promise<ReferencedQuotes> => {
+export const findQuotes = async (quotes: string[], db: Database.Database): Promise<ReferencedQuotes> => {
   const references = await Promise.all(
     quotes.map(async (quote) => {
-      const { objects } = await db.query(Query.select(Filter.text(quote, { type: 'vector' }))).run();
+      const objects = await db.query(Query.select(Filter.text(quote, { type: 'vector' }))).run();
       log.info('found quotes in vector index', { objects, quote });
       return objects.length > 0 ? [{ id: objects[0].id.toString(), quote }] : [];
     }),

package/src/functions/services.ts

@@ -9,14 +9,17 @@ import * as Effect from 'effect/Effect';
 import * as Layer from 'effect/Layer';
 import * as Record from 'effect/Record';
 import * as Schema from 'effect/Schema';
+import type * as SchemaAST from 'effect/SchemaAST';
 
 import { AiToolNotFoundError, ToolExecutionService, ToolResolverService } from '@dxos/ai';
 import { todo } from '@dxos/debug';
-import { Query } from '@dxos/echo';
-import { DatabaseService } from '@dxos/echo-db';
+import { Query, Type } from '@dxos/echo';
+import { Database } from '@dxos/echo';
 import { Function, FunctionDefinition, FunctionInvocationService } from '@dxos/functions';
 import { invariant } from '@dxos/invariant';
 
+import { RefFromLLM } from '../types';
+
 /**
  * Constructs a `ToolResolverService` whose `resolve(id)` looks up tools in the following order:
  *  1. Toolkit: return an existing tool from the provided `toolkit` if one is present under `id`.
@@ -26,27 +29,27 @@ import { invariant } from '@dxos/invariant';
  *
  * If none of the above yield a match, the effect fails with `AiToolNotFoundError`.
  *
- * Requires `DatabaseService` in the environment.
+ * Requires `Database.Service` in the environment.
  */
 export const makeToolResolverFromFunctions = (
   functions: FunctionDefinition.Any[],
   toolkit: Toolkit.Toolkit<any>,
-): Layer.Layer<ToolResolverService, never, DatabaseService> => {
+): Layer.Layer<ToolResolverService, never, Database.Service> => {
   return Layer.effect(
     ToolResolverService,
     Effect.gen(function* () {
-      const dbService = yield* DatabaseService;
+      const dbService = yield* Database.Service;
       return {
         resolve: (id): Effect.Effect<Tool.Any, AiToolNotFoundError> =>
           Effect.gen(function* () {
+            // TODO(dmaretskyi): Use FunctionInvocationService.resolveFunction().
+
             const tool = toolkit.tools[id];
             if (tool) {
               return tool;
             }
 
-            const {
-              objects: [dbFunction],
-            } = yield* DatabaseService.runQuery(Query.type(Function.Function, { key: id }));
+            const [dbFunction] = yield* Database.Service.runQuery(Query.type(Function.Function, { key: id }));
 
             const functionDef = dbFunction
               ? FunctionDefinition.deserialize(dbFunction)
@@ -57,7 +60,7 @@ export const makeToolResolverFromFunctions = (
             }
 
             return projectFunctionToTool(functionDef);
-          }).pipe(Effect.provideService(DatabaseService, dbService)),
+          }).pipe(Effect.provideService(Database.Service, dbService)),
       } satisfies Context.Tag.Service<ToolResolverService>;
     }),
   );
@@ -82,6 +85,7 @@ export const makeToolExecutionServiceFromFunctions = (
                 if (Tool.isProviderDefined(tool)) {
                   throw new Error('Attempted to call a provider-defined tool');
                 }
+
                 // TODO(wittjosiah): Everything is `never` here.
                 const { result } = yield* (toolkitHandler.handle as any)(tool.name, input);
                 return result;
@@ -149,7 +153,9 @@ const makeToolName = (name: string) => {
 const createStructFieldsFromSchema = (schema: Schema.Schema<any, any>): Record<string, Schema.Schema<any, any>> => {
   switch (schema.ast._tag) {
     case 'TypeLiteral':
-      return Object.fromEntries(schema.ast.propertySignatures.map((prop) => [prop.name, Schema.make(prop.type)]));
+      return Object.fromEntries(
+        schema.ast.propertySignatures.map((prop) => [prop.name, Schema.make(mapSchemaTypeForLLM(prop.type))]),
+      );
     case 'VoidKeyword':
       return {};
     default:
@@ -157,6 +163,18 @@ const createStructFieldsFromSchema = (schema: Schema.Schema<any, any>): Record<s
   }
 };
 
+/**
+ * Picks an LLM-friendly schema type for the given schema AST.
+ * The picked schema type decodes to the original schema type.
+ */
+const mapSchemaTypeForLLM = (ast: SchemaAST.AST): SchemaAST.AST => {
+  if (Type.Ref.isRefSchemaAST(ast)) {
+    return RefFromLLM.ast;
+  }
+
+  return ast;
+};
+
 const isHandlerLike = (value: unknown): value is Toolkit.WithHandler<Record<string, Tool.Any>> => {
   return typeof (value as any).tools === 'object' && typeof (value as any).handle === 'function';
 };