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

package/src/components/Toolbox/Toolbox.stories.tsx

@@ -3,16 +3,16 @@
 //
 
 import { type Meta, type StoryObj } from '@storybook/react-vite';
+import * as Effect from 'effect/Effect';
 import React from 'react';
 
-import { IntentPlugin, SettingsPlugin } from '@dxos/app-framework';
 import { withPluginManager } from '@dxos/app-framework/testing';
-import { capabilities } from '@dxos/assistant-toolkit';
+import { capabilities } from '@dxos/assistant-toolkit/testing';
 import { ChessPlugin } from '@dxos/plugin-chess';
 import { ClientPlugin } from '@dxos/plugin-client';
 import { MapPlugin } from '@dxos/plugin-map';
-import { SpacePlugin } from '@dxos/plugin-space';
 import { TablePlugin } from '@dxos/plugin-table';
+import { corePlugins } from '@dxos/plugin-testing';
 import { withTheme } from '@dxos/react-ui/testing';
 
 import { translations } from '../../translations';
@@ -20,7 +20,7 @@ import { translations } from '../../translations';
 import { Toolbox, type ToolboxProps } from './Toolbox';
 
 const DefaultStory = (props: ToolboxProps) => {
-  return <Toolbox {...props} classNames='w-[30rem] h-[15rem] rounded-sm border border-separator py-1' />;
+  return <Toolbox {...props} classNames='is-[30rem] bs-[15rem] rounded-sm border border-separator plb-1' />;
 };
 
 const meta = {
@@ -31,14 +31,13 @@ const meta = {
     withTheme,
     withPluginManager({
       plugins: [
+        ...corePlugins(),
         ClientPlugin({
-          onClientInitialized: async ({ client }) => {
-            await client.halo.createIdentity();
-          },
+          onClientInitialized: ({ client }) =>
+            Effect.gen(function* () {
+              yield* Effect.promise(() => client.halo.createIdentity());
+            }),
         }),
-        SpacePlugin({}),
-        SettingsPlugin(),
-        IntentPlugin(),
         ChessPlugin(),
         MapPlugin(),
         TablePlugin(),

package/src/components/Toolbox/Toolbox.tsx

@@ -5,13 +5,12 @@
 import React, { type FC, Fragment, useEffect, useState } from 'react';
 
 import { type Blueprint } from '@dxos/blueprints';
-import { type Ref } from '@dxos/echo';
+import { type Database, type Ref } from '@dxos/echo';
 import { Function } from '@dxos/functions';
 import { log } from '@dxos/log';
-import { Filter, type Space, useQuery } from '@dxos/react-client/echo';
-import { type ThemedClassName } from '@dxos/react-ui';
-import { useTranslation } from '@dxos/react-ui';
-import { mx } from '@dxos/react-ui-theme';
+import { Filter, useQuery } from '@dxos/react-client/echo';
+import { type ThemedClassName, useTranslation } from '@dxos/react-ui';
+import { mx } from '@dxos/ui-theme';
 
 import { meta } from '../../meta';
 import { type AiChatProcessor } from '../../processor';
@@ -21,7 +20,7 @@ export type ToolboxProps = ThemedClassName<{
   services?: { service: ServiceType }[];
   functions?: Function.Function[];
   // TODO(burdon): Combine into single array.
-  blueprints?: readonly Ref.Ref<Blueprint.Blueprint>[];
+  blueprints?: readonly Blueprint.Blueprint[];
   activeBlueprints?: readonly Ref.Ref<Blueprint.Blueprint>[];
 }>;
 
@@ -33,10 +32,10 @@ export const Toolbox = ({ classNames, functions, services, blueprints, activeBlu
       {blueprints && blueprints.length > 0 && (
         <Section
           title='Blueprints'
-          items={blueprints.map(({ target }) => ({
-            name: target?.name ?? '',
-            description: target?.description ?? '',
-            subitems: target?.tools.map((toolId) => ({ name: `∙ ${safeToolId(toolId)}` })),
+          items={blueprints.map(({ name, description, tools }) => ({
+            name,
+            description,
+            subitems: tools.map((toolId) => ({ name: `∙ ${safeToolId(toolId)}` })),
           }))}
         />
       )}
@@ -77,11 +76,11 @@ const Section: FC<{
 }> = ({ title, items, striped }) => {
   const stripeClassNames = 'odd:bg-neutral-50 dark:odd:bg-neutral-800';
   const gridClassNames = 'grid grid-cols-[8rem_1fr]';
-  const subGridClassNames = mx('col-span-full grid grid-cols-subgrid text-xs px-2', striped && stripeClassNames);
+  const subGridClassNames = mx('col-span-full grid grid-cols-subgrid text-xs pli-2', striped && stripeClassNames);
 
   return (
     <div>
-      <h1 className='px-2 text-sm'>{title}</h1>
+      <h1 className='pli-2 text-sm'>{title}</h1>
       <div className={gridClassNames}>
         {items.map(({ name, description, subitems }, i) => (
           <Fragment key={i}>
@@ -105,13 +104,13 @@ const Section: FC<{
 };
 
 export type ToolboxContainerProps = ThemedClassName<{
-  space?: Space;
+  db?: Database.Database;
   processor?: AiChatProcessor;
 }>;
 
-export const ToolboxContainer = ({ classNames, space, processor }: ToolboxContainerProps) => {
+export const ToolboxContainer = ({ classNames, db, processor }: ToolboxContainerProps) => {
   // Registered services.
-  const services = useQuery(space, Filter.type(ServiceType));
+  const services = useQuery(db, Filter.type(ServiceType));
   const [serviceTools, setServiceTools] = useState<{ service: ServiceType }[]>([]);
   useEffect(() => {
     log('creating service tools...', { services: services.length });
@@ -125,12 +124,12 @@ export const ToolboxContainer = ({ classNames, space, processor }: ToolboxContai
   }, [services]);
 
   // Deployed functions.
-  const functions = useQuery(space, Filter.type(Function.Function));
+  const functions = useQuery(db, Filter.type(Function.Function));
 
   return (
     <Toolbox
       classNames={classNames}
-      blueprints={processor?.context.blueprints.value}
+      blueprints={processor?.context.getBlueprints()}
       services={serviceTools}
       functions={functions}
     />

package/src/components/Typewriter/AssistantToolbar.tsx

@@ -0,0 +1,161 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import { type EditorView } from '@codemirror/view';
+import * as LanguageModel from '@effect/ai/LanguageModel';
+import * as Effect from 'effect/Effect';
+import type * as Runtime from 'effect/Runtime';
+import React, { useState } from 'react';
+
+import { runInRuntime } from '@dxos/effect';
+import { log } from '@dxos/log';
+import { Button, Icon, Input, ThemeProvider } from '@dxos/react-ui';
+
+export interface AssistantToolbarProps {
+  view: EditorView;
+  from: number;
+  to: number;
+  runtime: Runtime.Runtime<LanguageModel.LanguageModel>;
+}
+
+export const AssistantToolbar = ({ view, from, to, runtime }: AssistantToolbarProps) => {
+  const [isProcessing, setIsProcessing] = useState(false);
+  const [showPromptInput, setShowPromptInput] = useState(false);
+  const [customPrompt, setCustomPrompt] = useState('');
+
+  const runEffect = async (text: string, prompt: string) => {
+    setIsProcessing(true);
+    try {
+      const newText = await runInRuntime(
+        runtime,
+        Effect.gen(function* () {
+          // Construct meaningful system prompt context
+          const response = yield* LanguageModel.generateText({
+            prompt: `${prompt}:\n"${text}"`,
+          });
+          return response.text;
+        }),
+      );
+
+      // Apply change - check if selection is still valid/same?
+      // For simple demo, we just apply to the passed from/to, but ideally we check current selection.
+      // We can just rely on from/to passed as props, assuming the toolbar unmounts if selection changes.
+      view.dispatch({
+        changes: { from, to, insert: newText },
+      });
+    } catch (e) {
+      log.error('Assistant action failed', { error: e });
+    } finally {
+      setIsProcessing(false);
+      setShowPromptInput(false);
+    }
+  };
+
+  const currentText = view.state.sliceDoc(from, to);
+
+  const handleRegenerate = () => {
+    void runEffect(
+      currentText,
+      'You are a typewriting agent. Agent only outputs the result; one version; no explanation or formatting. Re-write the following text to be more clear and concise, keep the same meaning.',
+    );
+  };
+
+  const handleCasual = () => handleChangeTone('casual');
+  const handleFormal = () => handleChangeTone('formal');
+
+  const handleChangeTone = (tone: 'formal' | 'casual') => {
+    void runEffect(
+      currentText,
+      `You are a typewriting agent. Agent only outputs the result; one version; no explanation or formatting. Rewrite the following text to be more ${tone}`,
+    );
+  };
+
+  const handleEditWithPrompt = () => {
+    if (!customPrompt) return;
+    void runEffect(currentText, customPrompt);
+  };
+
+  return (
+    <ThemeProvider>
+      {showPromptInput ? (
+        <div
+          className='flex items-center gap-2 p-1 bg-white/90 dark:bg-neutral-800/90 backdrop-blur-md rounded-full shadow-xl border border-neutral-200 dark:border-neutral-700 animate-in fade-in zoom-in-95 duration-150 mt-2'
+          onMouseDown={(e) => e.preventDefault()}
+        >
+          <Input.Root>
+            <Input.TextInput
+              placeholder='Instruction...'
+              value={customPrompt}
+              onChange={(e) => setCustomPrompt(e.target.value)}
+              onKeyDown={(e) => {
+                if (e.key === 'Enter') handleEditWithPrompt();
+                if (e.key === 'Escape') setShowPromptInput(false);
+                e.stopPropagation();
+              }}
+              autoFocus
+              classNames='border-none bg-transparent focus-visible:ring-0 p-1 text-sm min-w-[200px]'
+            />
+          </Input.Root>
+          <Button
+            variant='ghost'
+            classNames='p-1 min-bs-0 bs-6 is-6 rounded-full'
+            onClick={handleEditWithPrompt}
+            disabled={!customPrompt || isProcessing}
+          >
+            <Icon icon='ph--arrows-clockwise--regular' classNames={isProcessing ? 'animate-spin' : ''} />
+          </Button>
+        </div>
+      ) : (
+        <div
+          className='flex items-center gap-1 p-1 bg-white/50 dark:bg-neutral-800/50 backdrop-blur-md rounded-xl shadow-xl border border-neutral-200 dark:border-neutral-700 animate-in fade-in zoom-in-95 duration-150 mt-2'
+          onMouseDown={(e) => e.preventDefault()}
+        >
+          <div className='flex items-center gap-1'>
+            <Button
+              variant='ghost'
+              classNames='flex flex-col items-center gap-1 p-1 h-auto w-auto rounded-lg hover:bg-neutral-100 dark:hover:bg-neutral-700'
+              onClick={handleCasual}
+              disabled={isProcessing}
+            >
+              <Icon icon='ph--caret-left--regular' size={4} />
+              <span className='text-[10px] leading-none text-neutral-500 font-medium'>Casual</span>
+            </Button>
+
+            <Button
+              variant='ghost'
+              classNames='flex flex-col items-center gap-1 p-1 h-auto w-auto rounded-lg hover:bg-neutral-100 dark:hover:bg-neutral-700'
+              onClick={handleRegenerate}
+              disabled={isProcessing}
+            >
+              <Icon icon='ph--arrows-clockwise--regular' size={4} classNames={isProcessing ? 'animate-spin' : ''} />
+              <span className='text-[10px] leading-none text-neutral-500 font-medium'>Different</span>
+            </Button>
+
+            <Button
+              variant='ghost'
+              classNames='flex flex-col gap-1 p-1 h-auto w-auto rounded-lg hover:bg-neutral-100 dark:hover:bg-neutral-700'
+              onClick={handleFormal}
+              disabled={isProcessing}
+            >
+              <Icon icon='ph--caret-right--regular' size={4} />
+              <span className='text-[10px] leading-none text-neutral-500 font-medium'>Formal</span>
+            </Button>
+          </div>
+
+          <div className='w-px h-8 bg-neutral-200 dark:bg-neutral-700 mx-1' />
+
+          <Button
+            variant='ghost'
+            classNames='flex flex-col gap-1 p-1 h-auto w-auto rounded-lg hover:bg-neutral-100 dark:hover:bg-neutral-700'
+            onClick={() => setShowPromptInput(true)}
+            disabled={isProcessing}
+          >
+            <Icon icon='ph--keyboard--regular' size={4} />
+            <span className='text-[10px] leading-none text-neutral-500 font-medium'>Instruct</span>
+          </Button>
+        </div>
+      )}
+    </ThemeProvider>
+  );
+};

package/src/components/Typewriter/Typewriter.stories.tsx

@@ -0,0 +1,83 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import type * as LanguageModel from '@effect/ai/LanguageModel';
+import type { Meta, StoryObj } from '@storybook/react-vite';
+import * as Layer from 'effect/Layer';
+import * as ManagedRuntime from 'effect/ManagedRuntime';
+import type * as Runtime from 'effect/Runtime';
+import React, { useEffect, useState } from 'react';
+
+import { AiService } from '@dxos/ai';
+import { AiServiceTestingPreset } from '@dxos/ai/testing';
+
+import { Typewriter } from './Typewriter';
+
+const meta: Meta<typeof Typewriter> = {
+  title: 'plugins/plugin-assistant/Typewriter',
+  component: Typewriter,
+  parameters: {
+    layout: 'centered',
+  },
+};
+
+export default meta;
+
+type Story = StoryObj<typeof Typewriter>;
+
+const sampleText = `# The Future of Artificial Intelligence
+
+Artificial Intelligence (AI) has rapidly evolved from a theoretical concept to a transformative force reshaping industries, economies, and daily life. At its core, AI refers to the development of computer systems capable of performing tasks that typically require human intelligence, such as visual perception, speech recognition, decision-making, and language translation.
+
+## Machine Learning and Neural Networks
+
+The resurgence of AI in recent years is largely driven by advancements in machine learning (ML) and deep learning. Neural networks, inspired by the human brain's architecture, enable computers to learn from vast amounts of data. This has led to breakthroughs in areas like natural language processing (NLP), allowing models like GPT-4 to generate human-like text and code.
+
+## Ethical Considerations
+
+As AI becomes more autonomous, ethical concerns have moved to the forefront. Issues such as algorithmic bias, data privacy, and the potential displacement of jobs are critical topics of debate. Ensuring that AI systems are fair, transparent, and aligned with human values is essential for their sustainable development.
+
+## Conclusion
+
+The journey of AI is just beginning. With responsible innovation and global collaboration, AI has the potential to solve some of humanity's most pressing challenges, from climate change to personalized medicine.
+
+(try selecting some text to see the AI assistant, or wait for the proofreader to find the typo below)
+
+The quick brown fox jumps over teh lazy dog. i think this is cool.
+`;
+
+export const Default: Story = {
+  render: () => {
+    const [rtEffect, setRtEffect] = useState<Runtime.Runtime<LanguageModel.LanguageModel>>();
+    useEffect(() => {
+      let disposed = false;
+      const rt = ManagedRuntime.make(
+        AiService.model('@anthropic/claude-haiku-4-5').pipe(
+          Layer.provide(AiServiceTestingPreset('edge-remote')),
+          Layer.orDie,
+        ),
+      );
+      queueMicrotask(async () => {
+        if (disposed) return;
+        const x = await rt.runtime();
+        if (disposed) return;
+        setRtEffect(x);
+      });
+      return () => {
+        disposed = true;
+        void rt.dispose();
+      };
+    }, []);
+
+    if (!rtEffect) {
+      return <div>...Loading</div>;
+    }
+
+    return (
+      <div className='w-[800px] p-4'>
+        <Typewriter initialContent={sampleText} runtime={rtEffect} />
+      </div>
+    );
+  },
+};

package/src/components/Typewriter/Typewriter.tsx

@@ -0,0 +1,54 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import { type Extension } from '@codemirror/state';
+import type * as LanguageModel from '@effect/ai/LanguageModel';
+import type * as Runtime from 'effect/Runtime';
+import React, { useMemo, useRef } from 'react';
+
+import { DensityProvider, ThemeProvider } from '@dxos/react-ui';
+import { Editor, type EditorController } from '@dxos/react-ui-editor';
+import {
+  createBasicExtensions,
+  createMarkdownExtensions,
+  createThemeExtensions,
+  decorateMarkdown,
+} from '@dxos/ui-editor';
+
+import { assistant } from './assistant-extension';
+
+export type TypewriterProps = {
+  initialContent?: string;
+  extensions?: Extension[];
+  runtime: Runtime.Runtime<LanguageModel.LanguageModel>;
+};
+
+export const Typewriter = ({ initialContent = '', extensions = [], runtime }: TypewriterProps) => {
+  const editorRef = useRef<EditorController>(null);
+  const extension = useMemo(
+    () => [
+      createBasicExtensions({ scrollPastEnd: true, search: true }),
+      createThemeExtensions({ syntaxHighlighting: true }),
+      createMarkdownExtensions(),
+      decorateMarkdown(),
+      assistant(runtime),
+      ...extensions,
+    ],
+    [runtime, extensions],
+  );
+
+  return (
+    <ThemeProvider>
+      <DensityProvider density='fine'>
+        <div className='is-full bs-full grid overflow-hidden'>
+          <Editor.Root ref={editorRef}>
+            <Editor.Viewport>
+              <Editor.Content initialValue={initialContent} extensions={extension} />
+            </Editor.Viewport>
+          </Editor.Root>
+        </div>
+      </DensityProvider>
+    </ThemeProvider>
+  );
+};

package/src/components/Typewriter/assistant-extension.tsx

@@ -0,0 +1,144 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+import { type Diagnostic, linter } from '@codemirror/lint';
+import { type Extension, StateField } from '@codemirror/state';
+import { EditorView, type Tooltip, showTooltip } from '@codemirror/view';
+import * as LanguageModel from '@effect/ai/LanguageModel';
+import * as Effect from 'effect/Effect';
+import * as Runtime from 'effect/Runtime';
+import React from 'react';
+import { createRoot } from 'react-dom/client';
+
+import { unwrapExit } from '@dxos/effect';
+import { log } from '@dxos/log';
+
+import { AssistantToolbar } from './AssistantToolbar';
+
+export const assistant = (runtime: Runtime.Runtime<LanguageModel.LanguageModel>): Extension[] => {
+  return [
+    assistantState(runtime),
+    assistantLinter(runtime),
+    EditorView.baseTheme({
+      '.cm-tooltip-assistant': {
+        backgroundColor: 'transparent !important',
+        border: 'none !important',
+        padding: '0 !important',
+      },
+    }),
+  ];
+};
+
+//
+// Toolbar State & Tooltip
+//
+
+const assistantState = (runtime: Runtime.Runtime<LanguageModel.LanguageModel>) =>
+  StateField.define<Tooltip | null>({
+    create: () => null,
+    update: (tooltip, tr): Tooltip | null => {
+      if (tr.selection || tr.docChanged) {
+        const { from, to } = tr.state.selection.main;
+        if (from === to) {
+          return null;
+        } // No selection
+
+        // Check if we should show tooltip
+        // We return a tooltip object
+        return {
+          pos: from,
+          above: false,
+          strictSide: true,
+          create: (view) => {
+            const dom = document.createElement('div');
+            dom.classList.add('cm-tooltip-assistant');
+            const root = createRoot(dom);
+
+            root.render(React.createElement(AssistantToolbar, { view, runtime, from, to }));
+
+            return {
+              dom,
+              destroy: () => {
+                // Delay unmount to allow animations if needed, but for now just cleanup
+                setTimeout(() => root.unmount(), 0);
+              },
+            };
+          },
+        };
+      }
+      return tooltip;
+    },
+    provide: (f) => showTooltip.from(f),
+  });
+
+//
+// Linter
+//
+
+const assistantLinter = (runtime: Runtime.Runtime<LanguageModel.LanguageModel>) =>
+  linter(async (view) => {
+    const text = view.state.doc.toString();
+    if (text.length < 10) {
+      return [];
+    }
+
+    // Run proofreading effect
+    const prompt = `Proofread the following text. Identify typos and simple grammar mistakes. Return ONLY a valid JSON array of objects with fields: "original" (string), "replacement" (string), "context" (string, 3-5 words around match). Text: \n\n${text}`;
+
+    // We mocked the model behavior in the story, but here we assume the model returns a string.
+    // In a real implementation we would parse the JSON.
+    // For this refactor, we just wire it up.
+
+    // NOTE: This runs on every doc view update that triggers linter. Linter is debounced by default (750ms).
+
+    try {
+      const result = unwrapExit(
+        await Effect.gen(function* () {
+          const generated = yield* LanguageModel.generateText({ prompt });
+          return generated.text;
+        }).pipe(Runtime.runPromiseExit(runtime)),
+      );
+
+      // Parse result - assuming simpler mock response format for now or we build a robust parser
+      // For the purpose of this task (UX demo), we can assume the mock returns a specific format or we parse it.
+      // Let's assume the mock returns a JSON string.
+
+      let suggestions: any[] = [];
+      try {
+        // Find JSON in response
+        const match = result.match(/\[.*\]/s);
+        if (match) {
+          suggestions = JSON.parse(match[0]);
+        }
+      } catch (e) {
+        // failed to parse
+      }
+
+      const diagnostics: Diagnostic[] = [];
+      for (const s of suggestions) {
+        // Naive finding of position
+        const idx = text.indexOf(s.original);
+        if (idx !== -1) {
+          diagnostics.push({
+            from: idx,
+            to: idx + s.original.length,
+            severity: 'info',
+            message: `Suggestion: ${s.replacement}`,
+            actions: [
+              {
+                name: 'Apply',
+                apply: (view, from, to) => {
+                  view.dispatch({ changes: { from, to, insert: s.replacement } });
+                },
+              },
+            ],
+          });
+        }
+      }
+      return diagnostics;
+    } catch (e) {
+      log.error('Proofreading failed', { error: e });
+      return [];
+    }
+  });

package/src/components/Typewriter/index.ts

@@ -0,0 +1,5 @@
+//
+// Copyright 2026 DXOS.org
+//
+
+export * from './Typewriter';

package/src/components/index.ts

@@ -6,12 +6,11 @@ import { lazy } from 'react';
 
 export * from './AssistantSettings';
 export * from './Chat';
-export * from './ChatProgress';
-export * from './ChatToolbar';
 export * from './TemplateEditor';
 export * from './Toolbox';
 
-export const BlueprintContainer = lazy(() => import('./BlueprintContainer'));
+export const BlueprintArticle = lazy(() => import('./BlueprintArticle'));
 export const ChatCompanion = lazy(() => import('./ChatCompanion'));
 export const ChatContainer = lazy(() => import('./ChatContainer'));
 export const ChatDialog = lazy(() => import('./ChatDialog'));
+export const PromptArticle = lazy(() => import('./PromptArticle'));

package/src/functions/index.ts

@@ -2,6 +2,6 @@
 // Copyright 2025 DXOS.org
 //
 
-export { default as analysis } from './analysis';
 export { default as list } from './object-list';
 export { default as load } from './object-load';
+export { default as search } from './object-search';

package/src/functions/object-list.ts

@@ -6,7 +6,8 @@ import * as Effect from 'effect/Effect';
 import * as Schema from 'effect/Schema';
 
 import { Filter, Obj, Query } from '@dxos/echo';
-import { DatabaseService, defineFunction } from '@dxos/functions';
+import { Database } from '@dxos/echo';
+import { defineFunction } from '@dxos/functions';
 import { trim } from '@dxos/util';
 
 // TODO(burdon): Move to toolkit (i.e., tool not function).
@@ -34,10 +35,10 @@ export default defineFunction({
   }),
   handler: Effect.fn(function* ({ data: { typename } }) {
     // TODO(wittjosiah): Typename query is not working for dynamic schemas.
-    const [schema] = yield* DatabaseService.runSchemaQuery({ typename });
+    const [schema] = yield* Database.Service.runSchemaQuery({ typename });
     const filter = schema ? Filter.type(schema) : Filter.typename(typename);
 
-    const { objects } = yield* DatabaseService.runQuery(Query.select(filter));
+    const objects = yield* Database.Service.runQuery(Query.select(filter));
     const results = objects.map((object) => ({
       dxn: Obj.getDXN(object).toString(),
       label: Obj.getLabel(object),

package/src/functions/object-load.ts

@@ -8,8 +8,9 @@ import * as Option from 'effect/Option';
 import * as Schema from 'effect/Schema';
 
 import { ArtifactId } from '@dxos/assistant';
-import { Obj } from '@dxos/echo';
-import { DatabaseService, defineFunction } from '@dxos/functions';
+import { Entity } from '@dxos/echo';
+import { Database } from '@dxos/echo';
+import { defineFunction } from '@dxos/functions';
 
 // TODO(burdon): Move to toolkit (i.e., tool not function).
 export default defineFunction({
@@ -17,6 +18,7 @@ export default defineFunction({
   name: 'Load object',
   description: 'Loads the object.',
   inputSchema: Schema.Struct({
+    // TODO(dmaretskyi): Use Ref.
     id: ArtifactId.annotations({
       description: 'The ID of the object.',
     }),
@@ -25,14 +27,14 @@ export default defineFunction({
     }),
   }),
   outputSchema: Schema.Struct({
-    // TODO(wittjosiah): Type.Obj.Any | Type.Relation.Any
+    // TODO(wittjosiah): Type.Entity.Any
     object: Schema.Any,
   }),
   handler: Effect.fn(function* ({ data: { id, typename } }) {
-    const object = yield* DatabaseService.resolve(ArtifactId.toDXN(id));
+    const object = yield* Database.Service.resolve(ArtifactId.toDXN(id));
     return yield* Function.pipe(
       Option.fromNullable(object),
-      Option.flatMap((object) => (Obj.getTypename(object) === typename ? Option.some(object) : Option.none())),
+      Option.flatMap((object) => (Entity.getTypename(object) === typename ? Option.some(object) : Option.none())),
       Option.match({
         onNone: () => Effect.fail('Object not found.'),
         onSome: (object) => {