Support for self-query with pgvector vector store

[vicotrbb]

Checked

• I searched existing ideas and did not find a similar one
• I added a very descriptive title
• I've clearly described the feature request and motivation for it

Feature request

Add support for self-query with pgvector vector store. Currently, the SelfQueryRetriever retriever tool doesn't work with pgvector.

Motivation

When trying to integrate SelfQueryRetriever with pgvector there is no working BaseTranslator working. I have tried with different ones searching for compatibility but all lead to the same issue while building the database query to search for the documents.

The issue is in the metadata query build, while translating the structure query into a working filter.

Proposal (If applicable)

I have wrote a class PGVectorTranslator extending from BaseTranslator that worked for my use case. With this, I managed to get my LLM to query the pgvector with dynamic metadata filters extracted and built from the query.

import {
  castValue,
  Comparator,
  Comparators,
  Comparison,
  isFilterEmpty,
  Operator,
  Operators,
  StructuredQuery,
} from 'langchain/chains/query_constructor/ir';
import { BaseTranslator } from 'langchain/retrievers/self_query';

export type Filter = { [k: string]: string | number };

export class PGVectorTranslator extends BaseTranslator {
  declare VisitOperationOutput: Filter;
  declare VisitComparisonOutput: Filter;
  declare VisitStructuredQueryOutput: { filter: Filter } | { [k: string]: never };

  allowedOperators: Operator[] = [Operators.and, Operators.or];
  allowedComparators: Comparator[] = [
    Comparators.eq,
    Comparators.ne,
    Comparators.gt,
    Comparators.gte,
    Comparators.lt,
    Comparators.lte,
  ];

  formatFunction(): string {
    throw new Error('Not implemented');
  }

  /**
   * Not required.
   */
  getComparatorFunction() {
    return null;
  }

  /**
   * Not required.
   */
  getOperatorFunction() {
    return null;
  }

  /**
   * Not required.
   */
  visitOperation() {
    return {};
  }

  /**
   * Visits the comparison part of a structured query and translates it into
   * a functional filter.
   * @param comparison The comparison part of a structured query.
   * @returns A function that takes a `Document` as an argument and returns a boolean based on the comparison.
   */
  visitComparison(comparison: Comparison): this['VisitComparisonOutput'] {
    const { comparator, attribute, value } = comparison;

    if (this.allowedComparators.includes(comparator)) {
      return { [attribute]: castValue(value) };
    } else {
      throw new Error('Comparator not allowed');
    }
  }

  /**
   * Visits a structured query and translates it into a functional filter.
   * @param query The structured query to translate.
   * @returns An object containing a `filter` property, which is a function that takes a `Document` as an argument and returns a boolean based on the structured query.
   */
  visitStructuredQuery(query: StructuredQuery): this['VisitStructuredQueryOutput'] {
    if (!query.filter) {
      return {};
    }

    const filterFunction = query.filter?.accept(this);
    return { filter: filterFunction as Filter };
  }

  /**
   * Merges two filters into one, based on the specified merge type.
   * @param defaultFilter The default filter function.
   * @param generatedFilter The generated filter function.
   * @param mergeType The type of merge to perform. Can be 'and', 'or', or 'replace'. Default is 'and'.
   * @returns A function that takes a `Document` as an argument and returns a boolean based on the merged filters, or `undefined` if both filters are empty.
   */
  mergeFilters(
    defaultFilter: Filter,
    generatedFilter: Filter,
    mergeType = 'and',
    forceDefaultFilter?: boolean,
  ): Filter {
    if (isFilterEmpty(defaultFilter) && isFilterEmpty(generatedFilter)) {
      return {};
    }

    if (isFilterEmpty(defaultFilter) || mergeType === 'replace') {
      if (isFilterEmpty(generatedFilter)) {
        return {};
      }

      return generatedFilter;
    }

    if (isFilterEmpty(generatedFilter)) {
      if (mergeType === 'and') {
        return {};
      }
      return defaultFilter;
    }

    if (forceDefaultFilter) {
      return defaultFilter;
    }

    if (mergeType === 'and') {
      return {
        ...defaultFilter,
        ...generatedFilter,
      };
    } else {
      throw new Error('Unknown merge type');
    }
  }
}

    const contextualziedRAGRetriever = SelfQueryRetriever.fromLLM({
      llm: llmModel,
      vectorStore: vectorStorage,
      documentContents,
      attributeInfo,
      verbose: isDevelopment(),
      structuredQueryTranslator: new PGVectorTranslator(),
      searchParams: {
        filter: {
          ...contextFilters,
        },
        forceDefaultFilter: false,
        mergeFiltersOperator: 'and',
        k: adjustKBasedOnQuery(attributes.userInput),
      },
    });

I believe this works as a baseline for the feature, however, I'm not sure if this implementation makes sense, giving the scope of a Translator class. In case this make sense, I can improve and finish it to submit it as a pull request if needed!