[FEATURE] Enable to use passage chunks from hybrid neural search result as RAG input

[reuschling]

I have implemented a hybrid search with according ingest and search pipeline, using text embeddings on document chunks, as the embedding models have input token size limitations of course.

The ingest pipeline follows https://opensearch.org/docs/latest/search-plugins/text-chunking/

The top results should now be used as input for RAG, I configured a search pipeline for this, following https://opensearch.org/docs/latest/search-plugins/conversational-search/ :

{
		"description": "Post and response processor for hybrid search and RAG",
		"phase_results_processors": [
			{
				"normalization-processor": {
					"normalization": {
						"technique": "min_max"
					},
					"combination": {
						"technique": "arithmetic_mean",
						"parameters": {
							"weights": [
								0.3,
								0.7
							]
						}
					}
				}
			}
		],
		"response_processors": [
			{
				"retrieval_augmented_generation": {
					"tag": "rag_pipeline",
					"description": "Pipeline using the configured LLM",
					"model_id": "{{ _.LlmId }}",
					"context_field_list": [
						"tns_body"
					],
					"system_prompt": "You are a helpful assistant",
                                        "user_instructions": "Generate a concise and informative answer in less than 100 words for the given question"
				}
			}
		]
}

Now I am able to send a search request using this search pipeline to {{ _.openSearchUrl }}/{{ _.openSearchIndex }}/_search?search_pipeline=hybrid-rag-pipeline , which is working:

{
	"_source": {
		"include": [
			"tns_body"
		]
	},
	"query": {
		"hybrid": {
			"queries": [
				{
					"bool": {
						"should": [
							{
								"match": {
									"tns_body": {
										"query": "${{query}}"
									}
								}
							},
							{
								"match": {
									"tns_body.ngram": {
										"query": "${{query}}"
									}
								}
							}
						]
					}
				},
				{
					"nested": {
						"score_mode": "max",
						"path": "embedding_chunked_512_paragraphs_chunks_tns_body",
						"query": {
							"neural": {
								"embedding_chunked_512_paragraphs_chunks_tns_body.knn": {
									"query_text": "{{ _.query }}",
									"model_id": "{{ _.embeddingsModelId }}",
									"k": 10
								}
							}
						}
					}
				}
			]
		}
	},
	"ext": {
		"generative_qa_parameters": {
			"llm_question": "${{query}}",
			"llm_model": "llama3",
			"context_size": 2,
			"message_size": 5,
			"timeout": 150
		}
	}
}

Now I am falling into the issue that the documents in my index are too long for my LLM input. In OpenSearch currently the context_size and the message_size is configurable, but when the first document exceeds the input token limit, OpenSearch sends a message to the LLM provider that can not be processed.

Two things comes into my mind now:

1. Add a 'max token' / 'max terms' config parameter parallel to context_size and message_size, that not depend on the document sizes.
2. Because the documents are such big, and it is working fine for embeddings with chunking, why not enable to use the matching chunks for RAG input instead of the whole documents. If the chunks are too small (because the embeddings input size is small) I can also think of some scripted field that considers the matching chunk together with neighbor chunks. Or maybe it would be possible to specify an according text field to a generated embedding vector at ingest time, that can comprise the neighbor chunks also. This then could be used for RAG afterwards.

Currently big documents are not only silently lost in RAG. Because the whole prompt exceeds the input token limit of the LLM, it is (in my setting at least) accidentally truncated, meaning that the question - which is the last part of the generated prompt - is lost. So the user question will not be answered at all.

[ylwu-amzn]

@reuschling, Can you share some sample data for your index ? I think you should save text chunks together with embeddings to index

[reuschling]

This is how a document looks like, returned by an empty query. I truncated the embedding vectors for better readability. I use dynamic mappings, thus the prefixes in the field names.

Do you mean to just save text chunks instead of the whole documents? In this case I can not search for the whole documents anymore. The hybrid search searches inside the whole documents in the 'classical' query part, only the embeddings rely on chunks, because the model input length is limited, and the semantic representation of the embeddings may become too generalized, depending on the used model.

If I chunk the body additionally regarding LLM input token size with an ingest pipeline, I get nested fields 'bodyChunk4LLM' to the origin document, as it is currently for the paragraph and paragraph->chunkSize4Embedding fields.

In the hybrid query with RAG postprocessing, following would be the situation:

1. Classy part of hybrid query searches inside origin body or in one of the three chunk fields (bodyChunk4LLM, paragraph, paragraph->chunkSize4Embedding)
2. Neuralsearch part of hybrid query will search inside the paragraph->chunkSize4Embedding field
3. Then the matched passages from bodyChunk4LLM as RAG input are needed. But because these are nested fields, I can not know which one it is. Further, the fields matched inside the hybrid query are normally other fields, and there exists no relationship. And the returned format must be valid input for the RAG processor.

