H-1767: Write service to insert embeddings into Graph (#3846)

Co-authored-by: Ciaran Morinan <[email protected]>
Co-authored-by: Ciaran Morinan <[email protected]>

Author: TimDiekmann

apps/hash-ai-worker-ts/package.json

@@ -38,6 +38,7 @@
     "@local/advanced-types": "0.0.0-private",
     "@local/hash-backend-utils": "0.0.0-private",
     "@local/hash-isomorphic-utils": "0.0.0-private",
+    "@local/hash-subgraph": "0.0.0-private",
     "@local/status": "0.0.0-private",
     "@temporalio/activity": "1.8.1",
     "@temporalio/worker": "1.8.1",
@@ -53,7 +54,6 @@
   "devDependencies": {
     "@local/eslint-config": "0.0.0-private",
     "@local/hash-graph-client": "0.0.0-private",
-    "@local/hash-subgraph": "0.0.0-private",
     "@local/tsconfig": "0.0.0-private",
     "@types/dedent": "0.7.0",
     "@types/dotenv-flow": "3.2.0",

apps/hash-ai-worker-ts/src/activities.ts

@@ -3,10 +3,19 @@ import type {
   InferEntitiesCallerParams,
   InferEntitiesReturn,
 } from "@local/hash-isomorphic-utils/ai-inference-types";
+import type {
+  BaseUrl,
+  EntityPropertiesObject,
+  PropertyTypeWithMetadata,
+} from "@local/hash-subgraph";
 import { ApplicationFailure } from "@temporalio/activity";
+import { CreateEmbeddingResponse } from "openai/resources";
 
+import { createEmbeddings } from "./activities/embeddings";
 import { inferEntities } from "./activities/infer-entities";
 
+export { createGraphActivities } from "./activities/graph";
+
 export const createAiActivities = ({
   graphApiClient,
 }: {
@@ -22,4 +31,14 @@ export const createAiActivities = ({
 
     return status;
   },
+
+  async createEmbeddingsActivity(params: {
+    entityProperties: EntityPropertiesObject;
+    propertyTypes: PropertyTypeWithMetadata[];
+  }): Promise<{
+    embeddings: { property?: BaseUrl; embedding: number[] }[];
+    usage: CreateEmbeddingResponse.Usage;
+  }> {
+    return createEmbeddings(params);
+  },
 });

apps/hash-ai-worker-ts/src/activities/embeddings.ts

@@ -0,0 +1,78 @@
+import type {
+  BaseUrl,
+  EntityPropertiesObject,
+  EntityPropertyValue,
+  PropertyTypeWithMetadata,
+} from "@local/hash-subgraph";
+import OpenAI from "openai/index";
+import { CreateEmbeddingResponse } from "openai/resources";
+import Usage = CreateEmbeddingResponse.Usage;
+
+const openai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY,
+});
+
+const createEmbeddingInput = (params: {
+  propertyType: PropertyTypeWithMetadata;
+  property: EntityPropertyValue;
+}): string => {
+  return `${params.propertyType.schema.title}: ${JSON.stringify(
+    params.property,
+  )}`;
+};
+
+export const createEmbeddings = async (params: {
+  entityProperties: EntityPropertiesObject;
+  propertyTypes: PropertyTypeWithMetadata[];
+}): Promise<{
+  embeddings: { property?: BaseUrl; embedding: number[] }[];
+  usage: Usage;
+}> => {
+  if (params.propertyTypes.length === 0) {
+    return {
+      embeddings: [],
+      usage: {
+        prompt_tokens: 0,
+        total_tokens: 0,
+      },
+    };
+  }
+
+  // sort property types by their base url
+  params.propertyTypes.sort((a, b) =>
+    a.metadata.recordId.baseUrl.localeCompare(b.metadata.recordId.baseUrl),
+  );
+
+  // We want to create embeddings for:
+  //   1. Each individual '[Property Title]: [Value]' pair, and
+  //   2. A list of all property key:value pairs
+  //
+  // We use the last item in the array to store the combined 'all properties' list.
+  const propertyEmbeddings = [];
+  let combinedEntityEmbedding = "";
+  for (const propertyType of params.propertyTypes) {
+    const property =
+      params.entityProperties[propertyType.metadata.recordId.baseUrl];
+    if (property === undefined) {
+      // `property` could be `null` or `false` as well but that is part of the semantic meaning of the property
+      // and should be included in the embedding.
+      continue;
+    }
+    const embeddingInput = createEmbeddingInput({ propertyType, property });
+    combinedEntityEmbedding += `${embeddingInput}\n`;
+    propertyEmbeddings.push(embeddingInput);
+  }
+
+  const response = await openai.embeddings.create({
+    input: [...propertyEmbeddings, combinedEntityEmbedding],
+    model: "text-embedding-ada-002",
+  });
+
+  return {
+    usage: response.usage,
+    embeddings: response.data.map((data, idx) => ({
+      property: params.propertyTypes[idx]?.metadata.recordId.baseUrl,
+      embedding: data.embedding,
+    })),
+  };
+};

apps/hash-ai-worker-ts/src/activities/graph.ts

@@ -0,0 +1,53 @@
+import type {
+  EntityEmbedding,
+  EntityTypeStructuralQuery,
+  GraphApi,
+} from "@local/hash-graph-client";
+import type {
+  AccountId,
+  EntityTypeRootType,
+  PropertyTypeWithMetadata,
+  Subgraph,
+} from "@local/hash-subgraph";
+import {
+  getPropertyTypes,
+  mapGraphApiSubgraphToSubgraph,
+} from "@local/hash-subgraph/stdlib";
+
+export const createGraphActivities = ({
+  graphApiClient,
+}: {
+  graphApiClient: GraphApi;
+}) => ({
+  async getEntityTypesByQuery(params: {
+    authentication: {
+      actorId: AccountId;
+    };
+    query: EntityTypeStructuralQuery;
+  }): Promise<Subgraph<EntityTypeRootType>> {
+    return graphApiClient
+      .getEntityTypesByQuery(params.authentication.actorId, params.query)
+      .then((response) => mapGraphApiSubgraphToSubgraph(response.data));
+  },
+
+  async updateEntityEmbeddings(params: {
+    authentication: {
+      actorId: AccountId;
+    };
+    embeddings: EntityEmbedding[];
+  }): Promise<void> {
+    await graphApiClient
+      .updateEntityEmbeddings(params.authentication.actorId, {
+        embeddings: params.embeddings,
+        reset: true,
+      })
+      .then((response) => response.data);
+  },
+
+  // eslint-disable-next-line @typescript-eslint/require-await
+  async getSubgraphPropertyTypes(params: {
+    subgraph: Subgraph;
+  }): Promise<PropertyTypeWithMetadata[]> {
+    return getPropertyTypes(params.subgraph);
+  },
+});

apps/hash-ai-worker-ts/src/main.ts

@@ -7,7 +7,7 @@ import { Logger } from "@local/hash-backend-utils/logger";
 import { NativeConnection, Worker } from "@temporalio/worker";
 import { config } from "dotenv-flow";
 
-import { createAiActivities } from "./activities";
+import { createAiActivities, createGraphActivities } from "./activities";
 
 export const monorepoRootDir = path.resolve(__dirname, "../../..");
 
@@ -59,9 +59,14 @@ async function run() {
 
   const worker = await Worker.create({
     ...workflowOption(),
-    activities: createAiActivities({
-      graphApiClient,
-    }),
+    activities: {
+      ...createAiActivities({
+        graphApiClient,
+      }),
+      ...createGraphActivities({
+        graphApiClient,
+      }),
+    },
     connection: await NativeConnection.connect({
       address: `${TEMPORAL_HOST}:${TEMPORAL_PORT}`,
     }),

apps/hash-ai-worker-ts/src/workflows.ts

@@ -1,7 +1,9 @@
 import type { InferEntitiesCallerParams } from "@local/hash-isomorphic-utils/ai-inference-types";
+import type { AccountId, Entity } from "@local/hash-subgraph";
 import { proxyActivities } from "@temporalio/workflow";
+import { CreateEmbeddingResponse } from "openai/resources";
 
-import { createAiActivities } from "./activities";
+import { createAiActivities, createGraphActivities } from "./activities";
 
 const aiActivities = proxyActivities<ReturnType<typeof createAiActivities>>({
   startToCloseTimeout: "1800 second",
@@ -10,5 +12,77 @@ const aiActivities = proxyActivities<ReturnType<typeof createAiActivities>>({
   },
 });
 
+const graphActivities = proxyActivities<
+  ReturnType<typeof createGraphActivities>
+>({
+  startToCloseTimeout: "10 second",
+  retry: {
+    maximumAttempts: 3,
+  },
+});
+
 export const inferEntities = (params: InferEntitiesCallerParams) =>
   aiActivities.inferEntitiesActivity(params);
+
+export const updateEntityEmbeddings = async (params: {
+  authentication: {
+    actorId: AccountId;
+  };
+  entity: Entity;
+}): Promise<CreateEmbeddingResponse.Usage> => {
+  const subgraph = await graphActivities.getEntityTypesByQuery({
+    authentication: params.authentication,
+    query: {
+      filter: {
+        equal: [
+          { path: ["versionedUrl"] },
+          { parameter: params.entity.metadata.entityTypeId },
+        ],
+      },
+      graphResolveDepths: {
+        inheritsFrom: { outgoing: 255 },
+        constrainsValuesOn: { outgoing: 0 },
+        constrainsPropertiesOn: { outgoing: 1 },
+        constrainsLinksOn: { outgoing: 0 },
+        constrainsLinkDestinationsOn: { outgoing: 0 },
+        isOfType: { outgoing: 0 },
+        hasLeftEntity: { incoming: 0, outgoing: 0 },
+        hasRightEntity: { incoming: 0, outgoing: 0 },
+      },
+      temporalAxes: {
+        pinned: {
+          axis: "transactionTime",
+          timestamp: null,
+        },
+        variable: {
+          axis: "decisionTime",
+          interval: {
+            start: null,
+            end: null,
+          },
+        },
+      },
+      includeDrafts: false,
+    },
+  });
+  const propertyTypes = await graphActivities.getSubgraphPropertyTypes({
+    subgraph,
+  });
+
+  const generatedEmbeddings = await aiActivities.createEmbeddingsActivity({
+    entityProperties: params.entity.properties,
+    propertyTypes,
+  });
+
+  if (generatedEmbeddings.embeddings.length > 0) {
+    await graphActivities.updateEntityEmbeddings({
+      authentication: params.authentication,
+      embeddings: generatedEmbeddings.embeddings.map((embedding) => ({
+        ...embedding,
+        entityId: params.entity.metadata.recordId.entityId,
+      })),
+    });
+  }
+
+  return generatedEmbeddings.usage;
+};