Resolve embed refs from query profile

1 files changed, 11 insertions, 10 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/schema/internal/TensorConverter.java b/container-search/src/main/java/com/yahoo/search/schema/internal/TensorConverter.java
index 94f92c7fd48..b1d6dbd5859 100644
--- a/container-search/src/main/java/com/yahoo/search/schema/internal/TensorConverter.java
+++ b/container-search/src/main/java/com/yahoo/search/schema/internal/TensorConverter.java
@@ -3,6 +3,7 @@ package com.yahoo.search.schema.internal;
 
 import com.yahoo.language.Language;
 import com.yahoo.language.process.Embedder;
+import com.yahoo.processing.request.Properties;
 import com.yahoo.tensor.Tensor;
 import com.yahoo.tensor.TensorType;
 
@@ -29,18 +30,18 @@ public class TensorConverter {
     }
 
     public Tensor convertTo(TensorType type, String key, Object value, Language language,
-                            Map<String, String> contextValues) {
+                            Map<String, String> contextValues, Properties properties) {
         var context = new Embedder.Context(key).setLanguage(language).setContextValues(contextValues);
-        Tensor tensor = toTensor(type, value, context);
+        Tensor tensor = toTensor(type, value, context, properties);
         if (tensor == null) return null;
         if (! tensor.type().isAssignableTo(type))
             throw new IllegalArgumentException("Require a tensor of type " + type);
         return tensor;
     }
 
-    private Tensor toTensor(TensorType type, Object value, Embedder.Context context) {
+    private Tensor toTensor(TensorType type, Object value, Embedder.Context context, Properties properties) {
         if (value instanceof Tensor) return (Tensor)value;
-        if (value instanceof String && isEmbed((String)value)) return embed((String)value, type, context);
+        if (value instanceof String && isEmbed((String)value)) return embed((String)value, type, context, properties);
         if (value instanceof String) return Tensor.from(type, (String)value);
         return null;
     }
@@ -49,7 +50,7 @@ public class TensorConverter {
         return value.startsWith("embed(");
     }
 
-    private Tensor embed(String s, TensorType type, Embedder.Context embedderContext) {
+    private Tensor embed(String s, TensorType type, Embedder.Context embedderContext, Properties properties) {
         if ( ! s.endsWith(")"))
             throw new IllegalArgumentException("Expected any string enclosed in embed(), but the argument does not end by ')'");
         String argument = s.substring("embed(".length(), s.length() - 1);
@@ -76,7 +77,7 @@ public class TensorConverter {
             embedderId = entry.getKey();
             embedder = entry.getValue();
         }
-        return embedder.embed(resolve(argument, embedderContext), embedderContext.copy().setEmbedderId(embedderId), type);
+        return embedder.embed(resolve(argument, properties), embedderContext.copy().setEmbedderId(embedderId), type);
     }
 
     private Embedder requireEmbedder(String embedderId) {
@@ -86,19 +87,19 @@ public class TensorConverter {
         return embedders.get(embedderId);
     }
 
-    private static String resolve(String s, Embedder.Context embedderContext) {
+    private static String resolve(String s, Properties properties) {
         if (s.startsWith("'") && s.endsWith("'"))
             return s.substring(1, s.length() - 1);
         if (s.startsWith("\"") && s.endsWith("\""))
             return s.substring(1, s.length() - 1);
         if (s.startsWith("@"))
-            return resolveReference(s, embedderContext);
+            return resolveReference(s, properties);
         return s;
     }
 
-    private static String resolveReference(String s, Embedder.Context embedderContext) {
+    private static String resolveReference(String s, Properties properties) {
         String referenceKey = s.substring(1);
-        String referencedValue = embedderContext.getContextValues().get(referenceKey);
+        String referencedValue = properties.getString(referenceKey);
         if (referencedValue == null)
             throw new IllegalArgumentException("Could not resolve query parameter reference '" + referenceKey +
                                                "' used in an embed() argument");