fix

LangStream · Sep 13, 2023 · e9dbb98 · e9dbb98
1 parent 6f67bc8
commit e9dbb98
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 28 deletions.
diff --git a/...am-api-gateway/src/main/java/ai/langstream/apigateway/websocket/handlers/ChatHandler.java b/...am-api-gateway/src/main/java/ai/langstream/apigateway/websocket/handlers/ChatHandler.java
@@ -80,7 +80,7 @@ protected List<String> getAllRequiredParameters(Gateway gateway) {
         }
         if (gateway.getChatOptions() != null && gateway.getChatOptions().getHeaders() != null) {
             for (Gateway.KeyValueComparison header : gateway.getChatOptions().getHeaders()) {
-                parameters.add(header.getKeyWithDefaultValue());
+                parameters.add(header.key());
             }
         }
         return parameters;
@@ -104,13 +104,7 @@ private void setupProducer(AuthenticatedGatewayRequestContext context) {
         final List<Gateway.KeyValueComparison> gwHeaders = chatOptions.getHeaders();
         if (gwHeaders != null) {
             for (Gateway.KeyValueComparison gwHeader : gwHeaders) {
-                final String key = gwHeader.getKeyWithDefaultValue();
-                headerConfig.add(
-                        new Gateway.KeyValueComparison(
-                                key,
-                                gwHeader.value(),
-                                gwHeader.valueFromParameters(),
-                                gwHeader.valueFromAuthentication()));
+                headerConfig.add(gwHeader);
             }
         }
         final List<Header> commonHeaders =
@@ -133,13 +127,7 @@ private void setupReader(AuthenticatedGatewayRequestContext context) throws Exce
         final List<Gateway.KeyValueComparison> gwHeaders = chatOptions.getHeaders();
         if (gwHeaders != null) {
             for (Gateway.KeyValueComparison gwHeader : gwHeaders) {
-                final String key = gwHeader.getKeyWithDefaultValue();
-                headerFilters.add(
-                        new Gateway.KeyValueComparison(
-                                key,
-                                gwHeader.value(),
-                                gwHeader.valueFromParameters(),
-                                gwHeader.valueFromAuthentication()));
+                headerFilters.add(gwHeader);
             }
         }
         final List<Function<Record, Boolean>> messageFilters =

diff --git a/langstream-api/src/main/java/ai/langstream/api/model/Gateway.java b/langstream-api/src/main/java/ai/langstream/api/model/Gateway.java
@@ -90,6 +90,19 @@ public record KeyValueComparison(
                 throw new IllegalArgumentException(
                         "Cannot specify both value-from-parameters and value-from-authentication");
             }
+            if (key == null) {
+                if (value != null) {
+                    key = value;
+                } else if (valueFromParameters != null) {
+                    key = valueFromParameters;
+                } else {
+                    key = valueFromAuthentication;
+                }
+                if (key == null) {
+                    throw new IllegalArgumentException("Not able to compute key: " + this);
+                }
+            }
+
         }
 
         public static KeyValueComparison value(String key, String value) {
@@ -106,20 +119,8 @@ public static KeyValueComparison valueFromAuthentication(
             return new KeyValueComparison(key, null, null, valueFromAuthentication);
         }
 
-        public String getKeyWithDefaultValue() {
-            if (key != null) {
-                return key;
-            }
+        private String getKeyWithDefaultValue() {
 
-            if (value != null) {
-                return value;
-            }
-            if (valueFromParameters != null) {
-                return valueFromParameters;
-            }
-            if (valueFromAuthentication != null) {
-                return valueFromAuthentication;
-            }
             throw new IllegalStateException();
         }
     }