@dxos/assistant-toolkit 0.8.4-main.66e292d → 0.8.4-main.6fa680abb7

package/src/crud/graph.ts

@@ -12,6 +12,7 @@ import * as Schema from 'effect/Schema';
 import * as SchemaAST from 'effect/SchemaAST';
 
 import { Entity, Filter, Obj, Query, Type } from '@dxos/echo';
+import { Database } from '@dxos/echo';
 import {
   ReferenceAnnotationId,
   RelationSourceDXNId,
@@ -22,10 +23,10 @@ import {
   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 { ContextQueueService } from '@dxos/functions';
 import { DXN, ObjectId } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { deepMapValues, isNonNullable, trim } from '@dxos/util';
@@ -39,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.AnyEntity;
   kind: 'reference' | 'relation';
 };
 
@@ -50,12 +51,9 @@ 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.AnyEntity): 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
@@ -71,8 +69,8 @@ export const findRelatedSchema = async (
     })
     .map(
       (schema): RelatedSchema => ({
-        schema,
         kind: 'relation',
+        schema,
       }),
     );
 };
@@ -81,7 +79,7 @@ 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.AnyEntity, dxn: DXN): boolean => {
   if (getTypeIdentifierAnnotation(schema) === dxn.toString()) {
     return true;
   }
@@ -108,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.runQuery(Query.select(Filter.text(query, { type: 'vector' })));
     const results = [...objects];
 
     const option = yield* Effect.serviceOption(ContextQueueService);
@@ -138,21 +136,21 @@ export const LocalSearchHandler = LocalSearchToolkit.toLayer({
 class GraphWriterSchema extends Context.Tag('@dxos/assistant/GraphWriterSchema')<
   GraphWriterSchema,
   {
-    schema: Schema.Schema.AnyNoContext[];
+    schema: Type.AnyEntity[];
   }
 >() {}
 
 /**
  * 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.AnyEntity[] }) => {
   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 })),
   );
 };
@@ -172,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);
@@ -187,7 +185,7 @@ export const makeGraphWriterHandler = (
 /**
  * Create a schema for structured data extraction.
  */
-export const createExtractionSchema = (types: Schema.Schema.AnyNoContext[]) => {
+export const createExtractionSchema = (types: Type.AnyEntity[]) => {
   return Schema.Struct({
     ...Object.fromEntries(
       types.map(preprocessSchema).map((schema, index) => [
@@ -200,18 +198,18 @@ export const createExtractionSchema = (types: Schema.Schema.AnyNoContext[]) => {
   });
 };
 
-export const getSanitizedSchemaName = (schema: Schema.Schema.AnyNoContext) => {
+export const getSanitizedSchemaName = (schema: Type.AnyEntity) => {
   return Type.getDXN(schema)!
     .asTypeDXN()!
     .type.replaceAll(/[^a-zA-Z0-9]+/g, '_');
 };
 
 export const sanitizeObjects = async (
-  types: Schema.Schema.AnyNoContext[],
+  types: Type.AnyEntity[],
   data: Record<string, readonly unknown[]>,
-  db: EchoDatabase,
+  db: Database.Database,
   queue?: Queue,
-): Promise<Obj.Any[]> => {
+): Promise<Obj.Unknown[]> => {
   const entries = types
     .map(
       (type) =>
@@ -292,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);
 

package/src/functions/READ_THIS.md

@@ -0,0 +1,5 @@
+###
+
+Can we move all functions into the corresponding blueprint folder?
+
+###

package/src/functions/agent/index.ts

@@ -2,10 +2,8 @@
 // Copyright 2025 DXOS.org
 //
 
-import { type FunctionDefinition } from '@dxos/functions';
+import Prompt from './prompt';
 
-import { default as prompt$ } from './prompt';
-
-export namespace Agent {
-  export const prompt: FunctionDefinition.Any = prompt$; // TODO(burdon): Temp fix for TS error.
-}
+export const AgentFunctions = {
+  Prompt,
+};

package/src/functions/agent/prompt.ts

@@ -5,26 +5,25 @@
 import * as Array from 'effect/Array';
 import * as Effect from 'effect/Effect';
 import * as Function from 'effect/Function';
-import * as Match from 'effect/Match';
 import * as Option from 'effect/Option';
 import * as Schema from 'effect/Schema';
 
 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, Obj, Ref } from '@dxos/echo';
+import { TracingService, defineFunction } from '@dxos/functions';
 import { log } from '@dxos/log';
 
 const DEFAULT_MODEL: ModelName = '@anthropic/claude-opus-4-0';
 
 export default defineFunction({
-  key: 'dxos.org/function/prompt',
+  key: 'org.dxos.function.prompt',
   name: 'Agent',
   description: 'Agentic worker that executes a provided prompt using blueprints and tools.',
   inputSchema: Schema.Struct({
-    prompt: Type.Ref(Prompt.Prompt),
-    systemPrompt: Type.Ref(Prompt.Prompt).pipe(Schema.optional),
+    prompt: Ref.Ref(Prompt.Prompt),
+    systemPrompt: Schema.optional(Ref.Ref(Prompt.Prompt)),
     /**
      * @default @anthropic/claude-opus-4-0
      */
@@ -45,26 +44,20 @@ 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.load(value);
     //     input[key] = Obj.toJSON(object);
     //   } else {
     //     input[key] = JSON.stringify(value);
     //   }
     // }
-    const input = yield* Match.value(data.input).pipe(
-      Match.when(
-        (value: any) => Ref.isRef(value),
-        Effect.fnUntraced(function* (ref) {
-          const object = yield* DatabaseService.load(ref);
-          return Obj.toJSON(object as Obj.Any);
-        }),
-      ),
-      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;
+    const input = yield* Ref.isRef(data.input)
+      ? Database.load(data.input).pipe(Effect.map(Obj.toJSON))
+      : Effect.succeed(data.input);
+
+    yield* Database.flush({ indexes: true });
+    const prompt = yield* Database.load(data.prompt);
+    const systemPrompt = data.systemPrompt ? yield* Database.load(data.systemPrompt) : undefined;
     yield* TracingService.emitStatus({ message: `Running ${prompt.id}` });
 
     log.info('starting agent', { prompt: prompt.id, input });
@@ -72,23 +65,24 @@ export default defineFunction({
     const blueprints = yield* Function.pipe(
       prompt.blueprints,
       Array.appendAll(systemPrompt?.blueprints ?? []),
-      Effect.forEach(DatabaseService.loadOption),
+      Effect.forEach(Database.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.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.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.load(systemPrompt.instructions.source) : undefined;
     const systemText = systemInstructions ? Template.process(systemInstructions.content, {}) : undefined;
 
     const session = new AiSession();
@@ -97,7 +91,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' })),
       })