Merge pull request #484 from M4xymm/RDBC-918

RDBC-918 - Allow to query vector searching using a document id

Author: M4xymm

src/Documents/Queries/VectorSearch/Common/VectorFieldBase.ts

@@ -1,12 +1,11 @@
-import { IVectorEmbeddingField, IVectorEmbeddingFieldFactoryAccessor } from "../../../Session/IVectorFieldFactory.js";
-import { VectorFieldBase } from "../Common/VectorFieldBase.js";
+import { IVectorEmbeddingField, IVectorEmbeddingFieldFactoryAccessor } from "../../../Session/VectorFieldFactory.js";
 import {VectorEmbeddingType} from "../VectorEmbeddingType.js";
 import { Field } from "../../../../Types/index.js";
 
-export class VectorEmbeddingField<T> extends VectorFieldBase<T> implements
+export class VectorEmbeddingField<T> implements
     IVectorEmbeddingField, 
     IVectorEmbeddingFieldFactoryAccessor<T> {
-
+    public fieldName: string;
     public sourceQuantizationType: VectorEmbeddingType;
     public destinationQuantizationType: VectorEmbeddingType;
     public isBase64Encoded: boolean;
@@ -15,20 +14,12 @@ export class VectorEmbeddingField<T> extends VectorFieldBase<T> implements
     constructor(fieldName: Field<T>,
                 sourceQuantizationType: VectorEmbeddingType = "Single",
                 isBase64Encoded: boolean = false) {
-        super(fieldName);
+        this.fieldName = fieldName;
         this.sourceQuantizationType = sourceQuantizationType;
         this.destinationQuantizationType = sourceQuantizationType;
         this.isBase64Encoded = isBase64Encoded;
-        this.updateFieldName();
     }
 
-    private updateFieldName(): void {
-        this.fieldName = this.getFormattedFieldName(
-            this.rawFieldName,
-            this.sourceQuantizationType,
-            this.destinationQuantizationType
-        );
-    }
 
     public targetQuantization(targetEmbeddingQuantization: VectorEmbeddingType): IVectorEmbeddingField {
         if (targetEmbeddingQuantization === "Text") {
@@ -43,7 +34,6 @@ export class VectorEmbeddingField<T> extends VectorFieldBase<T> implements
             throw new Error(`Cannot quantize already quantized embeddings. Source VectorEmbeddingType is ${this.sourceQuantizationType}; however the destination is ${this.destinationQuantizationType}.`);
         }
 
-        this.updateFieldName();
         return this;
     }
 }

src/Documents/Queries/VectorSearch/Fields/VectorEmbeddingField.ts

@@ -1,38 +1,31 @@
-import { IVectorEmbeddingFieldFactoryAccessor, IVectorEmbeddingTextField } from "../../../Session/IVectorFieldFactory.js";
-import { VectorFieldBase } from "../Common/VectorFieldBase.js";
-import {VectorEmbeddingType} from "../VectorEmbeddingType.js";
+import {
+    IVectorEmbeddingFieldFactoryAccessor,
+    IVectorEmbeddingTextField
+} from "../../../Session/VectorFieldFactory.js";
+import { VectorEmbeddingType } from "../VectorEmbeddingType.js";
 import { Field } from "../../../../Types/index.js";
 
-export class VectorEmbeddingTextField<T> extends VectorFieldBase<T> implements
+export class VectorEmbeddingTextField<T> implements
     IVectorEmbeddingTextField, 
     IVectorEmbeddingFieldFactoryAccessor<T> {
 
+    public fieldName: string;
     public sourceQuantizationType: VectorEmbeddingType = "Text";
     public destinationQuantizationType: VectorEmbeddingType = "Single";
     public isBase64Encoded: boolean = false;
     public embeddingsGenerationTaskIdentifier: string = "";
 
     constructor(fieldName: Field<T>) {
-        super(fieldName);
-        this.updateFieldName();
+        this.fieldName = fieldName;
     }
 
-    private updateFieldName(): void {
-        this.fieldName = this.getFormattedFieldName(
-            this.rawFieldName,
-            this.sourceQuantizationType,
-            this.destinationQuantizationType,
-            this.embeddingsGenerationTaskIdentifier
-        );
-    }
 
     public targetQuantization(targetEmbeddingQuantization: VectorEmbeddingType): IVectorEmbeddingTextField {
         if (targetEmbeddingQuantization === "Text") {
             throw new Error("Cannot quantize the embedding to Text. This option is only available for sourceQuantizationType.");
         }
 
         this.destinationQuantizationType = targetEmbeddingQuantization;
-        this.updateFieldName();
         return this;
     }
 
@@ -42,7 +35,6 @@ export class VectorEmbeddingTextField<T> extends VectorFieldBase<T> implements
         }
 
         this.embeddingsGenerationTaskIdentifier = embeddingsGenerationTaskIdentifier;
