Fix Vector KNN and Vector Store not using collection ID input port

Ironclad · May 9, 2024 · b75551b · b75551b
1 parent ea0d5e6
commit b75551b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/packages/core/src/model/nodes/VectorNearestNeighborsNode.ts b/packages/core/src/model/nodes/VectorNearestNeighborsNode.ts
@@ -12,8 +12,8 @@ import { type Inputs, type Outputs } from '../GraphProcessor.js';
 import { type InternalProcessContext } from '../ProcessContext.js';
 import { type DataValue, type EditorDefinition, type VectorDataValue } from '../../index.js';
 import { dedent } from 'ts-dedent';
-import { coerceTypeOptional } from '../../utils/coerceType.js';
 import { getIntegration } from '../../integrations/integrations.js';
+import { getInputOrData } from '../../utils/index.js';
 
 export type VectorNearestNeighborsNode = ChartNode<'vectorNearestNeighbors', VectorNearestNeighborsNodeData>;
 
@@ -143,19 +143,18 @@ export class VectorNearestNeighborsNodeImpl extends NodeImpl<VectorNearestNeighb
   }
 
   async process(inputs: Inputs, context: InternalProcessContext): Promise<Outputs> {
-    const integration = this.data.useIntegrationInput
-      ? coerceTypeOptional(inputs['integration' as PortId], 'string') ?? this.data.integration
-      : this.data.integration;
+    const integration = getInputOrData(this.data, inputs, 'integration');
     const vectorDb = getIntegration('vectorDatabase', integration, context);
 
-    const k = this.data.useKInput ? coerceTypeOptional(inputs['k' as PortId], 'number') ?? this.data.k : this.data.k;
+    const indexUrl = getInputOrData(this.data, inputs, 'collectionId');
+    const k = getInputOrData(this.data, inputs, 'k', 'number');
 
     if (inputs['vector' as PortId]?.type !== 'vector') {
       throw new Error(`Expected vector input, got ${inputs['vector' as PortId]?.type}`);
     }
 
     const results = await vectorDb.nearestNeighbors(
-      { type: 'string', value: this.data.collectionId },
+      { type: 'string', value: indexUrl },
       inputs['vector' as PortId] as VectorDataValue,
       k,
     );

diff --git a/packages/core/src/model/nodes/VectorStoreNode.ts b/packages/core/src/model/nodes/VectorStoreNode.ts
@@ -12,8 +12,9 @@ import { type Inputs, type Outputs } from '../GraphProcessor.js';
 import { type InternalProcessContext } from '../ProcessContext.js';
 import { type EditorDefinition, type VectorDataValue } from '../../index.js';
 import { dedent } from 'ts-dedent';
-import { coerceTypeOptional } from '../../utils/coerceType.js';
 import { getIntegration } from '../../integrations/integrations.js';
+import { getInputOrData } from '../../utils/index.js';
+import { coerceTypeOptional } from '../../utils/coerceType.js';
 
 export type VectorStoreNode = ChartNode<'vectorStore', VectorStoreNodeData>;
 
@@ -133,18 +134,17 @@ export class VectorStoreNodeImpl extends NodeImpl<VectorStoreNode> {
   }
 
   async process(inputs: Inputs, context: InternalProcessContext): Promise<Outputs> {
-    const integration = this.data.useIntegrationInput
-      ? coerceTypeOptional(inputs['integration' as PortId], 'string') ?? this.data.integration
-      : this.data.integration;
-
+    const integration = getInputOrData(this.data, inputs, 'integration');
     const vectorDb = getIntegration('vectorDatabase', integration, context);
 
+    const indexUrl = getInputOrData(this.data, inputs, 'collectionId');
+
     if (inputs['vector' as PortId]?.type !== 'vector') {
       throw new Error(`Expected vector input, got ${inputs['vector' as PortId]?.type}`);
     }
 
     await vectorDb.store(
-      { type: 'string', value: this.data.collectionId },
+      { type: 'string', value: indexUrl },
       inputs['vector' as PortId] as VectorDataValue,
       inputs['data' as PortId]!,
       {