@dxos/assistant 0.7.4-df37807

package/src/cypher/execution-plan.ts

@@ -0,0 +1,208 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import type { JsonPath } from '@dxos/echo-schema';
+import { invariant } from '@dxos/invariant';
+import { log } from '@dxos/log';
+
+import type { GraphPatternConnector, MatchClause, ReturnClause, WhereClause } from './ast';
+import type { ParsedQuery } from './parser';
+
+export type VariableName = string & { __VariableName: never };
+
+export type ExecutionPlanStep =
+  | {
+      type: 'NodeScan';
+      label: string;
+      into: VariableName;
+    }
+  | {
+      type: 'FilterByLabel';
+      variable: VariableName;
+      label: string;
+    }
+  | {
+      type: 'FilterByProperty';
+      variable: VariableName;
+      property: JsonPath;
+      value: unknown;
+    }
+  | {
+      /**
+       * Scans (anchor)-[relationshipInto:label]->(nodeInto).
+       */
+      type: 'RelationshipScan';
+      label: string;
+      anchor: VariableName;
+      relationshipInto: VariableName;
+      nodeInto: VariableName;
+
+      /**
+       * `in`: (anchor)<-[relationshipInto:label]-(nodeInto)
+       * `out`: (anchor)-[relationshipInto:label]->(nodeInto)
+       * `both`: (anchor)-[relationshipInto:label]-(nodeInto)
+       */
+      direction: 'in' | 'out' | 'both';
+    }
+  | {
+      type: 'ProduceResults';
+      fields: {
+        name: string;
+        variable: VariableName;
+        path: JsonPath;
+      }[];
+    };
+
+export type ExecutionPlan = {
+  steps: ExecutionPlanStep[];
+  output: {
+    fields: string[];
+  };
+};
+
+export const createExecutionPlan = (query: ParsedQuery): ExecutionPlan => {
+  return new ExecutionPlanBuilder().build(query);
+};
+
+class ExecutionPlanBuilder {
+  private counter = 0;
+
+  newVariable(): VariableName {
+    return `anon_${this.counter++}` as VariableName;
+  }
+
+  build(query: ParsedQuery): ExecutionPlan {
+    log('Creating execution plan for query', { query });
+    return {
+      steps: [
+        ...this.createStepsForMatchClause(query.ast.matchClause),
+        ...(query.ast.whereClause ? this.createStepsForWhereClause(query.ast.whereClause) : []),
+        ...this.createStepsForReturnClause(query.ast.returnClause),
+      ],
+      output: {
+        fields: this.getOutputFields(query.ast.returnClause),
+      },
+    };
+  }
+
+  createStepsForMatchClause(match: MatchClause): ExecutionPlanStep[] {
+    invariant(match.pattern.segments.length > 0);
+    const firstSegment = match.pattern.segments[0];
+    invariant(firstSegment.astKind === 'NodePattern');
+    let lastNodeVariable = (firstSegment.variable?.name as VariableName) ?? this.newVariable();
+
+    const steps: ExecutionPlanStep[] = [];
+    steps.push({
+      type: 'NodeScan',
+      label: firstSegment.label.name,
+      into: lastNodeVariable,
+    });
+
+    if (firstSegment.properties) {
+      for (const property of firstSegment.properties.properties) {
+        steps.push({
+          type: 'FilterByProperty',
+          variable: lastNodeVariable,
+          property: property.key.name as JsonPath,
+          value: property.value.value,
+        });
+      }
+    }
+
+    for (let i = 1; i < match.pattern.segments.length; i++) {
+      const relationshipSegment = match.pattern.segments[i++];
+      invariant(relationshipSegment.astKind === 'RelationshipPattern');
+      const nodeSegment = match.pattern.segments[i];
+      invariant(nodeSegment.astKind === 'NodePattern');
+      const leftSeparator = match.pattern.segments.separators[i - 2];
+      const rightSeparator = match.pattern.segments.separators[i - 1];
+      const direction = getDirection(leftSeparator.direction, rightSeparator.direction);
+
+      const relationshipVariable = (relationshipSegment.variable?.name as VariableName) ?? this.newVariable();
+      const nodeVariable = (nodeSegment.variable?.name as VariableName) ?? this.newVariable();
+
+      steps.push({
+        type: 'RelationshipScan',
+        label: relationshipSegment.label.name,
+        direction,
+        anchor: lastNodeVariable,
+        relationshipInto: relationshipVariable,
+        nodeInto: nodeVariable,
+      });
+
+      if (relationshipSegment.properties) {
+        for (const property of relationshipSegment.properties.properties) {
+          steps.push({
+            type: 'FilterByProperty',
+            variable: relationshipVariable,
+            property: property.key.name as JsonPath,
+            value: property.value.value,
+          });
+        }
+      }
+      if (nodeSegment.label) {
+        steps.push({
+          type: 'FilterByLabel',
+          variable: nodeVariable,
+          label: nodeSegment.label.name,
+        });
+      }
+      if (nodeSegment.properties) {
+        for (const property of nodeSegment.properties.properties) {
+          steps.push({
+            type: 'FilterByProperty',
+            variable: nodeVariable,
+            property: property.key.name as JsonPath,
+            value: property.value.value,
+          });
+        }
+      }
+
+      lastNodeVariable = nodeVariable;
+    }
+
+    return steps;
+  }
+
+  createStepsForWhereClause(where: WhereClause): ExecutionPlanStep[] {
+    return [];
+  }
+
+  createStepsForReturnClause(returnClause: ReturnClause): ExecutionPlanStep[] {
+    return [
+      {
+        type: 'ProduceResults',
+        fields: returnClause.fields.map((field) => ({
+          name: field.path.map((p) => p.name).join('_'),
+          variable: field.path[0].name as VariableName,
+          path:
+            field.path.length > 1
+              ? (field.path
+                  .slice(1)
+                  .map((p) => p.name)
+                  .join('.') as JsonPath)
+              : ('$' as JsonPath),
+        })),
+      },
+    ];
+  }
+
+  getOutputFields(returnClause: ReturnClause): string[] {
+    return returnClause.fields.map((field) => field.path.map((p) => p.name).join('_'));
+  }
+}
+
+const getDirection = (
+  left: GraphPatternConnector['direction'],
+  right: GraphPatternConnector['direction'],
+): 'in' | 'out' | 'both' => {
+  if (left === '-' && right === '->') {
+    return 'out';
+  } else if (left === '<-' && right === '-') {
+    return 'in';
+  } else if (left === '-' && right === '-') {
+    return 'both';
+  }
+  throw new Error(`Invalid direction: ()${left}[]${right}()`);
+};

