[RFC] Support pass query string to add in model input in ml inference search response processor

[mingshl]

Summary

Currently, the ML inference search response processor allows passing document fields as input to the machine learning model. However, in certain use cases, such as re-ranking, it is necessary to include the query text as part of the model input along with the document fields. This RFC proposes adding support for passing the query string as input to the ML inference search response processor.

Motivation

In re-ranking use cases, the machine learning model often needs to consider the query text in addition to the document fields to produce an accurate ranking score. By including the query text as input, the model can better understand the context and relevance of the documents to the query, leading to improved ranking results.

Proposed Solution

We propose adding a new configuration option to the model_config section of the ML inference search response processor. This option would allow specifying the query string as an input to the machine learning model.

Example configuration:

PUT /_search/pipeline/my_pipeline
{
  "response_processors": [
    {
      "ml_inference": {
        "tag": "ml_inference",
        "description": "This processor is going to run ml inference during search request",
        "model_id": "JtoA55ABzSn4tHMg5RDk",
        "function_name": "REMOTE",
        "input_map": [
          {
            "sentences": "dairy"
          }
        ],
        "output_map": [
          {
            "rank_score": "$.response.score"
          }
        ],
        "model_config": {
          "query_text": "$.query.term.dairy.value"
        },
        "ignore_missing": false,
        "ignore_failure": false
      }
    }
  ]
}

In the example above, the model_config section includes a new query_text field, which is set to the value of the path of query.term.dairy.value This template will be resolved to the actual query string during the search request.

Implementation Details

The implementation would involve the following steps:

Extend the model_config section in the ML inference search response processor to accept a query_text field.
During the search request processing, resolve the query_text json pat (if provided) to the actual query string.
Include the resolved query string as part of the input to the machine learning model, along with the other input fields specified in the input_map.
Backward Compatibility

This change should not break backward compatibility, as it introduces a new optional configuration option. Existing configurations without the query_text field will continue to work as before.

Security Considerations

There are no significant security considerations for this change, as it only involves passing the query string as input to the machine learning model, which is already part of the search request.

Adding support for passing the query string as input to the ML inference search response processor will enable better re-ranking capabilities by allowing machine learning models to consider the query context along with the document fields. This enhancement will improve the relevance of search results in re-ranking use cases.

Out of Scope

Currently ml inference search response processor doesn't support rescore and sorting functionality, a separate RFC is raised to address these functionality
opensearch-project/OpenSearch#15631

Related Issue

[META]#2878

[dhrubo-os]

So as a customer do I need to remember this query.term.dairy.value? Can you describe more about this value? Can customer give any other values? What are the other possible values for query_text?

[dhrubo-os]

May be we should also provide a detailed example output before and after providing this query_text for clearer understanding.

[mingshl]

So as a customer do I need to remember this query.term.dairy.value? Can you describe more about this value? Can customer give any other values? What are the other possible values for query_text?

the exact json path query.term.dairy.value will be able to look up the 'query_text',
but because we support json path look up, we can also use partial path to look up any json blob that contain value
for example, we can use the json path $..value

This is more for rerank use case, that we need the query_text to compare with the search documents to get text similarity score. for example. cross encoder model

{
    "query_text": "today is sunny",
    "text_docs": [
        "how are you",
        "today is sunny",
        "today is july fifth",
        "it is winter"
    ]
}

[mingshl]

@ylwu-amzn you proposed an idea to use search extension, so the search pipeline will carry the configuration and read from search extension,

"model_config": {
          "query_text": "$.ext.ml_inference_params.query_text"
        }

the pipeline config would be:

PUT /_search/pipeline/my_pipeline
{
  "response_processors": [
    {
      "ml_inference": {
        "tag": "ml_inference",
        "description": "This processor is going to run ml inference during search request",
        "model_id": "JtoA55ABzSn4tHMg5RDk",
        "function_name": "REMOTE",
        "input_map": [
          {
            "sentences": "dairy"
          }
        ],
        "output_map": [
          {
            "rank_score": "$.response.score"
          }
        ],
        "model_config": {
          "query_text": "$.ext.ml_inference_params.query_text"
        },
        "ignore_missing": false,
        "ignore_failure": false
      }
    }
  ]
}

the search request needs to always carry search extensions, or it won't find the model input for query_text,

this is the sample query that would work for this proposal:

GET /demo-index-1/_search?search_pipeline=my_sentimental_pipeline
{
  "query": {
    "match": {
      "label": "happy moments"
    }
  },
  "ext": {
    "ml_inference_params": {
      "query_text": "query.match.label"
    }
  }
}

the pros of this approach:

1. do not need to update search pipeline for different types of queries

the cons of this approach:

1. always need a search extension.
2. two times of writing json path, one time in pipeline configurations ext.ml_inference_params.query_text, another time in search extension query.match.label

[andrross]

[Catch All Triage - 1, 2, 3]

[austintlee]

@ylwu-amzn @mingshl

Here is an example from: https://opensearch.org/docs/latest/search-plugins/hybrid-search/

GET /my-nlp-index/_search?search_pipeline=nlp-search-pipeline
{
  "_source": {
    "exclude": [
      "passage_embedding"
    ]
  },
  "query": {
    "hybrid": {
      "queries": [
           {
             "match":{
                 "passage_text": "hello"
              }
           },
           {
             "term":{
              "passage_text":{
                 "value":"planet"
              }
             }
           }
      ]
    }
  }
}

Queries can be very complex. This is the reason I abandoned this approach for the RAG search processor and instead have a separate llmQuestion (or llmMessages) parameter that people can pass to the LLM.

I also think a lot of this work seems redundant and is already built into the RAG response processor...