I tried to generate a scripted field inside the search query where neural search returns the matched chunk offset with "inner_hits". But "inner_hits" are not returned with hybrid queries, and scripted fields doesn't have access to the inner_hits return values, only to the fields inside the origin document. I also doubt that the returned format would not be valid for the RAG processor, as the scripted fields are not appear in _source.

The second possibility would be to generate additional documents with one single bodyChunk4LLM field, i.e. for each input document N new documents with chunks and a reference to the origin doc. Then using these documents for RAG, and the big origin documents for the other searches (classy, neural + hybrid). But the chunk ingest processor doesn't generate separate documents, only new fields with N values inside the origin document instead. I don't see an ingest processor who can do this, maybe you have a hint? 😄

Third, I could generate both documents (origin and N according chunks) outside OpenSearch, but this means to change code inside all possible document providers, where we often don't have access to. It would be much better if this could be done entirely inside OpenSearch configs with e.g. ingest pipelines.

Also there would be a huge redundancy in terms of data size, because all embeddings, chunks, etc. must be generated on both documents, origin and bodyChunk4LLM, in order to support hybrid neural search in both cases. But this is another story.

"hits": [
    {
        "_index": "testfiles",
        "_id": "/home/reuschling/projectz/leech/resource/testData/example_files/HTML.html",
        "_score": 1.0,
        "_source": {
            "tns_description": "",
            "tk_source": "file:/home/reuschling/projectz/leech/resource/testData/example_files/HTML.html",
            "paragraphs_tns_title": [
                "Laptop power supplies are available in First Class only"
            ],
            "resourceName": "/home/reuschling/projectz/leech/resource/testData/example_files/HTML.html",
            "dataEntityContentFingerprint": "1330962029000",
            "paragraphs_tns_description": [],
            "paragraphs_tns_body": [
                "\n    Code, Write, Fly\n\n",
                "    This chapter is being written 11,000 meters above New Foundland. \n  "
            ],
            "tns_title": "Laptop power supplies are available in First Class only",
            "date_modified": "2012.03.05 16:40:29:000",
            "tns_body": "\n    Code, Write, Fly\n\n    This chapter is being written 11,000 meters above New Foundland. \n  ",
            "embedding_chunked_512_paragraphs_chunks_tns_body": [
                {
                    "knn": [
                        -0.016918883,
                        -0.0045324834,
                        ...
                    ]
                },
                {
                    "knn": [
                        0.03815532,
                        0.015174329,
                        ...
                    ]
                }
            ],
            "X-TIKA:Parsed-By": [
                "org.apache.tika.parser.CompositeParser",
                "de.dfki.km.leech.parser.HtmlCrawlerParser"
            ],
            "Content-Encoding": "windows-1252",
            "dataEntityId": "/home/reuschling/projectz/leech/resource/testData/example_files/HTML.html",
            "paragraphs_chunks_tns_title": [
                "Laptop power supplies are available in First Class only"
            ],
            "paragraphs_chunks_tns_body": [
                "\n    Code, Write, Fly\n\n",
                "    This chapter is being written 11,000 meters above New Foundland. \n  "
            ],
            "embedding_512_tns_body": [
                0.021791747,
                0.0016991429,
                ...
            ],
            "Content-Type": "text/html; charset=windows-1252",
            "paragraphs_chunks_tns_description": []
        }
    }
]

[yuye-aws]

The second possibility would be to generate additional documents with one single bodyChunk4LLM field, i.e. for each input document N new documents with chunks and a reference to the origin doc. Then using these documents for RAG, and the big origin documents for the other searches (classy, neural + hybrid). But the chunk ingest processor doesn't generate separate documents, only new fields with N values inside the origin document instead. I don't see an ingest processor who can do this, maybe you have a hint? 😄

You are right. To be the best of my knowledge, none of the ingest processor can return multiple documents. Can you take a look at logstash? After the text chunking and embedding, maybe you can use logstash to flatten the chunked list and embeddings into multiple documents.

[yuye-aws]

From my understanding, you only need the specific chunk, instead of the whole document. You use need a way for nested query to tell you which chunk gets matched. Please correct me if I am wrong. @reuschling

[reuschling]

yes exactly @yuye-aws , because the LLM has a limited text length that it can process. But there is still the problem that the chunks for the embedding model, that are matched against in the nested query, have normally other sizes as the chunks needed for the LLM, which are in general much bigger. Thus my though to get the neighbor chunks also.

So, in terms of control over the whole searching process, I would now assume it would be best if there is a possibility to have control over the LLM chunk size, which is currently only possible over the field content size. This can now achieved only by generating separate documents with LLM related chunk sizes, created outside (logstash, other document providers) or inside OpenSearch (ingest processor).