package/src/cypher/index.ts

@@ -0,0 +1,8 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+export * from './parser';
+export * from './query-executor';
+export * from './schema';
+export * from './echo-data-source';

package/src/cypher/parjs-shims.d.ts

@@ -0,0 +1,6 @@
+declare module '#parjs' {
+  import * as parjs from 'parjs';
+  import * as parjsCombinators from 'parjs/combinators';
+
+  export { parjsCombinators, parjs };
+}

package/src/cypher/parjs.ts

@@ -0,0 +1,10 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+// Workaround for https://github.com/vitest-dev/vitest/issues/7174.
+import { parjsCombinators, parjs } from '#parjs';
+
+import type { Parjser } from 'parjs';
+
+export { type Parjser, parjsCombinators, parjs };

package/src/cypher/parser.ts

@@ -0,0 +1,19 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import type { CypherQuery } from './ast';
+import { cypherQuery } from './combinators';
+
+export type ParsedQuery = {
+  ast: CypherQuery;
+};
+
+export const parseCypherQuery = (query: string): ParsedQuery => {
+  const parseResult = cypherQuery.parse(query);
+  if (parseResult.isOk) {
+    return { ast: parseResult.value };
+  } else {
+    throw new Error(`Failed to parse query: ${parseResult}`);
+  }
+};

package/src/cypher/query-executor.test.ts

@@ -0,0 +1,47 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { test } from 'vitest';
+
+import { log } from '@dxos/log';
+
+import { executeQuery } from './query-executor';
+import { createTestData } from '../testing';
+
+test('run', async ({ expect }) => {
+  const dataSource = createTestData();
+  // for (const node of dataSource.nodes) {
+  //   log.info('node', { id: node.id, label: node.label, name: node.properties.name });
+  // }
+  // for (const relationship of dataSource.relationships) {
+  //   log.info('relationship', {
+  //     source: relationship.source.id,
+  //     target: relationship.target.id,
+  //     label: relationship.label,
+  //   });
+  // }
+
+  const result = await executeQuery(
+    dataSource,
+    `
+      MATCH (org:Org {name: 'DXOS'})-[:ORG_EMPLOYEES]->(c:Contact)<-[:TASK_ASSIGNEE]-(t:Task)-[:TASK_PROJECT]->(p:Project {name: 'Composer'})
+      RETURN c.name, t.name
+    `,
+  );
+
+  expect(result).toEqual([
+    { c_name: 'Rich', t_name: 'Create form builder' },
+    {
+      c_name: 'Josiah',
+      t_name: 'Optimize startup performance',
+    },
+    { c_name: 'Rich', t_name: 'Implement dark mode' },
+    {
+      c_name: 'Josiah',
+      t_name: 'Add support for custom themes',
+    },
+  ]);
+
+  log.info('result', { result });
+});

