@dxos/assistant-toolkit 0.8.4-main.72ec0f3 → 0.8.4-main.937b3ca

package/src/crud/graph.ts

@@ -11,28 +11,23 @@ import * as Option from 'effect/Option';
 import * as Schema from 'effect/Schema';
 import * as SchemaAST from 'effect/SchemaAST';
 
-import { Obj, type Relation } from '@dxos/echo';
-import { Filter, Query } from '@dxos/echo';
+import { Entity, Filter, Obj, Query, Type } from '@dxos/echo';
+import { Database } from '@dxos/echo';
 import {
-  EntityKind,
-  ObjectId,
   ReferenceAnnotationId,
   RelationSourceDXNId,
   RelationSourceId,
   RelationTargetDXNId,
   RelationTargetId,
-  create,
-  getEntityKind,
-  getSchemaDXN,
-  getSchemaTypename,
+  createObject,
   getTypeAnnotation,
   getTypeIdentifierAnnotation,
 } from '@dxos/echo/internal';
-import { type EchoDatabase, type Queue } from '@dxos/echo-db';
+import { type Queue } from '@dxos/echo-db';
 import { isEncodedReference } from '@dxos/echo-protocol';
 import { mapAst } from '@dxos/effect';
-import { ContextQueueService, DatabaseService } from '@dxos/functions';
-import { DXN } from '@dxos/keys';
+import { ContextQueueService } from '@dxos/functions';
+import { DXN, ObjectId } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { deepMapValues, isNonNullable, trim } from '@dxos/util';
 
