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

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

@@ -5,9 +5,9 @@
 import { beforeAll, describe, test } from 'vitest';
 
 import { todo } from '@dxos/debug';
-import { FunctionExecutor } from '@dxos/functions';
+import { FunctionExecutor } from '@dxos/functions-runtime';
 import { log } from '@dxos/log';
-import { createTestData } from '@dxos/schema/testing';
+import { createTestData } from '@dxos/types/testing';
 import { range } from '@dxos/util';
 
 import { processTranscriptMessage } from '../extraction';

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

@@ -2,20 +2,9 @@
 // Copyright 2025 DXOS.org
 //
 
-// ISSUE(burdon): defineFunction
-// @ts-nocheck
-
-import { AiService } from '@dxos/ai';
-import { Obj } from '@dxos/echo';
-import { create } from '@dxos/echo/internal';
 import { type FunctionDefinition, defineFunction } from '@dxos/functions';
-import { type ContentBlock, DataType } from '@dxos/schema';
 
-import { AiSession } from '../../session';
 import { ExtractionInput, ExtractionOutput } from '../extraction';
-import { ReferencedQuotes, insertReferences } from '../quotes';
-
-import PROMPT from './instructions.tpl?raw';
 
 export const extractionAnthropicFunction: FunctionDefinition<ExtractionInput, ExtractionOutput> = defineFunction({
   key: 'dxos.org/function/extraction/extract-entities',
@@ -24,49 +13,50 @@ export const extractionAnthropicFunction: FunctionDefinition<ExtractionInput, Ex
   inputSchema: ExtractionInput,
   outputSchema: ExtractionOutput,
   handler: async ({ data: { message, objects }, context }) => {
-    const startTime = performance.now();
-    const ai = context.getService(AiService.AiService);
-    const session = new AiSession({ operationModel: 'configured' });
-    const result = await session.runStructured(ReferencedQuotes, {
-      generationOptions: {
-        model: '@anthropic/claude-3-5-haiku-20241022',
-      },
-      client: ai.client,
-      systemPrompt: PROMPT,
-      history: [
-        Obj.make(DataType.Message, {
-          created: new Date().toISOString(),
-          sender: { role: 'user' },
-          blocks: [
-            {
-              _tag: 'text',
-              text: `<context>${JSON.stringify(objects)}</context>`,
-            } satisfies ContentBlock.Text,
-            {
-              _tag: 'text',
-              text: `<transcript>${JSON.stringify(message.blocks)}</transcript>`,
-            } satisfies ContentBlock.Text,
-          ],
-        }),
-      ],
-      artifacts: [],
-      prompt: '',
-      tools: [],
-    } as any); // TODO(burdon): Rewrite test.
+    throw new Error('not implemented');
+    // const startTime = performance.now();
+    // const ai = context.getService(AiService.AiService);
+    // const session = new AiSession({ operationModel: 'configured' });
+    // const result = await session.runStructured(ReferencedQuotes, {
+    //   generationOptions: {
+    //     model: '@anthropic/claude-3-5-haiku-20241022',
+    //   },
+    //   client: ai.client,
+    //   systemPrompt: PROMPT,
+    //   history: [
+    //     Obj.make(DataType.Message.Message, {
+    //       created: new Date().toISOString(),
+    //       sender: { role: 'user' },
+    //       blocks: [
+    //         {
+    //           _tag: 'text',
+    //           text: `<context>${JSON.stringify(objects)}</context>`,
+    //         } satisfies DataType.ContentBlock.Text,
+    //         {
+    //           _tag: 'text',
+    //           text: `<transcript>${JSON.stringify(message.blocks)}</transcript>`,
+    //         } satisfies DataType.ContentBlock.Text,
+    //       ],
+    //     }),
+    //   ],
+    //   artifacts: [],
+    //   prompt: '',
+    //   tools: [],
+    // } as any); // TODO(burdon): Rewrite test.
 
-    return {
-      timeElapsed: performance.now() - startTime,
-      message: create(DataType.Message, {
-        ...message,
-        blocks: message.blocks.map((block) =>
-          block._tag !== 'transcript'
-            ? block
-            : {
-                ...block,
-                text: insertReferences(block.text, result),
-              },
-        ),
-      }),
-    };
+    // return {
+    //   timeElapsed: performance.now() - startTime,
+    //   message: create(DataType.Message.Message, {
+    //     ...message,
+    //     blocks: message.blocks.map((block) =>
+    //       block._tag !== 'transcript'
+    //         ? block
+    //         : {
+    //             ...block,
+    //             text: insertReferences(block.text, result),
+    //           },
+    //     ),
+    //   }),
+    // };
   },
 });

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