package/src/cypher/query-executor.ts

@@ -0,0 +1,145 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import type { JsonPath } from '@dxos/echo-schema';
+import { log } from '@dxos/log';
+import { mapValues } from '@dxos/util';
+
+import { createExecutionPlan, type ExecutionPlan, type VariableName } from './execution-plan';
+import { parseCypherQuery } from './parser';
+
+export interface DataSource {
+  getNodes: (params: { label?: string }) => Promise<Node[]>;
+  getRelationships: (params: { label?: string }) => Promise<Relationship[]>;
+}
+
+export type Node = {
+  id: string;
+  kind: 'node';
+  label: string;
+  properties: Record<string, unknown>;
+};
+
+export type Relationship = {
+  id: string;
+  kind: 'relationship';
+  label: string;
+  properties: Record<string, unknown>;
+  source: Node;
+  target: Node;
+};
+
+export const executeQueryPlan = async (
+  dataSource: DataSource,
+  plan: ExecutionPlan,
+): Promise<Record<string, unknown>[]> => {
+  let buffer: Record<VariableName, Node | Relationship>[] = [];
+  const results: Record<string, unknown>[] = [];
+
+  for (const step of plan.steps) {
+    log('begin step', { step });
+
+    switch (step.type) {
+      case 'NodeScan': {
+        const nodes = await dataSource.getNodes({ label: step.label });
+        buffer.push(...nodes.map((node) => ({ [step.into]: node })));
+        break;
+      }
+      case 'FilterByLabel': {
+        buffer = buffer.filter((item) => {
+          const node = item[step.variable];
+          return node && node.kind === 'node' && node.label === step.label;
+        });
+        break;
+      }
+      case 'FilterByProperty': {
+        buffer = buffer.filter((item) => {
+          const node = item[step.variable];
+          return node && node.kind === 'node' && node.properties[step.property] === step.value;
+        });
+        break;
+      }
+      case 'RelationshipScan': {
+        const relationships = await dataSource.getRelationships({ label: step.label });
+        for (const tuple of [...buffer]) {
+          if (tuple[step.anchor].kind !== 'node') {
+            continue;
+          }
+          const filteredRelationships = relationships.filter((relationship) => {
+            switch (step.direction) {
+              case 'in':
+                return relationship.target.id === tuple[step.anchor].id;
+              case 'out':
+                return relationship.source.id === tuple[step.anchor].id;
+              case 'both':
+                return (
+                  relationship.source.id === tuple[step.anchor].id || relationship.target.id === tuple[step.anchor].id
+                );
+              default:
+                throw new Error(`Unknown direction: ${step.direction}`);
+            }
+          });
+          for (let i = 0; i < filteredRelationships.length; i++) {
+            const relationship = filteredRelationships[i];
+            const newTuple = i === 0 ? tuple : { ...tuple };
+            if (newTuple !== tuple) {
+              buffer.push(newTuple);
+            }
+            newTuple[step.relationshipInto] = relationship;
+            switch (step.direction) {
+              case 'in':
+                newTuple[step.nodeInto] = relationship.source;
+                break;
+              case 'out':
+                newTuple[step.nodeInto] = relationship.target;
+                break;
+              case 'both':
+                newTuple[step.nodeInto] =
+                  relationship.source.id === newTuple[step.anchor].id ? relationship.target : relationship.source;
+                break;
+            }
+          }
+        }
+        break;
+      }
+      case 'ProduceResults': {
+        for (const tuple of buffer) {
+          const result: Record<string, unknown> = {};
+          for (const field of step.fields) {
+            result[field.name] = jsonPathGet(tuple[field.variable].properties, field.path);
+          }
+          results.push(result);
+        }
+        break;
+      }
+    }
+
+    log('finish step', { buffer: buffer.map((tuple) => mapValues(tuple, (x) => x.id)), results });
+  }
+
+  return results;
+};
+
+export const executeQuery = (dataSource: DataSource, query: string): Promise<Record<string, unknown>[]> => {
+  const plan = createExecutionPlan(parseCypherQuery(query));
+  return executeQueryPlan(dataSource, plan);
+};
+
+// TODO(dmaretskyi): Extract.
+const jsonPathGet = (obj: Record<string, unknown>, path: JsonPath): unknown => {
+  if (path === '' || path === '$') {
+    return obj;
+  }
+
+  const parts = path.split('.');
+  if (parts[0] === '$') {
+    parts.shift();
+  }
+
+  let result: any = obj;
+  for (const part of parts) {
+    result = result?.[part];
+  }
+  return result;
+};

