Update semantic search CCS to support generic query vector builders#142254
Conversation
…erence validation logic
…nce validation logic
|
@elasticmachine update branch |
server/src/main/java/org/elasticsearch/search/vectors/QueryVectorBuilder.java
Outdated
Show resolved
Hide resolved
| // Other query vector builder types always require validation | ||
| queryVectorBuilder.validate(); |
There was a problem hiding this comment.
We don't need this new interface. You still do the instanceof check before it, it isn't really providing anything.
There was a problem hiding this comment.
Hmmm, I think I see how to do that. Some special handling would be required for TextEmbeddingQueryVectorBuilder, but all other query vector builders would be assumed to be complete and able to build vectors.
There was a problem hiding this comment.
The only special handling is that we inject the model to TextEmbeddingQueryVectorBuilder right? If thats the case, that is the only thing that ever requires special handling. Everything else should be treated as "it should just work and if you don't give us a vector, shame on you"
There was a problem hiding this comment.
Essentially yes. The special handling for TextEmbeddingQueryVectorBuilder is that if a model ID isn't specified, we pull the model text from it and then infer the inference ID(s) from the semantic text fields queried.
All other query vector builders should "just work" though, and if they don't, they should throw an error when buildVector is called.
There was a problem hiding this comment.
++ yes, I think this is good. I think the logic is that "rewrite as normal if NOT TextEmbeddingQueryVectorBuilder and if you are, let's handle our special case"
|
Pinging @elastic/es-search-relevance (Team:Search Relevance) |
benwtrent
left a comment
There was a problem hiding this comment.
@Mikep86 its difficult to parse what is "refactoring" vs. actually fixing the bug. Could you extract the refactoring pieces into their own PR, that don't adjust behavior at all? This SEEMS like the new "QueryRewriteAction" things.
|
@benwtrent Sure, but the refactoring pieces would just be moving the |
| QueryVectorBuilder queryVectorBuilder = originalQuery.queryVectorBuilder(); | ||
| if (queryVectorBuilder != null) { | ||
| boolean registerAction = false; | ||
| if (queryVectorBuilder instanceof TextEmbeddingQueryVectorBuilder tevb) { | ||
| // TextEmbeddingQueryVectorBuilder is a special case. If a model ID is set, we register an action to generate | ||
| // the query vector. If not, the model text will be returned via getQuery() so that InferenceQueryUtils can | ||
| // generate the appropriate inference results for the inferred inference ID(s). | ||
| if (tevb.getModelId() != null) { | ||
| registerAction = true; | ||
| } | ||
| } else { | ||
| // We register an action to generate the query vector for all other query vector builders. If they cannot, buildVector() | ||
| // should throw an error indicating why. | ||
| registerAction = true; | ||
| } | ||
|
|
||
| if (registerAction) { | ||
| SetOnce<float[]> newQueryVectorSupplier = new SetOnce<>(); | ||
| queryRewriteContext.registerUniqueAsyncAction( | ||
| new QueryVectorBuilderAsyncAction(queryVectorBuilder), | ||
| newQueryVectorSupplier::set | ||
| ); | ||
| return new InterceptedInferenceKnnVectorQueryBuilder(queryBuilder, originalQuery, newQueryVectorSupplier); | ||
| } | ||
| } | ||
|
|
||
| return queryBuilder; | ||
| } |
benwtrent
left a comment
There was a problem hiding this comment.
looking at the knn specific things. This looks good.
I don't know anything about the changes to the InferenceQueryUtils class. But seems fall out from getInferenceIdOverride. I am not sure why we ever needed that interface. Where does the logic reside now? Was it ever needed for the fqdn inference id?
Inference ID override was how we used to handle when the user provided an inference ID to the the query vector builder or This has been simplified with this updated implementation. Now, |
|
@elasticmachine update branch |
The current semantic search CCS implementation assumes that the only
QueryVectorBuilderimplementation isTextEmbeddingQueryVectorBuilder, However, this will soon not be the case. There is already a PR to add a new "lookup" query vector builder, and more variants may be added later.This PR updates the semantic search CCS implementation to support generic query vector builders. The high-level approach is:
QueryVectorBuildercan be used to build a query vector on the coordinator node. If it cannot, the query vector builder'sbuildVectormethod is responsible for throwing an exception.knnqueries to rewrite a query vector builder to a query vector on the coordinator node. The resulting query vector is stored in the original query that was intercepted.TextEmbeddingQueryVectorBuilders still need special handling, because they can be used for semantic search even when incomplete. When aTextEmbeddingQueryVectorBuilderdoes not specify an inference ID, we extract the model text from it and use that to generate the inference results necessary to query the specifiedsemantic_textfield(s).This approach also allows us to simplify the intercepted query logic overall. We can remove the concept of a general "inference ID override" and replace it with custom coordinator node rewrite actions that generate query vectors as necessary.