-        this.updateFieldName();
         return this;
     }
 }

src/Documents/Queries/VectorSearch/Fields/VectorEmbeddingTextField.ts

@@ -1,11 +1,10 @@
-import { IVectorField } from "../../../Session/IVectorFieldFactory.js";
-import { VectorFieldBase } from "../Common/VectorFieldBase.js";
+import { IVectorField } from "../../../Session/VectorFieldFactory.js";
 import { Field } from "../../../../Types/index.js";
 
-export class VectorField<T> extends VectorFieldBase<T> implements IVectorField {
+export class VectorField<T> implements IVectorField {
+    public fieldName: string;
+
     constructor(fieldName: Field<T>) {
-        super(fieldName);
-        this._byFieldMethodUsed = true;
         this.fieldName = fieldName;
     }
 }

src/Documents/Queries/VectorSearch/Fields/VectorField.ts

@@ -3,7 +3,7 @@ import {
     IVectorEmbeddingTextField,
     IVectorField,
     IVectorFieldFactory
-} from "../../Session/IVectorFieldFactory.js";
+} from "../../Session/VectorFieldFactory.js";
 import { VectorField } from "./Fields/VectorField.js";
 import { VectorEmbeddingField } from "./Fields/VectorEmbeddingField.js";
 import { VectorEmbeddingTextField } from "./Fields/VectorEmbeddingTextField.js";

src/Documents/Queries/VectorSearch/VectorEmbeddingFieldFactory.ts

@@ -94,10 +94,11 @@ import {
     IVectorEmbeddingFieldFactoryAccessor, IVectorEmbeddingTextField,
     IVectorField,
     IVectorFieldFactory, IVectorFieldValueFactory, VectorEmbeddingFieldValueFactory
-} from "./IVectorFieldFactory.js";
+} from "./VectorFieldFactory.js";
 import { VectorEmbeddingFieldFactory } from "../Queries/VectorSearch/VectorEmbeddingFieldFactory.js";
 import { Field } from "../../Types/index.js";
 import { JsonSerializer } from "../../Mapping/Json/Serializer.js";
