feat: store http headers through an interceptors invocation chain for…

… analytics
epam · Oct 28, 2024 · 23bf90b · 23bf90b
1 parent 3b35864
commit 23bf90b
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 2 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/data/ApiKeyData.java b/server/src/main/java/com/epam/aidial/core/server/data/ApiKeyData.java
@@ -1,6 +1,7 @@
 package com.epam.aidial.core.server.data;
 
 import com.epam.aidial.core.config.Key;
+import com.epam.aidial.core.server.Proxy;
 import com.epam.aidial.core.server.ProxyContext;
 import com.epam.aidial.core.server.security.ExtractedClaims;
 import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -10,6 +11,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /**
  * The container keeps data associated with API key.
@@ -23,6 +27,10 @@
  */
 @Data
 public class ApiKeyData {
+    private static final List<String> HTTP_HEADERS_TO_STORE = Stream.of(
+            Proxy.HEADER_JOB_TITLE,
+            Proxy.HEADER_CONVERSATION_ID
+    ).map(String::toLowerCase).toList();
     // per request key is available with during the request lifetime. It's generated in runtime
     private String perRequestKey;
     // the key of root request initiator
@@ -47,6 +55,8 @@ public class ApiKeyData {
     // deployment triggers interceptors
     private String initialDeployment;
     private String initialDeploymentApi;
+    // HTTP headers to store during an interceptor invocation chain
+    private Map<String, String> httpHeaders = new HashMap<>();
 
     public ApiKeyData() {
     }
@@ -57,6 +67,7 @@ public static void initFromContext(ApiKeyData proxyApiKeyData, ProxyContext cont
         proxyApiKeyData.setInterceptors(context.getInterceptors());
         proxyApiKeyData.setInitialDeployment(context.getInitialDeployment());
         proxyApiKeyData.setInitialDeploymentApi(context.getInitialDeploymentApi());
+        proxyApiKeyData.setHttpHeaders(collectHttpHeaders(context));
 
         if (apiKeyData.getPerRequestKey() == null) {
             proxyApiKeyData.setOriginalKey(context.getKey());
@@ -80,4 +91,13 @@ public static void initFromContext(ApiKeyData proxyApiKeyData, ProxyContext cont
     public boolean isInterceptor() {
         return perRequestKey != null && interceptors != null && interceptorIndex >= 0 && interceptorIndex < interceptors.size();
     }
+
+    private static Map<String, String> collectHttpHeaders(ProxyContext context) {
+        if (!context.getApiKeyData().getHttpHeaders().isEmpty()) {
+            return context.getApiKeyData().getHttpHeaders();
+        }
+        return context.getRequest().headers().entries().stream()
+                .filter(h -> HTTP_HEADERS_TO_STORE.contains(h.getKey().toLowerCase()))
+                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+    }
 }
diff --git a/server/src/main/java/com/epam/aidial/core/server/log/GfLogStore.java b/server/src/main/java/com/epam/aidial/core/server/log/GfLogStore.java
@@ -3,6 +3,7 @@
 import com.epam.aidial.core.config.Upstream;
 import com.epam.aidial.core.server.Proxy;
 import com.epam.aidial.core.server.ProxyContext;
+import com.epam.aidial.core.server.data.ApiKeyData;
 import com.epam.aidial.core.server.token.TokenUsage;
 import com.epam.aidial.core.server.upstream.UpstreamRoute;
 import com.epam.aidial.core.server.util.MergeChunks;
@@ -67,9 +68,13 @@ private Void doSave(ProxyContext context) {
     private void append(ProxyContext context, LogEntry entry) throws JsonProcessingException {
         HttpServerRequest request = context.getRequest();
         HttpServerResponse response = context.getResponse();
+        ApiKeyData apiKeyData = Optional.ofNullable(context.getProxyApiKeyData())
+                .orElse(context.getApiKeyData());
 
         append(entry, "{\"apiType\":\"DialOpenAI\",\"chat\":{\"id\":\"", false);
-        append(entry, request.getHeader(Proxy.HEADER_CONVERSATION_ID), true);
+        String conversationId = Optional.ofNullable(request.getHeader(Proxy.HEADER_CONVERSATION_ID))
+                .orElse(apiKeyData.getHttpHeaders().get(Proxy.HEADER_CONVERSATION_ID));
+        append(entry, conversationId, true);
 
         append(entry, "\"},\"project\":{\"id\":\"", false);
         append(entry, context.getProject(), true);
@@ -78,7 +83,9 @@ private void append(ProxyContext context, LogEntry entry) throws JsonProcessingE
         append(entry, context.getUserHash(), true);
 
         append(entry, "\",\"title\":\"", false);
-        append(entry, request.getHeader(Proxy.HEADER_JOB_TITLE), true);
+        String jobTitle = Optional.ofNullable(request.getHeader(Proxy.HEADER_JOB_TITLE))
+                .orElse(apiKeyData.getHttpHeaders().get(Proxy.HEADER_JOB_TITLE));
+        append(entry, jobTitle, true);
         append(entry, "\"}", false);
 
         TokenUsage tokenUsage = context.getTokenUsage();