@equinor/fusion-framework-cli-plugin-ai-index 2.0.0 → 2.1.0

package/src/bin/apply-schema.test.ts

@@ -0,0 +1,170 @@
+import { describe, it, expect } from 'vitest';
+import { z } from 'zod';
+import { of, lastValueFrom } from 'rxjs';
+import type { VectorStoreDocument } from '@equinor/fusion-framework-module-ai/lib';
+
+import { defineIndexSchema } from '../schema.js';
+import { applySchema } from './apply-schema.js';
+
+/** Helper to create a minimal VectorStoreDocument for testing. */
+function makeDocument(
+  overrides: Partial<VectorStoreDocument> & { metadata: VectorStoreDocument['metadata'] },
+): VectorStoreDocument {
+  return {
+    id: 'test-id',
+    pageContent: 'test content',
+    ...overrides,
+  };
+}
+
+describe('defineIndexSchema', () => {
+  it('returns the same config object (type-narrowing only)', () => {
+    const shape = z.object({ type: z.string() });
+    const resolve = () => ({ type: 'tsdoc' });
+
+    const schema = defineIndexSchema({ shape, resolve });
+
+    expect(schema.shape).toBe(shape);
+    expect(schema.resolve).toBe(resolve);
+  });
+});
+
+describe('applySchema', () => {
+  const schema = defineIndexSchema({
+    shape: z.object({
+      pkg_name: z.string().optional(),
+      type: z.string(),
+      tags: z.array(z.string()).default([]),
+      source_dir: z.string(),
+    }),
+    resolve: (doc) => ({
+      pkg_name: doc.metadata.attributes?.pkg_name as string | undefined,
+      type: (doc.metadata.attributes?.type as string) ?? 'unknown',
+      tags: (doc.metadata.attributes?.tags as string[]) ?? [],
+      source_dir: doc.metadata.source.split('/')[0],
+    }),
+  });
+
+  it('passes through unchanged when schema is undefined', async () => {
+    const doc = makeDocument({
+      metadata: { source: 'packages/foo/src/index.ts', attributes: { type: 'tsdoc' } },
+    });
+    const docs$ = of([doc]);
+
+    const result = await lastValueFrom(applySchema(docs$, undefined));
+
+    expect(result).toEqual([doc]);
+  });
+
+  it('resolves promoted fields and stores them on metadata.schemaFields', async () => {
+    const doc = makeDocument({
+      metadata: {
+        source: 'packages/foo/src/index.ts',
+        attributes: {
+          type: 'tsdoc',
+          pkg_name: '@equinor/fusion-framework',
+          tags: ['package', 'react'],
+          other_attr: 'keep-me',
+        },
+      },
+    });
+    const docs$ = of([doc]);
+
+    const result = await lastValueFrom(applySchema(docs$, schema));
+
+    expect(result[0].metadata.schemaFields).toEqual({
+      pkg_name: '@equinor/fusion-framework',
+      type: 'tsdoc',
+      tags: ['package', 'react'],
+      source_dir: 'packages',
+    });
+  });
+
+  it('removes promoted keys from attributes to avoid duplication', async () => {
+    const doc = makeDocument({
+      metadata: {
+        source: 'packages/foo/src/index.ts',
+        attributes: {
+          type: 'tsdoc',
+          pkg_name: '@equinor/fusion-framework',
+          tags: ['package'],
+          git_commit_hash: 'abc123',
+        },
+      },
+    });
+    const docs$ = of([doc]);
+
+    const result = await lastValueFrom(applySchema(docs$, schema));
+
+    // Promoted keys removed, non-promoted keys preserved
+    expect(result[0].metadata.attributes).toEqual({ git_commit_hash: 'abc123' });
+  });
+
+  it('handles documents with no attributes gracefully', async () => {
+    const doc = makeDocument({
+      metadata: { source: 'cookbooks/app-react/src/App.tsx' },
+    });
+    const docs$ = of([doc]);
+
+    const result = await lastValueFrom(applySchema(docs$, schema));
+
+    expect(result[0].metadata.schemaFields).toEqual({
+      pkg_name: undefined,
+      type: 'unknown',
+      tags: [],
+      source_dir: 'cookbooks',
+    });
+    expect(result[0].metadata.attributes).toEqual({});
+  });
+
+  it('throws when resolved values fail Zod validation', async () => {
+    const badSchema = defineIndexSchema({
+      shape: z.object({ type: z.string().min(1) }),
+      resolve: () => ({ type: '' }), // Empty string fails min(1)
+    });
+
+    const doc = makeDocument({
+      metadata: { source: 'test.ts', attributes: {} },
+    });
+    const docs$ = of([doc]);
+
+    await expect(lastValueFrom(applySchema(docs$, badSchema))).rejects.toThrow();
+  });
+
+  it('runs prepareAttributes before resolve to enrich attributes', async () => {
+    const schemaWithPrepare = defineIndexSchema({
+      shape: z.object({
+        tags: z.array(z.string()).default([]),
+        type: z.string(),
+      }),
+      prepareAttributes: (attrs, doc) => {
+        // Type-safe: attrs.tags is string[] | undefined
+        attrs.tags ??= [];
+        if (doc.metadata.source.includes('packages/')) {
+          attrs.tags.push('package');
+        }
+        return attrs;
+      },
+      resolve: (doc) => ({
+        tags: (doc.metadata.attributes?.tags as string[]) ?? [],
+        type: (doc.metadata.attributes?.type as string) ?? 'unknown',
+      }),
+    });
+
+    const doc = makeDocument({
+      metadata: {
+        source: 'packages/framework/src/index.ts',
+        attributes: { type: 'tsdoc' },
+      },
+    });
+    const docs$ = of([doc]);
+
+    const result = await lastValueFrom(applySchema(docs$, schemaWithPrepare));
+
+    // prepareAttributes added 'package' tag before resolve consumed it
+    expect(result[0].metadata.schemaFields).toEqual({
+      tags: ['package'],
+      type: 'tsdoc',
+    });
+  });
+});

