From f88c451cca07a8c7b1205f2d87893dc5c23b81d3 Mon Sep 17 00:00:00 2001
From: Aleksey Vdovenko <aleksey_vdovenko@epam.com>
Date: Tue, 29 Oct 2024 11:06:31 +0100
Subject: [PATCH] feat: store http headers through an interceptors invocation
 chain (#550)

---
 .../epam/aidial/core/server/ProxyContext.java | 14 ++++++++++++-
 .../aidial/core/server/data/ApiKeyData.java   | 20 +++++++++++++++++++
 .../aidial/core/server/log/GfLogStore.java    |  4 ++--
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/server/src/main/java/com/epam/aidial/core/server/ProxyContext.java b/server/src/main/java/com/epam/aidial/core/server/ProxyContext.java
index 9ae39bdbd..4ff01f2db 100644
--- a/server/src/main/java/com/epam/aidial/core/server/ProxyContext.java
+++ b/server/src/main/java/com/epam/aidial/core/server/ProxyContext.java
@@ -26,6 +26,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -217,4 +218,15 @@ public ProxyContext exposeHeaders() {
 
         return this;
     }
-}
\ No newline at end of file
+
+    public String getRequestHeader(String name) {
+        String value = request.getHeader(name);
+        if (value != null) {
+            return value;
+        }
+        return Optional.ofNullable(apiKeyData)
+                .map(ApiKeyData::getHttpHeaders)
+                .map(h -> h.get(name))
+                .orElse(null);
+    }
+}
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
index 427c33f5c..c764f4c7d 100644
--- 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 = List.of(
+            Proxy.HEADER_JOB_TITLE,
+            Proxy.HEADER_CONVERSATION_ID
+    );
     // 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;
+    // Original HTTP headers to be stored 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()))
+                .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
index f3db0ea83..848a1a806 100644
--- 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
@@ -69,7 +69,7 @@ private void append(ProxyContext context, LogEntry entry) throws JsonProcessingE
         HttpServerResponse response = context.getResponse();
 
         append(entry, "{\"apiType\":\"DialOpenAI\",\"chat\":{\"id\":\"", false);
-        append(entry, request.getHeader(Proxy.HEADER_CONVERSATION_ID), true);
+        append(entry, context.getRequestHeader(Proxy.HEADER_CONVERSATION_ID), true);
 
         append(entry, "\"},\"project\":{\"id\":\"", false);
         append(entry, context.getProject(), true);
@@ -78,7 +78,7 @@ 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);
+        append(entry, context.getRequestHeader(Proxy.HEADER_JOB_TITLE), true);
         append(entry, "\"}", false);
 
         TokenUsage tokenUsage = context.getTokenUsage();