+import { VectorSearchToken } from "./Tokens/VectorSearchToken.js";
 
 /**
  * A query against a Raven index
@@ -2577,56 +2578,63 @@ export abstract class AbstractDocumentQuery<T extends object, TSelf extends Abst
     ) {
         this._assertMethodIsCurrentlySupported("vectorSearch");
 
+        const fieldAccessor = this._resolveVectorSearchFieldAccessor(fieldName);
+        const {value, isDocumentId} = this._resolveVectorSearchValueFactory(valueOrFactory);
+
+        const tokens = this._getCurrentWhereTokens();
+        this._appendOperatorIfNeeded(tokens);
+        this._negateIfNeeded(tokens, fieldAccessor.fieldName);
+
+        const sourceQuantizationType = VectorSearchToken.getSourceQuantizationType(fieldAccessor);
+        const targetQuantizationType = VectorSearchToken.getTargetQuantizationType(fieldAccessor);
+        const taskIdentifier = VectorSearchToken.getTaskIdentifier(fieldAccessor);
+
+        const parameterName = this._addQueryParameter(value);
+
+        const vectorSearchToken = new VectorSearchToken(
+            fieldAccessor.fieldName,
+            parameterName,
+            sourceQuantizationType,
+            targetQuantizationType,
+            options?.similarity || null,
+            options?.numberOfCandidates || null,
+            options?.isExact || VectorSearchToken.DEFAULT_IS_EXACT,
+            isDocumentId,
+            taskIdentifier
+        );
+
+        tokens.push(vectorSearchToken);
+    }
+
+
+    private _resolveVectorSearchFieldAccessor(fieldName: Field<T> | ((factory: IVectorFieldFactory<T>) => IVectorField | IVectorEmbeddingField | IVectorEmbeddingTextField)): IVectorEmbeddingFieldFactoryAccessor<T> {
         const vectorFactory = new VectorEmbeddingFieldFactory<T>();
-        let fieldAccessor: IVectorEmbeddingFieldFactoryAccessor<T>;
-    
+
         if (typeof fieldName === "string") {
-            fieldAccessor = vectorFactory.withField(fieldName) as IVectorEmbeddingFieldFactoryAccessor<T>;
+            return vectorFactory.withField(fieldName) as IVectorEmbeddingFieldFactoryAccessor<T>;
         } else if (typeof fieldName === "function") {
-            fieldAccessor = fieldName(vectorFactory) as IVectorEmbeddingFieldFactoryAccessor<T>;
+            return fieldName(vectorFactory) as IVectorEmbeddingFieldFactoryAccessor<T>;
         } else {
-            throwError("InvalidArgumentException",
-                "fieldName must be either a string or a function that selects a vector field");
+            throwError("InvalidArgumentException", "fieldName must be either a string or a function that selects a vector field");
         }
-    
-        const whereParams = new WhereParams();
-        whereParams.fieldName = fieldAccessor.fieldName;
+    }
 
-        // Handle value or valueFactory
+    private _resolveVectorSearchValueFactory(valueOrFactory: number[] | string | ((factory: IVectorFieldValueFactory) => void)) {
         if (typeof valueOrFactory === "function") {
             const fieldValueFactory = new VectorEmbeddingFieldValueFactory();
             valueOrFactory(fieldValueFactory);
 
-            if (fieldValueFactory.embeddings) {
-                whereParams.value = fieldValueFactory.embeddings;
-            } else if (fieldValueFactory.embedding) {
-                whereParams.value = fieldValueFactory.embedding;
-            } else if (fieldValueFactory.text) {
-                whereParams.value = fieldValueFactory.text;
-            } else if (fieldValueFactory.texts) {
-                whereParams.value = fieldValueFactory.texts;
-            } else {
+            const value = fieldValueFactory.embedding || fieldValueFactory.embeddings ||
+                fieldValueFactory.text || fieldValueFactory.texts || fieldValueFactory.byId;
+
+            if (!value) {
                 throwError("InvalidOperationException", "No value was provided in the valueFactory");
             }
+
+            return {value, isDocumentId: !!fieldValueFactory.byId};
         } else {
-            whereParams.value = valueOrFactory;
+            return {value: valueOrFactory, isDocumentId: false};
         }
-
-        whereParams.allowWildcards = true;
-        const transformToEqualValue = this._transformValue(whereParams);
-
-        const tokens = this._getCurrentWhereTokens();
-        this._appendOperatorIfNeeded(tokens);
-        this._negateIfNeeded(tokens, whereParams.fieldName);
-
-        const whereToken = WhereToken.create(
-            "VectorSearch",
-            whereParams.fieldName,
-            this._addQueryParameter(transformToEqualValue),
-            new WhereOptions({vectorSearch: options, exact: options?.isExact})
-        );
-
-        tokens.push(whereToken);
     }
 }
 

src/Documents/Session/AbstractDocumentQuery.ts

@@ -64,7 +64,7 @@ import {
     IVectorField,
     IVectorFieldFactory,
     IVectorFieldValueFactory
-} from "./IVectorFieldFactory.js";
+} from "./VectorFieldFactory.js";
 
 export const NESTED_OBJECT_TYPES_PROJECTION_FIELD = "__PROJECTED_NESTED_OBJECT_TYPES__";
 

src/Documents/Session/DocumentQuery.ts

@@ -27,7 +27,7 @@ import {
     IVectorEmbeddingTextField,
     IVectorField,
     IVectorFieldFactory, IVectorFieldValueFactory
-} from "./IVectorFieldFactory.js";
+} from "./VectorFieldFactory.js";
 
 /**
  * A query against a Raven index