@@ -45,7 +40,7 @@ export const Subgraph = Schema.Struct({
 export interface Subgraph extends Schema.Schema.Type<typeof Subgraph> {}
 
 export type RelatedSchema = {
-  schema: Schema.Schema.AnyNoContext;
+  schema: Type.Entity.Any;
   kind: 'reference' | 'relation';
 };
 
@@ -56,17 +51,14 @@ export type RelatedSchema = {
  * @param schema
  * @returns
  */
-export const findRelatedSchema = async (
-  db: EchoDatabase,
-  anchor: Schema.Schema.AnyNoContext,
-): Promise<RelatedSchema[]> => {
+export const findRelatedSchema = async (db: Database.Database, anchor: Type.Entity.Any): Promise<RelatedSchema[]> => {
   // TODO(dmaretskyi): Query stored schemas.
-  const allSchemas = [...db.graph.schemaRegistry.schemas];
+  const allSchemas = await db.graph.schemaRegistry.query().run();
 
   // TODO(dmaretskyi): Also do references.
   return allSchemas
     .filter((schema) => {
-      if (getTypeAnnotation(schema)?.kind !== EntityKind.Relation) {
+      if (getTypeAnnotation(schema)?.kind !== Entity.Kind.Relation) {
         return false;
       }
 
@@ -77,8 +69,8 @@ export const findRelatedSchema = async (
     })
     .map(
       (schema): RelatedSchema => ({
-        schema,
         kind: 'relation',
+        schema,
       }),
     );
 };
@@ -87,14 +79,14 @@ export const findRelatedSchema = async (
  * Non-strict DXN comparison.
  * Returns true if the DXN could be resolved to the schema.
  */
-const isSchemaAddressableByDxn = (schema: Schema.Schema.AnyNoContext, dxn: DXN): boolean => {
+const isSchemaAddressableByDxn = (schema: Type.Entity.Any, dxn: DXN): boolean => {
   if (getTypeIdentifierAnnotation(schema) === dxn.toString()) {
     return true;
   }
 
   const t = dxn.asTypeDXN();
   if (t) {
-    return t.type === getSchemaTypename(schema);
+    return t.type === Type.getTypename(schema);
   }
 
   return false;
@@ -114,13 +106,13 @@ export const LocalSearchToolkit = Toolkit.make(
     },
     success: Schema.Unknown,
     failure: Schema.Never,
-    dependencies: [DatabaseService],
+    dependencies: [Database.Service],
   }),
 );
 
 export const LocalSearchHandler = LocalSearchToolkit.toLayer({
   search_local_search: Effect.fn(function* ({ query }) {
-    const { objects } = yield* DatabaseService.runQuery(Query.select(Filter.text(query, { type: 'vector' })));
+    const objects = yield* Database.Service.runQuery(Query.select(Filter.text(query, { type: 'vector' })));
     const results = [...objects];
 
     const option = yield* Effect.serviceOption(ContextQueueService);
@@ -144,21 +136,21 @@ export const LocalSearchHandler = LocalSearchToolkit.toLayer({
 class GraphWriterSchema extends Context.Tag('@dxos/assistant/GraphWriterSchema')<
   GraphWriterSchema,
   {
-    schema: Schema.Schema.AnyNoContext[];
+    schema: Type.Entity.Any[];
   }
 >() {}
 
 /**
  * Forms typed objects that can be written to the graph database.
  */
-export const makeGraphWriterToolkit = ({ schema }: { schema: Schema.Schema.AnyNoContext[] }) => {
+export const makeGraphWriterToolkit = ({ schema }: { schema: Type.Entity.Any[] }) => {
   return Toolkit.make(
     Tool.make('graph_writer', {
       description: 'Write to the local graph database',
       parameters: createExtractionSchema(schema).fields,
       success: Schema.Unknown,
       failure: Schema.Never,
-      dependencies: [DatabaseService, ContextQueueService],
+      dependencies: [Database.Service, ContextQueueService],
     }).annotateContext(Context.make(GraphWriterSchema, { schema })),
   );
 };
@@ -178,10 +170,10 @@ export const makeGraphWriterHandler = (
 
   return toolkit.toLayer({
     graph_writer: Effect.fn(function* (input) {
-      const { db } = yield* DatabaseService;
+      const { db } = yield* Database.Service;
       const { queue } = yield* ContextQueueService;
       const data = yield* Effect.promise(() => sanitizeObjects(schema, input as any, db, queue));
-      yield* Effect.promise(() => queue.append(data as Obj.Any[]));
+      yield* Effect.promise(() => queue.append(data as Obj.Unknown[]));
 
       const dxns = data.map((obj) => Obj.getDXN(obj));
       onAppend?.(dxns);
@@ -193,31 +185,31 @@ export const makeGraphWriterHandler = (
 /**
  * Create a schema for structured data extraction.
  */
-export const createExtractionSchema = (types: Schema.Schema.AnyNoContext[]) => {
+export const createExtractionSchema = (types: Type.Entity.Any[]) => {
   return Schema.Struct({
     ...Object.fromEntries(
       types.map(preprocessSchema).map((schema, index) => [
         `objects_${getSanitizedSchemaName(types[index])}`,
         Schema.optional(Schema.Array(schema)).annotations({
-          description: `The objects of type: ${getSchemaDXN(types[index])?.asTypeDXN()!.type}. ${SchemaAST.getDescriptionAnnotation(types[index].ast).pipe(Option.getOrElse(() => ''))}`,
+          description: `The objects of type: ${Type.getDXN(types[index])?.asTypeDXN()!.type}. ${SchemaAST.getDescriptionAnnotation(types[index].ast).pipe(Option.getOrElse(() => ''))}`,
         }),
       ]),
     ),
   });
 };
 
-export const getSanitizedSchemaName = (schema: Schema.Schema.AnyNoContext) => {
-  return getSchemaDXN(schema)!
+export const getSanitizedSchemaName = (schema: Type.Entity.Any) => {
+  return Type.getDXN(schema)!
     .asTypeDXN()!
     .type.replaceAll(/[^a-zA-Z0-9]+/g, '_');
 };
 
 export const sanitizeObjects = async (
-  types: Schema.Schema.AnyNoContext[],
+  types: Type.Entity.Any[],
   data: Record<string, readonly unknown[]>,
-  db: EchoDatabase,
+  db: Database.Database,
   queue?: Queue,
-): Promise<Obj.Any[]> => {
+): Promise<Obj.Unknown[]> => {
   const entries = types
     .map(
       (type) =>
@@ -230,7 +222,7 @@ export const sanitizeObjects = async (
 
   const idMap = new Map<string, string>();
   const existingIds = new Set<ObjectId>();
-  const enitties = new Map<ObjectId, Obj.Any | Relation.Any>();
+  const enitties = new Map<ObjectId, Entity.Unknown>();
 
   const resolveId = (id: string): DXN | undefined => {
     if (ObjectId.isValid(id)) {
@@ -275,7 +267,7 @@ export const sanitizeObjects = async (
         return recurse(value);
       });
 
-      if (getEntityKind(entry.schema) === 'relation') {
+      if (Entity.getKind(entry.schema) === 'relation') {
         const sourceDxn = resolveId(data.source);
         if (!sourceDxn) {
           log.warn('source not found', { source: data.source });
@@ -298,7 +290,7 @@ export const sanitizeObjects = async (
     .filter((object) => !existingIds.has(object.data.id)); // TODO(dmaretskyi): This dissallows updating existing objects.
 
   // TODO(dmaretskyi): Use ref resolver.
-  const { objects: dbObjects } = await db.query(Query.select(Filter.ids(...existingIds))).run();
+  const dbObjects = await db.query(Query.select(Filter.id(...existingIds))).run();
   const queueObjects = (await queue?.getObjectsById([...existingIds])) ?? [];
   const objects = [...dbObjects, ...queueObjects].filter(isNonNullable);
 
@@ -332,7 +324,7 @@ export const sanitizeObjects = async (
       }
     }
     if (!skip) {
-      const obj = create(schema, data);
+      const obj = createObject(schema, data);
       enitties.set(obj.id, obj);
       return [obj];
     }
@@ -347,7 +339,7 @@ const SoftRef = Schema.Struct({
 });
 
 const preprocessSchema = (schema: Schema.Schema.AnyNoContext) => {
-  const isRelationSchema = getEntityKind(schema) === 'relation';
+  const isRelationSchema = Entity.getKind(schema) === 'relation';
 
   const go = (ast: SchemaAST.AST, visited = new Set<SchemaAST.AST>()): SchemaAST.AST => {
     if (visited.has(ast)) {

package/src/experimental/feed.test.ts

@@ -13,7 +13,7 @@ import * as Redacted from 'effect/Redacted';
 import { AiService } from '@dxos/ai';
 import { AiServiceTestingPreset, EXA_API_KEY } from '@dxos/ai/testing';
 import { makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '@dxos/assistant';
-import { TestHelpers } from '@dxos/effect';
+import { TestHelpers } from '@dxos/effect/testing';
 import { CredentialsService, FunctionInvocationService } from '@dxos/functions';
 import { TracingServiceExt } from '@dxos/functions-runtime';
 import {

package/src/functions/agent/prompt.ts

@@ -13,7 +13,8 @@ import { AiService, ConsolePrinter, ModelName } from '@dxos/ai';
 import { AiSession, GenerationObserver, createToolkit } from '@dxos/assistant';
 import { Prompt, Template } from '@dxos/blueprints';
 import { Obj, Ref, Type } from '@dxos/echo';
-import { DatabaseService, TracingService, defineFunction } from '@dxos/functions';
+import { Database } from '@dxos/echo';
+import { TracingService, defineFunction } from '@dxos/functions';
 import { log } from '@dxos/log';
 
 const DEFAULT_MODEL: ModelName = '@anthropic/claude-opus-4-0';
@@ -45,7 +46,7 @@ export default defineFunction({
     // for (const key of Object.keys(data.input)) {
     //   const value = data.input[key];
     //   if (Ref.isRef(value)) {
-    //     const object = yield* DatabaseService.load(value);
+    //     const object = yield* Database.Service.load(value);
     //     input[key] = Obj.toJSON(object);
     //   } else {
     //     input[key] = JSON.stringify(value);
@@ -55,16 +56,16 @@ export default defineFunction({
       Match.when(
         (value: any) => Ref.isRef(value),
         Effect.fnUntraced(function* (ref) {
-          const object = yield* DatabaseService.load(ref);
-          return Obj.toJSON(object as Obj.Any);
+          const object = yield* Database.Service.load(ref);
+          return Obj.toJSON(object as Obj.Unknown);
         }),
       ),
       Match.orElse(() => Effect.succeed(data.input)),
     );
 
-    yield* DatabaseService.flush({ indexes: true });
-    const prompt = yield* DatabaseService.load(data.prompt);
-    const systemPrompt = data.systemPrompt ? yield* DatabaseService.load(data.systemPrompt) : undefined;
+    yield* Database.Service.flush({ indexes: true });
+    const prompt = yield* Database.Service.load(data.prompt);
+    const systemPrompt = data.systemPrompt ? yield* Database.Service.load(data.systemPrompt) : undefined;
     yield* TracingService.emitStatus({ message: `Running ${prompt.id}` });
 
     log.info('starting agent', { prompt: prompt.id, input });
@@ -72,23 +73,26 @@ export default defineFunction({
     const blueprints = yield* Function.pipe(
       prompt.blueprints,
       Array.appendAll(systemPrompt?.blueprints ?? []),
-      Effect.forEach(DatabaseService.loadOption),
+      Effect.forEach(Database.Service.loadOption),
       Effect.map(Array.filter(Option.isSome)),
       Effect.map(Array.map((option) => option.value)),
     );
+    const toolkit = yield* createToolkit({ blueprints });
+
     const objects = yield* Function.pipe(
       prompt.context,
       Array.appendAll(systemPrompt?.context ?? []),
-      Effect.forEach(DatabaseService.loadOption),
+      Effect.forEach(Database.Service.loadOption),
       Effect.map(Array.filter(Option.isSome)),
       Effect.map(Array.map((option) => option.value)),
     );
-    const toolkit = yield* createToolkit({ blueprints });
 
-    const promptInstructions = yield* DatabaseService.load(prompt.instructions.source);
+    const promptInstructions = yield* Database.Service.load(prompt.instructions.source);
     const promptText = Template.process(promptInstructions.content, input);
 
-    const systemInstructions = systemPrompt ? yield* DatabaseService.load(systemPrompt.instructions.source) : undefined;
+    const systemInstructions = systemPrompt
+      ? yield* Database.Service.load(systemPrompt.instructions.source)
+      : undefined;
     const systemText = systemInstructions ? Template.process(systemInstructions.content, {}) : undefined;
 
     const session = new AiSession();
@@ -97,7 +101,7 @@ export default defineFunction({
         prompt: promptText,
         system: systemText,
         blueprints,
-        objects: objects as Obj.Any[],
+        objects: objects as Obj.Unknown[],
         toolkit,
         observer: GenerationObserver.fromPrinter(new ConsolePrinter({ tag: 'agent' })),
       })

package/src/functions/discord/fetch-messages.test.ts

@@ -12,7 +12,7 @@ import * as Layer from 'effect/Layer';
 import { AiService } from '@dxos/ai';
 import { AiServiceTestingPreset } from '@dxos/ai/testing';
 import { makeToolExecutionServiceFromFunctions, makeToolResolverFromFunctions } from '@dxos/assistant';
-import { TestHelpers } from '@dxos/effect';
+import { TestHelpers } from '@dxos/effect/testing';
 import { CredentialsService, FunctionInvocationService, TracingService } from '@dxos/functions';
 import { FunctionInvocationServiceLayerTestMocked, TestDatabaseLayer } from '@dxos/functions-runtime/testing';
 

package/src/functions/discord/fetch-messages.ts

@@ -133,7 +133,7 @@ export default defineFunction({
         throw new Error('no channels found');
       }
       for (const channel of channels) {
-        console.log(channel.id, 'name' in channel ? channel.name : undefined);
+        log.info('channel', { id: channel.id, name: 'name' in channel ? channel.name : undefined });
       }
 
       yield* TracingService.emitStatus({ message: `Will fetch from channels: ${channels.length}` });

package/src/functions/document/read.ts

@@ -6,7 +6,8 @@ import * as Effect from 'effect/Effect';
 import * as Schema from 'effect/Schema';
 
 import { ArtifactId } from '@dxos/assistant';
-import { DatabaseService, defineFunction } from '@dxos/functions';
+import { Database } from '@dxos/echo';
+import { defineFunction } from '@dxos/functions';
 import { Markdown } from '@dxos/plugin-markdown/types';
 
 export default defineFunction({
@@ -22,8 +23,8 @@ export default defineFunction({
     content: Schema.String,
   }),
   handler: Effect.fn(function* ({ data: { id } }) {
-    const doc = yield* DatabaseService.resolve(ArtifactId.toDXN(id), Markdown.Document);
-    const { content } = yield* DatabaseService.load(doc.content);
+    const doc = yield* Database.Service.resolve(ArtifactId.toDXN(id), Markdown.Document);
+    const { content } = yield* Database.Service.load(doc.content);
     return { content };
   }),
 });

package/src/functions/document/update.ts

@@ -6,7 +6,8 @@ import * as Effect from 'effect/Effect';
 import * as Schema from 'effect/Schema';
 
 import { ArtifactId } from '@dxos/assistant';
-import { DatabaseService, defineFunction } from '@dxos/functions';
+import { Database, Obj } from '@dxos/echo';
+import { defineFunction } from '@dxos/functions';
 import { Markdown } from '@dxos/plugin-markdown/types';
 
 export default defineFunction({
@@ -23,8 +24,10 @@ export default defineFunction({
   }),
   outputSchema: Schema.Void,
   handler: Effect.fn(function* ({ data: { id, content } }) {
-    const doc = yield* DatabaseService.resolve(ArtifactId.toDXN(id), Markdown.Document);
-    const text = yield* DatabaseService.load(doc.content);
-    text.content = content;
+    const doc = yield* Database.Service.resolve(ArtifactId.toDXN(id), Markdown.Document);
+    const text = yield* Database.Service.load(doc.content);
+    Obj.change(text, (t) => {
+      t.content = content;
+    });
   }),
 });