@vertesia/workflow 1.2.0 → 1.3.0-dev.20260620.061059Z

package/src/activities/generateEmbeddings.test.ts

@@ -0,0 +1,114 @@
+import { MockActivityEnvironment } from '@temporalio/testing';
+import type { VertesiaClient } from '@vertesia/client';
+import {
+    ContentEventName,
+    type ContentObject,
+    type DSLActivityExecutionPayload,
+    SupportedEmbeddingTypes,
+} from '@vertesia/common';
+import { beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import type { ActivityContext } from '../dsl/setup/ActivityContext.js';
+import { type GenerateEmbeddingsParams, generateEmbeddings } from './generateEmbeddings.js';
+
+vi.mock('../dsl/setup/ActivityContext.js', async (importOriginal) => {
+    const actual = await importOriginal<typeof import('../dsl/setup/ActivityContext.js')>();
+    return { ...actual, setupActivity: vi.fn() };
+});
+
+let testEnv: MockActivityEnvironment;
+
+beforeAll(() => {
+    testEnv = new MockActivityEnvironment();
+});
+
+beforeEach(() => {
+    vi.clearAllMocks();
+});
+
+const createPayload = (params: GenerateEmbeddingsParams): DSLActivityExecutionPayload<GenerateEmbeddingsParams> => {
+    const activityParams: GenerateEmbeddingsParams & Record<string, unknown> = { ...params };
+    return {
+        auth_token: 'mock-token',
+        account_id: 'test-account',
+        project_id: 'test-project',
+        params,
+        config: { studio_url: 'http://mock-studio', store_url: 'http://mock-store' },
+        workflow_name: 'StandardDocumentIntake',
+        event: ContentEventName.create,
+        objectIds: ['properties-only-object'],
+        input: { inputType: 'objectIds', objectIds: ['properties-only-object'] },
+        vars: {},
+        activity: { name: 'generateEmbeddings', params: activityParams },
+    };
+};
+
+describe('generateEmbeddings', () => {
+    it('should generate property embeddings for an object without content', async () => {
+        const { setupActivity } = await import('../dsl/setup/ActivityContext.js');
+        const document = {
+            id: 'properties-only-object',
+            properties: {
+                title: 'Properties-only object',
+                category: 'metadata',
+            },
+        } satisfies Partial<ContentObject>;
+        const embeddingResponse = {
+            model: 'embedding-model',
+            results: [{ outputs: [{ values: [0.1, 0.2, 0.3] }] }],
+        };
+        const client = {
+            objects: {
+                retrieve: vi.fn().mockResolvedValue(document),
+                setEmbedding: vi.fn().mockResolvedValue(undefined),
+            },
+            environments: {
+                embeddings: vi.fn().mockResolvedValue(embeddingResponse),
+            },
+        } as unknown as VertesiaClient;
+        const params = {
+            type: SupportedEmbeddingTypes.properties,
+            force: false,
+        } satisfies GenerateEmbeddingsParams;
+
+        vi.mocked(setupActivity).mockResolvedValue({
+            client,
+            objectId: document.id,
+            params,
+            fetchProject: vi.fn().mockResolvedValue({
+                name: 'Test Project',
+                namespace: 'test-project',
+                configuration: {
+                    embeddings: {
+                        [SupportedEmbeddingTypes.properties]: {
+                            enabled: true,
+                            environment: 'test-environment',
+                            max_tokens: 8000,
+                        },
+                    },
+                },
+            }),
+        } as unknown as ActivityContext<GenerateEmbeddingsParams>);
+
+        const result = await testEnv.run(generateEmbeddings, createPayload(params));
+
+        expect(result).toEqual({
+            id: document.id,
+            type: SupportedEmbeddingTypes.properties,
+            status: 'completed',
+            len: 3,
+        });
+        expect(client.objects.retrieve).toHaveBeenCalledWith(
+            document.id,
+            '+text +parts +embeddings +tokens +properties',
+        );
+        expect(client.environments.embeddings).toHaveBeenCalledWith('test-environment', {
+            inputs: [{ type: 'text', text: JSON.stringify(document.properties) }],
+            model: undefined,
+        });
+        expect(client.objects.setEmbedding).toHaveBeenCalledWith(document.id, SupportedEmbeddingTypes.properties, {
+            values: embeddingResponse.results[0].outputs[0].values,
+            model: embeddingResponse.model,
+            etag: expect.any(String),
+        });
+    });
+});

package/src/activities/generateEmbeddings.ts

@@ -1,19 +1,19 @@
-import { EmbeddingsResult } from "@llumiverse/common";
-import { log } from "@temporalio/activity";
-import { VertesiaClient, ZenoClientNotFoundError } from "@vertesia/client";
+import { log } from '@temporalio/activity';
+import { type VertesiaClient, ZenoClientNotFoundError } from '@vertesia/client';
 import {
-    ContentObject,
-    DSLActivityExecutionPayload,
-    DSLActivitySpec,
+    type ContentObject,
+    type DSLActivityExecutionPayload,
+    type DSLActivitySpec,
+    type EmbeddingsApiResult,
     ImageRenditionFormat,
-    ProjectConfigurationEmbeddings,
+    type ProjectConfigurationEmbedding,
     SupportedEmbeddingTypes,
-} from "@vertesia/common";
-import { setupActivity } from "../dsl/setup/ActivityContext.js";
-import { DocumentNotFoundError } from "../errors.js";
-import { fetchBlobAsBase64, md5 } from "../utils/blobs.js";
-import { DocPart } from "../utils/chunks.js";
-import { countTokens } from "../utils/tokens.js";
+} from '@vertesia/common';
+import { setupActivity } from '../dsl/setup/ActivityContext.js';
+import { DocumentNotFoundError } from '../errors.js';
+import { fetchBlobAsBase64, md5 } from '../utils/blobs.js';
+import type { DocPart } from '../utils/chunks.js';
+import { countTokens } from '../utils/tokens.js';
 
 export interface GenerateEmbeddingsParams {
     /**
@@ -44,31 +44,25 @@ export interface GenerateEmbeddingsParams {
     parts?: DocPart[];
 }
 
-export interface GenerateEmbeddings
-    extends DSLActivitySpec<GenerateEmbeddingsParams> {
-    name: "generateEmbeddings";
+export interface GenerateEmbeddings extends DSLActivitySpec<GenerateEmbeddingsParams> {
+    name: 'generateEmbeddings';
 }
 
-export async function generateEmbeddings(
-    payload: DSLActivityExecutionPayload<GenerateEmbeddingsParams>,
-) {
-    const { params, client, objectId, fetchProject } =
-        await setupActivity<GenerateEmbeddingsParams>(payload);
+export async function generateEmbeddings(payload: DSLActivityExecutionPayload<GenerateEmbeddingsParams>) {
+    const { params, client, objectId, fetchProject } = await setupActivity<GenerateEmbeddingsParams>(payload);
     const { force, type } = params;
 
     const projectData = await fetchProject();
     const config = projectData?.configuration.embeddings[type];
     if (!projectData) {
-        throw new DocumentNotFoundError("Project not found", [payload.project_id]);
+        throw new DocumentNotFoundError('Project not found', [payload.project_id]);
     }
     if (!config) {
-        throw new DocumentNotFoundError("Embeddings configuration not found", [
-            objectId,
-        ]);
+        throw new DocumentNotFoundError('Embeddings configuration not found', [objectId]);
     }
 
     if (!projectData) {
-        throw new DocumentNotFoundError("Project not found", [payload.project_id]);
+        throw new DocumentNotFoundError('Project not found', [payload.project_id]);
     }
 
     if (!projectData?.configuration.embeddings[type]?.enabled) {
@@ -78,7 +72,7 @@ export async function generateEmbeddings(
         );
         return {
             id: objectId,
-            status: "skipped",
+            status: 'skipped',
             message: `Embeddings generation disabled for type ${type}`,
         };
     }
@@ -90,16 +84,13 @@ export async function generateEmbeddings(
 
     if (!config.environment) {
         throw new Error(
-            "No environment found in project configuration. Set environment in project configuration to generate embeddings.",
+            'No environment found in project configuration. Set environment in project configuration to generate embeddings.',
         );
     }
 
-    let document;
+    let document: Awaited<ReturnType<typeof client.objects.retrieve>>;
     try {
-        document = await client.objects.retrieve(
-            objectId,
-            "+text +parts +embeddings +tokens +properties",
-        );
+        document = await client.objects.retrieve(objectId, '+text +parts +embeddings +tokens +properties');
     } catch (error) {
         if (error instanceof ZenoClientNotFoundError) {
             throw new DocumentNotFoundError(`Document not found: ${objectId}`, [objectId]);
@@ -108,14 +99,13 @@ export async function generateEmbeddings(
     }
 
     if (!document) {
-        throw new DocumentNotFoundError("Document not found", [objectId]);
-    }
-
-    if (!document.content) {
-        throw new DocumentNotFoundError("Document content not found", [objectId]);
+        throw new DocumentNotFoundError('Document not found', [objectId]);
     }
 
-    let res;
+    let res:
+        | Awaited<ReturnType<typeof generateTextEmbeddings>>
+        | Awaited<ReturnType<typeof generateImageEmbeddings>>
+        | { id: string; status: string; message: string };
 
     switch (type) {
         case SupportedEmbeddingTypes.text:
@@ -148,7 +138,7 @@ export async function generateEmbeddings(
         default:
             res = {
                 id: objectId,
-                status: "failed",
+                status: 'failed',
                 message: `unsupported embedding type: ${type}`,
             };
     }
@@ -160,138 +150,118 @@ interface ExecuteGenerateEmbeddingsParams {
     document: ContentObject;
     client: VertesiaClient;
     type: SupportedEmbeddingTypes;
-    config: ProjectConfigurationEmbeddings;
+    config: ProjectConfigurationEmbedding;
     property?: string;
     force?: boolean;
 }
 
-async function generateTextEmbeddings(
-    { document, client, type, config, force }: ExecuteGenerateEmbeddingsParams
-) {
-
+async function generateTextEmbeddings({ document, client, type, config, force }: ExecuteGenerateEmbeddingsParams) {
     if (!document) {
-        return { status: "error", message: "document is null or undefined" };
+        return { status: 'error', message: 'document is null or undefined' };
     }
 
-    if (
-        type !== SupportedEmbeddingTypes.text &&
-        type !== SupportedEmbeddingTypes.properties
-    ) {
+    if (type !== SupportedEmbeddingTypes.text && type !== SupportedEmbeddingTypes.properties) {
         return {
             id: document.id,
-            status: "failed",
+            status: 'failed',
             message: `unsupported embedding type: ${type}`,
         };
     }
 
-    if (type === SupportedEmbeddingTypes.text && !document.text) {
-        return { id: document.id, status: "failed", message: "no text found" };
-    }
-    if (type === SupportedEmbeddingTypes.properties && !document?.properties) {
+    const sourceText =
+        type === SupportedEmbeddingTypes.text
+            ? document.text
+            : document.properties
+              ? JSON.stringify(document.properties)
+              : undefined;
+
+    if (!sourceText) {
         return {
             id: document.id,
-            status: "failed",
-            message: "no properties found",
+            type,
+            status: 'skipped',
+            message: type === SupportedEmbeddingTypes.text ? 'no text found' : 'no properties found',
         };
     }
 
     const { environment } = config;
+    if (!environment) {
+        throw new Error(
+            'No environment found in project configuration. Set environment in project configuration to generate embeddings.',
+        );
+    }
 
-    // Compute text etag for comparison
-    const textEtag = document.text_etag ?? (document.text ? md5(document.text) : undefined);
+    const sourceEtag =
+        type === SupportedEmbeddingTypes.text ? (document.text_etag ?? md5(sourceText)) : md5(sourceText);
 
     // Skip if embeddings already exist with matching etag (unless force=true)
     const existingEmbedding = document.embeddings?.[type];
-    if (!force && existingEmbedding?.etag && textEtag && existingEmbedding.etag === textEtag) {
+    if (!force && existingEmbedding?.etag && existingEmbedding.etag === sourceEtag) {
         log.debug(`Skipping ${type} embeddings for document ${document.id} - etag unchanged`);
         return {
             id: document.id,
             type,
-            status: "skipped",
-            message: "embeddings already exist with matching etag",
+            status: 'skipped',
+            message: 'embeddings already exist with matching etag',
         };
     }
 
     // Count tokens if needed, do not rely on existing token count
-    let tokenCount : number | undefined = undefined;
-    if (type === SupportedEmbeddingTypes.text && document.text) {
-        tokenCount = countTokens(document.text).count;
-    }
-
-    if (type === SupportedEmbeddingTypes.properties && document.properties) {
-        const propertiesText = JSON.stringify(document.properties);
-        tokenCount = countTokens(propertiesText).count;
-    }
+    const tokenCount = countTokens(sourceText).count;
 
     const maxTokens = config.max_tokens ?? 8000;
 
     //generate embeddings for the main doc if document isn't too large
     log.debug(`Generating ${type} embeddings for document ${document.id}`);
-    if (
-        tokenCount !== undefined && tokenCount > maxTokens
-    ) {
+    if (tokenCount !== undefined && tokenCount > maxTokens) {
         //TODO: Review strategy for large documents
-        log.warn(
-            `Document too large for ${type} embeddings generation, skipping (${tokenCount} tokens)`,
-        );
+        log.warn(`Document too large for ${type} embeddings generation, skipping (${tokenCount} tokens)`);
         return {
             id: document.id,
-            status: "skipped",
+            status: 'skipped',
             message: `${type} embeddings generation, skipped for large document (${tokenCount} tokens)`,
-        }
+        };
     } else {
         log.debug(`Generating ${type} embeddings for document`);
 
-        const res = await generateEmbeddingsFromStudio(
-            JSON.stringify(document[type]),
-            environment,
-            client,
-        );
-        if (!res || !res.values) {
+        const res = await generateEmbeddingsFromStudio(sourceText, environment, client);
+        const values = res?.results?.[0]?.outputs?.[0]?.values;
+        if (!values) {
             return {
                 id: document.id,
-                status: "failed",
-                message: "no embeddings generated",
+                status: 'failed',
+                message: 'no embeddings generated',
             };
         }
 
         log.debug(`${type} embeddings generated for document ${document.id}`, {
-            len: res.values.length,
+            len: values.length,
         });
         await client.objects.setEmbedding(document.id, type, {
-            values: res.values,
+            values,
             model: res.model,
-            etag: textEtag,
+            etag: sourceEtag,
         });
 
         return {
             id: document.id,
             type,
-            status: "completed",
-            len: res.values.length,
+            status: 'completed',
+            len: values.length,
         };
     }
 }
 
-async function generateImageEmbeddings({
-    document,
-    client,
-    type,
-    config,
-    force,
-}: ExecuteGenerateEmbeddingsParams) {
-    log.debug("Generating image embeddings for document " + document.id, {
+async function generateImageEmbeddings({ document, client, type, config, force }: ExecuteGenerateEmbeddingsParams) {
+    log.debug(`Generating image embeddings for document ${document.id}`, {
         content: document.content,
     });
-    if (
-        !document.content?.type?.startsWith("image/") &&
-        !document.content?.type?.includes("pdf")
-    ) {
+    if (!document.content?.type?.startsWith('image/') && !document.content?.type?.includes('pdf')) {
         return {
             id: document.id,
             type,
-            status: "failed",
-            message: "content is not an image",
+            status: 'failed',
+            message: 'content is not an image',
         };
     }
 
@@ -305,12 +275,17 @@ async function generateImageEmbeddings({
         return {
             id: document.id,
             type,
-            status: "skipped",
-            message: "embeddings already exist with matching etag",
+            status: 'skipped',
+            message: 'embeddings already exist with matching etag',
         };
     }
 
     const { environment, model } = config;
+    if (!environment) {
+        throw new Error(
+            'No environment found in project configuration. Set environment in project configuration to generate embeddings.',
+        );
+    }
 
     const resRnd = await client.store.objects.getRendition(document.id, {
         format: ImageRenditionFormat.jpeg,
@@ -318,60 +293,58 @@ async function generateImageEmbeddings({
         sign_url: true,
     });
 
-    if (resRnd.status === "generating") {
-        throw new Error("Rendition is generating, will retry later");
-    } else if (
-        resRnd.status === "failed" ||
-        !resRnd.renditions ||
-        !resRnd.renditions.length
-    ) {
-        throw new DocumentNotFoundError("Rendition retrieval failed", [document.id]);
+    if (resRnd.status === 'generating') {
+        throw new Error('Rendition is generating, will retry later');
+    } else if (resRnd.status === 'failed' || !resRnd.renditions || !resRnd.renditions.length) {
+        throw new DocumentNotFoundError('Rendition retrieval failed', [document.id]);
     }
 
     const renditions = resRnd.renditions;
     if (!renditions?.length) {
-        throw new DocumentNotFoundError("No source found in rendition", [
-            document.id,
-        ]);
+        throw new DocumentNotFoundError('No source found in rendition', [document.id]);
     }
 
     const rendition = renditions[0];
     const image = await fetchBlobAsBase64(client, rendition);
 
+    // TODO(task_type): Document embedding task — image inputs do not carry task_type, but this is a document-indexing call.
+    // Revisit if task_type support is added to ImageEmbeddingInput in a later PR.
     const res = await client.environments
         .embeddings(environment, {
-            image,
+            inputs: [
+                {
+                    type: 'image',
+                    source: { base64: image, mime_type: 'image/jpeg' },
+                },
+            ],
             model,
         })
         .then((res) => res)
         .catch((e) => {
-            log.error("Error generating embeddings for image", { error: e });
+            log.error('Error generating embeddings for image', { error: e });
             throw e;
         });
 
-    if (!res || !res.values) {
+    const values = res?.results?.[0]?.outputs?.[0]?.values;
+    if (!values) {
         return {
             id: document.id,
-            status: "failed",
-            message: "no embeddings generated",
+            status: 'failed',
+            message: 'no embeddings generated',
         };
     }
 
-    await client.objects.setEmbedding(
-        document.id,
-        SupportedEmbeddingTypes.image,
-        {
-            values: res.values,
-            model: res.model,
-            etag: contentEtag,
-        },
-    );
+    await client.objects.setEmbedding(document.id, SupportedEmbeddingTypes.image, {
+        values,
+        model: res.model,
+        etag: contentEtag,
+    });
 
     return {
         id: document.id,
         type,
-        status: "completed",
-        len: res.values.length,
+        status: 'completed',
+        len: values.length,
     };
 }
 
@@ -380,19 +353,18 @@ async function generateEmbeddingsFromStudio(
     env: string,
     client: VertesiaClient,
     model?: string,
-): Promise<EmbeddingsResult> {
-    log.debug(
-        `Generating embeddings for text of ${text.length} chars with environment ${env}`,
-    );
+): Promise<EmbeddingsApiResult> {
+    log.debug(`Generating embeddings for text of ${text.length} chars with environment ${env}`);
 
+    // TODO(task_type): Document embedding task — add task_type: "document" once task_type support is validated end-to-end.
     return client.environments
         .embeddings(env, {
-            text,
+            inputs: [{ type: 'text', text }],
             model,
         })
         .then((res) => res)
         .catch((e) => {
-            log.error("Error generating embeddings for text", { error: e });
+            log.error('Error generating embeddings for text', { error: e });
             throw e;
         });
 }