package/src/bin/apply-schema.ts

@@ -0,0 +1,86 @@
+import { map } from 'rxjs';
+import type { Observable } from 'rxjs';
+import type { VectorStoreDocument } from '@equinor/fusion-framework-module-ai/lib';
+import type { IndexSchemaConfig } from '../schema.js';
+
+/**
+ * Creates an RxJS operator that resolves promoted schema fields for each
+ * document and separates them from the generic `attributes` bag.
+ *
+ * For each document in the batch:
+ * 1. Runs the optional `prepareAttributes` callback to enrich attributes
+ *    with type-safe access to schema-declared fields
+ * 2. Calls the schema resolver to compute promoted field values
+ * 3. Validates the resolved values against the Zod shape
+ * 4. Stores promoted fields on `metadata.schemaFields`
+ * 5. Removes promoted keys from `metadata.attributes` to avoid duplication
+ *
+ * When no schema is configured, the stream passes through unchanged.
+ *
+ * @param document$ - Stream of document batches from the metadata enrichment step.
+ * @param schema - The index schema config, if defined. When `undefined`, documents pass through unchanged.
+ * @returns Stream of document batches with promoted fields resolved and stored.
+ */
+export function applySchema(
+  document$: Observable<VectorStoreDocument[]>,
+  schema: IndexSchemaConfig | undefined,
+): Observable<VectorStoreDocument[]> {
+  // No schema configured — pass through unchanged (backward compatible)
+  if (!schema) {
+    return document$;
+  }
+
+  const promotedKeys = new Set(Object.keys(schema.shape.shape as Record<string, unknown>));
+
+  return document$.pipe(
+    map((documents) =>
+      documents.map((document) => {
+        // Run typed attribute processor before schema resolution so the
+        // resolver receives fully enriched attributes
+        let enrichedDocument = document;
+        if (schema.prepareAttributes) {
+          const enrichedAttributes = schema.prepareAttributes(
+            (document.metadata.attributes ?? {}) as Record<string, unknown>,
+            document,
+          );
+          enrichedDocument = {
+            ...document,
+            metadata: {
+              ...document.metadata,
+              attributes: enrichedAttributes as Record<string, unknown>,
+            },
+          };
+        }
+
+        // Resolve promoted field values from the fully enriched document
+        const resolved = schema.resolve(enrichedDocument);
+
+        // Validate against the Zod shape — throws on invalid data with
+        // a clear error message pointing to the offending field
+        const validated = schema.shape.parse(resolved) as Record<string, unknown>;
+
+        // Remove promoted keys from attributes to avoid storing them
+        // in both top-level fields and the generic attributes array
+        const currentAttributes = (enrichedDocument.metadata.attributes ?? {}) as Record<
+          string,
+          unknown
+        >;
+        const remainingAttributes: Record<string, unknown> = {};
+        for (const [key, value] of Object.entries(currentAttributes)) {
+          if (!promotedKeys.has(key)) {
+            remainingAttributes[key] = value;
+          }
+        }
+
+        return {
+          ...enrichedDocument,
+          metadata: {
+            ...enrichedDocument.metadata,
+            attributes: remainingAttributes,
+            schemaFields: validated,
+          },
+        };
+      }),
+    ),
+  );
+}

package/src/bin/delete-removed-files.ts

@@ -33,7 +33,7 @@ export function createDeleteRemovedFilesStream(
         console.log('Removing entry from vector store', file.relativePath);
       }
       if (!options.dryRun) {
-        const vectorStoreService = framework.ai.getService('search', options.azureSearchIndexName);
+        const vectorStoreService = framework.ai.useIndex(options.indexName);
         // Single batch deletion - one file can produce multiple document chunks
         await vectorStoreService.deleteDocuments({
           filter: { filterExpression: filterExpression ?? undefined },