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

package/src/conversation/conversation.ts

@@ -2,15 +2,20 @@
 // Copyright 2025 DXOS.org
 //
 
+import type * as Tool from '@effect/ai/Tool';
 import type * as Toolkit from '@effect/ai/Toolkit';
-import * as Array from 'effect/Array';
+import type { Registry } from '@effect-atom/atom-react';
+import * as Context from 'effect/Context';
 import * as Effect from 'effect/Effect';
-import * as Option from 'effect/Option';
+import * as Layer from 'effect/Layer';
 
+import { type ToolExecutionService, type ToolResolverService } from '@dxos/ai';
 import { Resource } from '@dxos/context';
 import { Obj } from '@dxos/echo';
 import { type Queue } from '@dxos/echo-db';
-import { DatabaseService } from '@dxos/echo-db';
+import { acquireReleaseResource } from '@dxos/effect';
+import { QueueService } from '@dxos/functions';
+import { invariant } from '@dxos/invariant';
 import { log } from '@dxos/log';
 import { Message } from '@dxos/types';
 
@@ -24,49 +29,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 +78,126 @@ 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));
   }
 
+  getTools(): Effect.Effect<Record<string, Tool.Any>, never, ToolExecutionService | ToolResolverService> {
+    return Effect.gen(this, function* () {
+      const blueprints = this.context.getBlueprints();
+      const tookit = yield* createToolkit({ toolkit: this._toolkit, blueprints });
+      return tookit.tools;
+    }).pipe(Effect.orDie);
+  }
+
   /**
    * 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)),
-      );
-
       // Create toolkit.
-      const toolkit = yield* createToolkit({ blueprints });
-      this._toolkit = toolkit;
+      const blueprints = this.context.getBlueprints();
+      const toolkit = yield* createToolkit({
+        toolkit: this._toolkit,
+        blueprints,
+      });
+
+      // Context objects.
+      const objects = this.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(
-        Effect.provideService(AiContextService, {
-          binder: this.context,
-        }),
+      const session = new AiSession();
+      const messages = yield* session.run({ history, blueprints, toolkit, objects, ...params }).pipe(
+        Effect.provide(
+          Layer.mergeAll(
+            Layer.succeed(AiContextService, {
+              binder: this.context,
+            }),
+            Layer.succeed(AiConversationService, this),
+          ),
+        ),
       );
 
-      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));
       return messages;
-    }).pipe(Effect.withSpan('AiConversation.request'));
+    });
   }
 }
+
+/**
+ * Gives access to the ai conversation.
+ */
+export class AiConversationService extends Context.Tag('@dxos/assistant/AiConversationService')<
+  AiConversationService,
+  AiConversation
+>() {
+  /**
+   * Create a new conversation layer from options.
+   */
+  static layer = (options: AiConversationOptions): Layer.Layer<AiConversationService | AiContextService> =>
+    aiContextFromConversation.pipe(
+      Layer.provideMerge(
+        Layer.scoped(
+          AiConversationService,
+          Effect.gen(function* () {
+            const conversation = yield* acquireReleaseResource(() => new AiConversation(options));
+            return conversation;
+          }),
+        ),
+      ),
+    );
+
+  /**
+   * Create a new conversation with a new queue.
+   */
+  static layerNewQueue = (
+    options?: Omit<AiConversationOptions, 'queue'>,
+  ): Layer.Layer<AiConversationService | AiContextService, never, QueueService> =>
+    Layer.unwrapScoped(
+      Effect.gen(function* () {
+        return AiConversationService.layer({
+          ...options,
+          queue: yield* QueueService.createQueue<Message.Message | ContextBinding>(),
+        });
+      }),
+    );
+
+  /**
+   * Run a prompt in the current conversation.
+   */
+  static run = (
+    params: AiConversationRunProps,
+  ): Effect.Effect<Message.Message[], AiSessionRunError, AiSessionRunRequirements | AiConversationService> =>
+    Effect.gen(function* () {
+      const conversation = yield* AiConversationService;
+      return yield* conversation.createRequest(params);
+    });
+}
+
+const aiContextFromConversation = Layer.effect(
+  AiContextService,
+  Effect.gen(function* () {
+    const conversation = yield* AiConversationService;
+    return {
+      binder: conversation.context,
+    };
+  }),
+);

package/src/diff.test.ts

@@ -2,7 +2,7 @@
 // Copyright 2025 DXOS.org
 //
 
-import { describe, expect, it } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
 import { createDocAccessor } from '@dxos/echo-db';
 import { EchoTestBuilder } from '@dxos/echo-db/testing';