@@ -4,12 +4,13 @@
 
 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';
+import { FunctionExecutor, ServiceContainer } from '@dxos/functions-runtime';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
-import { Testing, createTestData } 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';
 
 import { processTranscriptMessage } from '../extraction';
@@ -18,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: DataType.Message[];
-    transcriptWoflram: DataType.Message[];
-    transcriptMessages: DataType.Message[];
-    documents: Testing.DocumentType[];
-    contacts: Record<string, DataType.Person>;
-    organizations: Record<string, DataType.Organization>;
+    transcriptJosiah: Message.Message[];
+    transcriptWoflram: Message.Message[];
+    transcriptMessages: Message.Message[];
+    documents: TestSchema.DocumentType[];
+    contacts: Record<string, Person.Person>;
+    organizations: Record<string, Organization.Organization>;
   };
 
-  const TYPES = [DataType.Organization, DataType.Person, Testing.Contact, 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 } });
     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

@@ -2,18 +2,9 @@
 // Copyright 2025 DXOS.org
 //
 
-// ISSUE(burdon): defineFunction
-// @ts-nocheck
-
-import { asyncTimeout } from '@dxos/async';
-import { create } from '@dxos/echo/internal';
-import { DatabaseService, defineFunction } from '@dxos/functions';
-import { invariant } from '@dxos/invariant';
-import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { defineFunction } from '@dxos/functions';
 
 import { ExtractionInput, ExtractionOutput } from '../extraction';