package/src/cypher/schema.test.ts

@@ -0,0 +1,17 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { test } from 'vitest';
+
+import { toJsonSchema } from '@dxos/echo-schema';
+
+import { formatJsonSchemaForLLM } from './schema';
+import { Contact, Org, Project, Task } from '../testing';
+
+test('Org-Project-Task-Contact schema', async ({ expect }) => {
+  const jsonSchema = [Org, Project, Task, Contact].map((schema) => toJsonSchema(schema));
+  const result = formatJsonSchemaForLLM(jsonSchema);
+  // log.info('result', { result });
+  expect(result).toHaveLength(8);
+});

package/src/cypher/schema.ts

@@ -0,0 +1,82 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { raise } from '@dxos/debug';
+import { JSON_SCHEMA_ECHO_REF_ID, type JsonSchemaType } from '@dxos/echo-schema';
+import { DXN } from '@dxos/keys';
+
+import { trim } from '../util';
+
+/**
+ * Formats ECHO schema to be consumed by LLM.
+ * Schema is formatted in cypher format.
+ * Objects are represented as nodes.
+ * References are represented as relationships.
+ */
+export const formatJsonSchemaForLLM = (schemaTypes: JsonSchemaType[]): string[] => {
+  return schemaTypes.flatMap((schema) => [
+    formatNodeSchema(schema),
+    ...formatInferredRelationships(schema, schemaTypes),
+  ]);
+};
+
+const formatNodeSchema = (schema: JsonSchemaType) => {
+  const label = formatNodeLabel(schema.typename ?? raise(new Error('Schema must have $id')));
+  const properties = Object.entries(schema.properties!).filter(([key, prop]) => key !== 'id' && !isReference(prop));
+
+  return trim`
+    <node>
+      (:${label})
+
+      <description>${schema.description}</description>
+
+      <properties>
+        ${properties.map(([key, value]) => {
+          return trim`
+            <${key}>
+              <type>${value.type}</type>
+              <description>${value.description}</description>
+            </${key}>
+          `;
+        })}
+      </properties>
+    </node>
+  `;
+};
+
+const formatInferredRelationships = (schema: JsonSchemaType, allSchema: JsonSchemaType[]): string[] => {
+  const nodeLabel = formatNodeLabel(schema.typename ?? raise(new Error('Schema must have $id')));
+  const relationships = Object.entries(schema.properties!).filter(([key, prop]) => isReference(prop));
+
+  return relationships.map(([key, value]) => {
+    const relationshipLabel = formatInferredRelationshipLabel(
+      schema.typename ?? raise(new Error('Schema must have $id')),
+      key,
+    );
+
+    // TODO(dmaretskyi): Registered schema don't get reference types updated.
+    const targetSchema = allSchema.find((schema) => `dxn:type/${schema.typename}` === value.reference?.schema.$ref);
+    const targetLabel = targetSchema ? formatNodeLabel(targetSchema?.typename ?? '') : '';
+
+    return trim`
+      <relationship>
+        (:${nodeLabel})-[:${relationshipLabel}]->(:${targetLabel})
+
+        <description>${value.description}</description>
+      </relationship>
+    `;
+  });
+};
+
+export const formatNodeLabel = (typename: string) => {
+  return typename.replace(/^example\.com\/type\//g, '');
+};
+
+export const formatInferredRelationshipLabel = (typename: string, property: string) => {
+  const nodeLabel = formatNodeLabel(typename);
+  return `${nodeLabel}_${property}`.toUpperCase();
+};
+
+export const isReference = (schema: JsonSchemaType) =>
+  schema.$id === JSON_SCHEMA_ECHO_REF_ID || (schema.type === 'array' && schema.items?.$id === JSON_SCHEMA_ECHO_REF_ID);

package/src/index.ts

@@ -0,0 +1,7 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+export * from './ai-service';
+export * from './cypher';
+export * from './conversation';

package/src/sanity.test.ts

@@ -0,0 +1,9 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+import { test } from 'vitest';
+
+test('sanity', ({ expect }) => {
+  expect(true).toBe(true);
+});

package/src/testing/index.ts

@@ -0,0 +1,8 @@
+//
+// Copyright 2024 DXOS.org
+//
+
+export * from './logger';
+export * from './query-promts';
+export * from './test-data';
+export * from './test-schema';