@@ -12,10 +12,17 @@ import { trim } from '@dxos/util';
 import { applyDiffs } from './diff';
 
 describe('diff', () => {
+  let builder: EchoTestBuilder;
+  beforeEach(async () => {
+    builder = await new EchoTestBuilder().open();
+  });
+
+  afterEach(async () => {
+    await builder.close();
+  });
+
   it('should apply diffs', async () => {
-    const builder = new EchoTestBuilder();
-    const { db, graph } = await builder.createDatabase();
-    graph.schemaRegistry.addSchema([Text.Text]);
+    const { db } = await builder.createDatabase({ types: [Text.Text] });
 
     const document = trim`
       # Hello World
@@ -40,4 +47,28 @@ describe('diff', () => {
       But not this one.
     `);
   });
+
+  it('should apply diffs that append content at the end', async () => {
+    const { db } = await builder.createDatabase({ types: [Text.Text] });
+
+    const document = '# Shopping list';
+
+    const text = db.add(Text.make(document));
+    const accessor = createDocAccessor(text, ['content']);
+    const result = applyDiffs(accessor, ['- # Shopping list', '+ # Shopping list\n- Milk']);
+
+    expect(result).toBe('# Shopping list\n- Milk');
+  });
+
+  it('should append text to an empty document', async () => {
+    const builder = new EchoTestBuilder();
+    const { db, graph } = await builder.createDatabase();
+    await graph.schemaRegistry.register([Text.Text]);
+
+    const text = db.add(Text.make(''));
+    const accessor = createDocAccessor(text, ['content']);
+    const result = applyDiffs(accessor, ['+ # Shopping list']);
+
+    expect(result).toBe('# Shopping list');
+  });
 });

package/src/diff.ts

@@ -47,9 +47,10 @@ export const reduceDiffs = (diffs: readonly string[]): Diff[] => {
         }
         return acc;
       }
+      default: {
+        throw new Error(`Invalid diff format: each line must start with "- " or "+ ", got: ${JSON.stringify(diff)}`);
+      }
     }
-
-    return acc;
   }, []);
 };
 
@@ -65,7 +66,11 @@ export const computeDiffsWithCursors = <T>(accessor: DocAccessor<T>, diffs: read
     .filter(isNonNullable);
 };
 