By far preferable would be inside OpenSearch - there exists millions of existing document providers in parallel to logstash, and they all have to be adjusted otherwise. I personally would look now if there is a possibility to create an own ingest processor (maybe script based?) that would create chunks like the current text chunking processor, but creates separate documents instead of additional fields instead.

The idea to use logstash as a postprocessing step is also not so bad, but it is not so easy to realize for a non-static index where new document content is added frequently.

[yuye-aws]

So, in terms of control over the whole searching process, I would now assume it would be best if there is a possibility to have control over the LLM chunk size, which is currently only possible over the field content size. This can now achieved only by generating separate documents with LLM related chunk sizes, created outside (logstash, other document providers) or inside OpenSearch (ingest processor).

Aligning the chunk size with the LLM is a good practice for search relevance. I have created an RFC on model-based tokenizer: opensearch-project/neural-search#794. Do you think you concern will be addressed if we can register the same tokenizer as the LLM?

By the way, I have listed a few options under the RFC, along with their pros and cons. Would you like to share your valuable opinions? Any comments will be appreciated.

[yuye-aws]

I personally would look now if there is a possibility to create an own ingest processor (maybe script based?) that would create chunks like the current text chunking processor, but creates separate documents instead of additional fields instead.

I am afraid not. The ingest processor in OpenSearch only performs certain actions upon a single document. We cannot create multiple documents based on a single document.

@ylwu-amzn Can you provide some suggestions? Maybe we can have a discussion together.

[yuye-aws]

I tried to generate a scripted field inside the search query where neural search returns the matched chunk offset with "inner_hits". But "inner_hits" are not returned with hybrid queries, and scripted fields doesn't have access to the inner_hits return values, only to the fields inside the origin document. I also doubt that the returned format would not be valid for the RAG processor, as the scripted fields are not appear in _source.

How about supporting inner_hits in hybrid queries? Feel free to create an RFC so that we can discuss next-step plans.

[reuschling]

How about supporting inner_hits in hybrid queries? Feel free to create an RFC so that we can discuss next-step plans.

There is already a feature request for this: opensearch-project/neural-search#718.
Nevertheless, the follow-up problems still would resist (scripted fields doesn't have access to inner_hits, and rag processor can not deal with scripted fields, as far as I know).

Otherwise this could be also a valid solution, because there would be no need for extra chunking for the LLM anymore. It would be possible to build the chunk for the LLM out of the matched embedding chunk with its neighbor chunks. The only lack is that the match would be only rely on the embedding part of the hybrid query.

[reuschling]

Aligning the chunk size with the LLM is a good practice for search relevance. I have created an RFC on model-based tokenizer: opensearch-project/neural-search#794. Do you think you concern will be addressed if we can register the same tokenizer as the LLM?

This is a better way to specify valid chunk lengths for the current model, right? I currently use the formula terms*0,75≈tokenlimit, as documented at https://opensearch.org/docs/latest/ingest-pipelines/processors/text-chunking/. I think this is valid, with a specified overlap it feels to me that there is no relevant information loss. Specifying the real token limit as token limit (as it is known from the model description) would be much easier of course. But I am not sure if every (future) model follows the same rules for tokenization. If not, things would become complex against a bit redundancy in the data with the current approach. I will have a look on your RFC, thanks for the hint.

[yuye-aws]

Aligning the chunk size with the LLM is a good practice for search relevance. I have created an RFC on model-based tokenizer: opensearch-project/neural-search#794. Do you think you concern will be addressed if we can register the same tokenizer as the LLM?

This is a better way to specify valid chunk lengths for the current model, right? I currently use the formula terms*0,75≈tokenlimit, as documented at https://opensearch.org/docs/latest/ingest-pipelines/processors/text-chunking/. I think this is valid, with a specified overlap it feels to me that there is no relevant information loss. Specifying the real token limit as token limit (as it is known from the model description) would be much easier of course. But I am not sure if every (future) model follows the same rules for tokenization. If not, things would become complex against a bit redundancy in the data with the current approach. I will have a look on your RFC, thanks for the hint.

Replied in opensearch-project/neural-search#794

[yuye-aws]

There is already a feature request for this: opensearch-project/neural-search#718.
Nevertheless, the follow-up problems still would resist (scripted fields doesn't have access to inner_hits, and rag processor can not deal with scripted fields, as far as I know).

I see. Supporting inner hits from hybrid query does not suffice to resolve your problem. It may take us sometime to investigate and check valid solutions to your problem. Thanks for your patience.

If we enable the nested query to return the specific chunk. Does that resolve your problem?

[yuye-aws]

For your information: https://github.com/opensearch-project/OpenSearch/blob/main/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java