-import { findQuotes, insertReferences } from '../quotes';
 
 /**
  * Named Entity Recognition model
@@ -26,39 +17,40 @@ export const extractionNerFunction = defineFunction({
   inputSchema: ExtractionInput,
   outputSchema: ExtractionOutput,
   handler: async ({ data: { message, options }, context }) => {
-    log.info('input', { message, options });
-    const startTime = performance.now();
-    const { db } = context.getService(DatabaseService);
-    const { extractFullEntities } = await import('./named-entity-recognition');
-
-    const entitiesPromise = Promise.all(
-      message.blocks.map(async (block) => {
-        invariant(block._tag === 'transcript' || block._tag === 'text');
-        return extractFullEntities(block.text);
-      }),
-    ).then((entities) => entities.flat());
-    const entities = options?.timeout ? await asyncTimeout(entitiesPromise, options.timeout) : await entitiesPromise;
-
-    log.info('entities', { entities });
-    const quoteReferences = await findQuotes(
-      entities.map((entity) => entity.word),
-      db,
-    );
-
-    const blocksWithReferences = message.blocks.map((block) => {
-      invariant(block._tag === 'transcript' || block._tag === 'text');
-      return { ...block, text: insertReferences(block._tag, quoteReferences) };
-    });
-
-    const messageWithReferences = create(DataType.Message, {
-      ...message,
-      blocks: blocksWithReferences,
-    });
-
-    log.info('output', { messageWithReferences });
-    return {
-      message: messageWithReferences,
-      timeElapsed: performance.now() - startTime,
-    };
+    throw new Error('not implemented');
+    // log.info('input', { message, options });
+    // const startTime = performance.now();
+    // const { db } = context.getService(Database.Service);
+    // const { extractFullEntities } = await import('./named-entity-recognition');
+
+    // const entitiesPromise = Promise.all(
+    //   message.blocks.map(async (block) => {
+    //     invariant(block._tag === 'transcript' || block._tag === 'text');
+    //     return extractFullEntities(block.text);
+    //   }),
+    // ).then((entities) => entities.flat());
+    // const entities = options?.timeout ? await asyncTimeout(entitiesPromise, options.timeout) : await entitiesPromise;
+
+    // log.info('entities', { entities });
+    // const quoteReferences = await findQuotes(
+    //   entities.map((entity) => entity.word),
+    //   db,
+    // );
+
+    // const blocksWithReferences = message.blocks.map((block) => {
+    //   invariant(block._tag === 'transcript' || block._tag === 'text');
+    //   return { ...block, text: insertReferences(block._tag, quoteReferences) };
+    // });
+
+    // const messageWithReferences = create(Message.Message, {
+    //   ...message,
+    //   blocks: blocksWithReferences,
+    // });
+
+    // log.info('output', { messageWithReferences });
+    // return {
+    //   message: messageWithReferences,
+    //   timeElapsed: performance.now() - startTime,
+    // };
   },
 });

package/src/extraction/extraction-ner-function/named-entity-recognition.test.ts

@@ -6,7 +6,8 @@ import { type TokenClassificationOutput, type TokenClassificationPipelineType }
 import { beforeAll, describe, expect, test } from 'vitest';
 
 import { log } from '@dxos/log';
-import { createTestData } from '@dxos/schema/testing';
+import { type ContentBlock, type Message } from '@dxos/types';
+import { createTestData } from '@dxos/types/testing';
 
 import { combineNerTokens, createTokenGroups, extractFullEntities, getNer } from './named-entity-recognition';
 
@@ -57,8 +58,10 @@ describe.skip('NamedEntityRecognition', () => {
   test('own data', async () => {
     const { transcriptWoflram, transcriptJosiah } = await createTestData();
 
-    const blocks = [...transcriptWoflram, ...transcriptJosiah].flatMap((message) =>
-      message.blocks.flatMap((block) => (block._tag === 'transcript' ? [block.text] : [])),
+    const blocks = [...transcriptWoflram, ...transcriptJosiah].flatMap((message: Message.Message) =>
+      message.blocks.flatMap((block: ContentBlock.Any) =>
+        block._tag === 'transcript' ? [(block as ContentBlock.Transcript).text] : [],
+      ),
     );
 
     for (const block of blocks) {
@@ -70,8 +73,10 @@ describe.skip('NamedEntityRecognition', () => {
   test('more own data', async () => {
     const { transcriptMessages } = await createTestData();
 
-    const blocks = transcriptMessages.flatMap((message) =>
-      message.blocks.flatMap((block) => (block._tag === 'transcript' ? [block.text] : [])),
+    const blocks = transcriptMessages.flatMap((message: Message.Message) =>
+      message.blocks.flatMap((block: ContentBlock.Any) =>
+        block._tag === 'transcript' ? [(block as ContentBlock.Transcript).text] : [],
+      ),
     );
 
     for (const block of blocks) {

package/src/extraction/extraction.ts

@@ -5,14 +5,15 @@
 import * as Schema from 'effect/Schema';
 
 import { asyncTimeout } from '@dxos/async';
-import { Expando } from '@dxos/echo/internal';
-import { type FunctionDefinition, type FunctionExecutor } from '@dxos/functions';
+import { Type } from '@dxos/echo';
+import { type FunctionDefinition } from '@dxos/functions';
+import { type FunctionExecutor } from '@dxos/functions-runtime';
 import { log } from '@dxos/log';
-import { DataType } from '@dxos/schema';
+import { Message } from '@dxos/types';
 
 export const ExtractionInput = Schema.Struct({
-  message: DataType.Message,
-  objects: Schema.optional(Schema.Array(Expando)),
+  message: Message.Message,
+  objects: Schema.optional(Schema.Array(Type.Obj)),
   options: Schema.optional(
     Schema.Struct({
       timeout: Schema.optional(Schema.Number),
@@ -23,12 +24,12 @@ export const ExtractionInput = Schema.Struct({
 export interface ExtractionInput extends Schema.Schema.Type<typeof ExtractionInput> {}
 
 export const ExtractionOutput = Schema.Struct({
-  message: DataType.Message,
+  message: Message.Message,
   timeElapsed: Schema.Number,
 });
 export interface ExtractionOutput extends Schema.Schema.Type<typeof ExtractionOutput> {}
 
-export type ProcessTranscriptMessageParams = {
+export type ProcessTranscriptMessageProps = {
   input: ExtractionInput;
   function: FunctionDefinition<ExtractionInput, ExtractionOutput>;
   executor: FunctionExecutor;
@@ -54,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,13 +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, Function, FunctionDefinition, FunctionInvocationService } from '@dxos/functions';
+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`.
@@ -25,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, any>[],
+  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)
@@ -56,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>;
     }),
   );
@@ -81,8 +85,10 @@ export const makeToolExecutionServiceFromFunctions = (
                 if (Tool.isProviderDefined(tool)) {
                   throw new Error('Attempted to call a provider-defined tool');
                 }
+
                 // TODO(wittjosiah): Everything is `never` here.
-                return yield* (toolkitHandler.handle as any)(tool.name, input);
+                const { result } = yield* (toolkitHandler.handle as any)(tool.name, input);
+                return result;
               }
 
               const { definition: functionDef } = Context.get(FunctionToolAnnotation)(tool.annotations as any);
@@ -104,24 +110,26 @@ export const makeToolExecutionServiceFromFunctions = (
 
 class FunctionToolAnnotation extends Context.Tag('@dxos/assistant/FunctionToolAnnotation')<
   FunctionToolAnnotation,
-  { definition: FunctionDefinition<any, any> }
+  { definition: FunctionDefinition.Any }
 >() {}
 
-const toolCache = new WeakMap<FunctionDefinition<any, any>, Tool.Any>();
+const toolCache = new WeakMap<FunctionDefinition.Any, Tool.Any>();
 
 /**
  * Projects a `FunctionDefinition` into an `AiTool`.
  *
  * @param fn The function definition to project into a tool.
  * @param meta Optional projection metadata.
- * @param meta.deployedFunctionId Backend deployment ID used for remote invocation when present. This is the
- *        EDGE service's function deployment identifier (not the ECHO object ID/DXN and not `FunctionDefinition.key`).
+ * @param meta.deployedFunctionId Backend deployment ID used for remote invocation when present.
+ *    This is the EDGE service's function deployment identifier (not the ECHO object ID/DXN and not `FunctionDefinition.key`).
  */
-const projectFunctionToTool = (fn: FunctionDefinition<any, any>): Tool.Any => {
+const projectFunctionToTool = (fn: FunctionDefinition.Any): Tool.Any => {
   if (toolCache.has(fn)) {
     return toolCache.get(fn)!;
   }
 
+  // TODO(burdon): Use and map function.key?
+  // TODO(dmaretskyi): Use function key instead.
   const tool = Tool.make(makeToolName(fn.name), {
     description: fn.description,
     parameters: createStructFieldsFromSchema(fn.inputSchema),
@@ -136,8 +144,8 @@ const projectFunctionToTool = (fn: FunctionDefinition<any, any>): Tool.Any => {
  * @returns Tool name produced from function name by escaping invalid characters.
  */
 const makeToolName = (name: string) => {
-  const toolName = name.replace(/[^a-zA-Z0-9]/g, '_');
-  invariant(toolName.match(/^[a-zA-Z_][a-zA-Z0-9-_]*$/));
+  const toolName = name.toLowerCase().replace(/[^a-zA-Z0-9]/g, '-');
+  invariant(toolName.match(/^[a-z_][a-z0-9-_]*$/));
   return toolName;
 };
 
@@ -145,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:
@@ -153,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';
 };