-export const applyDiffs = <T>(accessor: DocAccessor<T>, diffs: readonly string[]): string => {
+export const applyDiffs = <T>(
+  accessor: DocAccessor<T>,
+  diffs: readonly string[],
+  { errorOnNotFound = false }: { errorOnNotFound?: boolean } = {},
+): string => {
   for (const diff of reduceDiffs(diffs)) {
     accessor.handle.change((doc: Doc<T>) => {
       const text = DocAccessor.getValue<string>(accessor);
@@ -75,7 +80,11 @@ export const applyDiffs = <T>(accessor: DocAccessor<T>, diffs: readonly string[]
         // NOTE: This only replaces the first match.
         A.splice(doc, accessor.path as A.Prop[], idx, diff.match.length, diff.replace);
       } else {
-        log.warn('diff not found', diff);
+        if (errorOnNotFound) {
+          throw new Error(`Diff not found: ${diff.match}`);
+        } else {
+          log.warn('diff not found', diff);
+        }
       }
     });
   }

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.test.ts

@@ -6,7 +6,7 @@ import { describe, expect, it } from '@effect/vitest';
 
 import { AgentStatus } from '@dxos/ai';
 import { Obj } from '@dxos/echo';
-import { ObjectId } from '@dxos/echo/internal';
+import { ObjectId } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { type ContentBlock, Message } from '@dxos/types';
 

package/src/execution-graph.ts

@@ -3,7 +3,7 @@
 //
 
 import { AgentStatus } from '@dxos/ai';
-import { Obj, type Ref } from '@dxos/echo';
+import { type Entity, Obj, type Ref } from '@dxos/echo';
 import { MESSAGE_PROPERTY_TOOL_CALL_ID } from '@dxos/functions';
 import { type ObjectId } from '@dxos/keys';
 import { LogLevel } from '@dxos/log';
@@ -65,7 +65,7 @@ export class ExecutionGraph {
   /**
    * Adds events to the graph.
    */
-  addEvents(events: Obj.Any[]) {
+  addEvents(events: Entity.Unknown[]) {
     const sortedEvents = this.sortObjectsByCreated(events);
     for (const event of sortedEvents) {
       if (Obj.instanceOf(Message.Message, event)) {
@@ -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,11 +4,11 @@
 
 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';
-import { Testing } from '@dxos/schema/testing';
+import { TestSchema } from '@dxos/schema/testing';
 import { type Message, Organization, Person } from '@dxos/types';
 import { createTestData } from '@dxos/types/testing';
 import { range } from '@dxos/util';
@@ -19,25 +19,25 @@ 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[];
     transcriptWoflram: Message.Message[];
     transcriptMessages: Message.Message[];
-    documents: Testing.DocumentType[];
+    documents: TestSchema.DocumentType[];
     contacts: Record<string, Person.Person>;
     organizations: Record<string, Organization.Organization>;
   };
 
-  const TYPES = [Organization.Organization, Person.Person, Testing.Person, Testing.DocumentType];
+  const TYPES = [Organization.Organization, Person.Person, TestSchema.Person, TestSchema.DocumentType];
 
   beforeAll(async () => {
     // TODO(dmaretskyi): Helper to scaffold this from a config.
     builder = await new EchoTestBuilder().open();
-    const { db: db1 } = await builder.createDatabase({ indexing: { vector: true } });
+    const { db: db1 } = await builder.createDatabase();
     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

@@ -5,7 +5,7 @@
 import * as Schema from 'effect/Schema';
 
 import { asyncTimeout } from '@dxos/async';
-import { Expando } from '@dxos/echo/internal';
+import { Type } from '@dxos/echo';
 import { type FunctionDefinition } from '@dxos/functions';
 import { type FunctionExecutor } from '@dxos/functions-runtime';
 import { log } from '@dxos/log';
@@ -13,7 +13,7 @@ import { Message } from '@dxos/types';
 
 export const ExtractionInput = Schema.Struct({
   message: Message.Message,
-  objects: Schema.optional(Schema.Array(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,9 +4,8 @@
 
 import * as Schema from 'effect/Schema';
 
-import { Filter, Query } from '@dxos/echo';
-import { ObjectId } from '@dxos/echo/internal';
-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 { Function, FunctionDefinition, FunctionInvocationService } from '@dxos/functions';
+import { Type } from '@dxos/echo';
+import { type FunctionDefinition, FunctionInvocationService } from '@dxos/functions';
 import { invariant } from '@dxos/invariant';
 
+import { GenericToolkit } from '../session';
+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,51 +29,50 @@ 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> => {
+export const makeToolResolverFromFunctions = (): Layer.Layer<
+  ToolResolverService,
+  never,
+  GenericToolkit.Provider | FunctionInvocationService
+> => {
   return Layer.effect(
     ToolResolverService,
     Effect.gen(function* () {
-      const dbService = yield* DatabaseService;
+      const toolkitProvider = yield* GenericToolkit.Provider;
+      const functionInvocationService = yield* FunctionInvocationService;
       return {
         resolve: (id): Effect.Effect<Tool.Any, AiToolNotFoundError> =>
           Effect.gen(function* () {
-            const tool = toolkit.tools[id];
+            const toolkit = toolkitProvider.getToolkit();
+
+            const tool = toolkit.toolkit.tools[id];
             if (tool) {
               return tool;
             }
 
-            const {
-              objects: [dbFunction],
-            } = yield* DatabaseService.runQuery(Query.type(Function.Function, { key: id }));
-
-            const functionDef = dbFunction
-              ? FunctionDefinition.deserialize(dbFunction)
-              : functions.find((fn) => fn.key === id);
-
-            if (!functionDef) {
-              return yield* Effect.fail(new AiToolNotFoundError(id));
-            }
-
-            return projectFunctionToTool(functionDef);
-          }).pipe(Effect.provideService(DatabaseService, dbService)),
+            return yield* functionInvocationService.resolveFunction(id).pipe(
+              Effect.map(projectFunctionToTool),
+              Effect.catchTag('FunctionNotFound', () => Effect.fail(new AiToolNotFoundError(id))),
+            );
+          }),
       } satisfies Context.Tag.Service<ToolResolverService>;
     }),
   );
 };
 
-export const makeToolExecutionServiceFromFunctions = (
-  toolkit: Toolkit.Toolkit<any>,
-  handlersLayer: Layer.Layer<Tool.Handler<any>, never, never>,
-): Layer.Layer<ToolExecutionService, never, FunctionInvocationService> => {
+export const makeToolExecutionServiceFromFunctions = (): Layer.Layer<
+  ToolExecutionService,
+  never,
+  FunctionInvocationService | GenericToolkit.Provider
+> => {
   return Layer.effect(
     ToolExecutionService,
     Effect.gen(function* () {
-      const toolkitHandler = yield* toolkit.pipe(Effect.provide(handlersLayer));
+      const toolkitProvider = yield* GenericToolkit.Provider;
+      const toolkit = toolkitProvider.getToolkit();
+
+      const toolkitHandler = yield* toolkit.toolkit.pipe(Effect.provide(toolkit.layer));
       invariant(isHandlerLike(toolkitHandler));
       const functionInvocationService = yield* FunctionInvocationService;
 
@@ -82,6 +84,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;
@@ -104,6 +107,11 @@ export const makeToolExecutionServiceFromFunctions = (
   );
 };
 
+export const ToolExecutionServices = Layer.mergeAll(
+  makeToolResolverFromFunctions(),
+  makeToolExecutionServiceFromFunctions(),
+);
+
 class FunctionToolAnnotation extends Context.Tag('@dxos/assistant/FunctionToolAnnotation')<
   FunctionToolAnnotation,
   { definition: FunctionDefinition.Any }
@@ -149,7 +157,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 +167,21 @@ 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)) {
+    const description = ast.annotations.description
+      ? ast.annotations.description + '\n' + RefFromLLM.ast.annotations.description
+      : (RefFromLLM.ast.annotations.description as string);
+    return RefFromLLM.annotations({ description }).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';
 };