From bd82f97d33a5bf55408cc77d39d74f44ccbfaf47 Mon Sep 17 00:00:00 2001
From: Michael Yang <fuhai999@gmail.com>
Date: Mon, 5 Aug 2024 09:45:32 +0800
Subject: [PATCH] feat: close
 https://github.com/agents-flex/agents-flex/issues/5

---
 .../com/agentsflex/core/llm/LlmConfig.java    | 82 +++++++++++--------
 .../com/agentsflex/llm/openai/OpenAiLlm.java  |  6 ++
 2 files changed, 53 insertions(+), 35 deletions(-)

diff --git a/agents-flex-core/src/main/java/com/agentsflex/core/llm/LlmConfig.java b/agents-flex-core/src/main/java/com/agentsflex/core/llm/LlmConfig.java
index ebf41f6..afcead1 100644
--- a/agents-flex-core/src/main/java/com/agentsflex/core/llm/LlmConfig.java
+++ b/agents-flex-core/src/main/java/com/agentsflex/core/llm/LlmConfig.java
@@ -17,56 +17,68 @@
 package com.agentsflex.core.llm;
 
 import java.io.Serializable;
+import java.util.Map;
+import java.util.function.Consumer;
 
 public class LlmConfig implements Serializable {
 
-	private String model;
+    private String model;
 
-	private String endpoint;
+    private String endpoint;
 
-	private String apiKey;
+    private String apiKey;
 
-	private String apiSecret;
+    private String apiSecret;
 
-	private boolean debug;
+    private boolean debug;
 
-	public boolean isDebug() {
-		return debug;
-	}
+    private Consumer<Map<String, String>> headersConfig;
 
-	public void setDebug(boolean debug) {
-		this.debug = debug;
-	}
+    public boolean isDebug() {
+        return debug;
+    }
 
-	public String getModel() {
-		return model;
-	}
+    public void setDebug(boolean debug) {
+        this.debug = debug;
+    }
 
-	public void setModel(String model) {
-		this.model = model;
-	}
+    public String getModel() {
+        return model;
+    }
 
-	public String getEndpoint() {
-		return endpoint;
-	}
+    public void setModel(String model) {
+        this.model = model;
+    }
 
-	public void setEndpoint(String endpoint) {
-		this.endpoint = endpoint;
-	}
+    public String getEndpoint() {
+        return endpoint;
+    }
 
-	public String getApiKey() {
-		return apiKey;
-	}
+    public void setEndpoint(String endpoint) {
+        this.endpoint = endpoint;
+    }
 
-	public void setApiKey(String apiKey) {
-		this.apiKey = apiKey;
-	}
+    public String getApiKey() {
+        return apiKey;
+    }
 
-	public String getApiSecret() {
-		return apiSecret;
-	}
+    public void setApiKey(String apiKey) {
+        this.apiKey = apiKey;
+    }
 
-	public void setApiSecret(String apiSecret) {
-		this.apiSecret = apiSecret;
-	}
+    public String getApiSecret() {
+        return apiSecret;
+    }
+
+    public void setApiSecret(String apiSecret) {
+        this.apiSecret = apiSecret;
+    }
+
+    public Consumer<Map<String, String>> getHeadersConfig() {
+        return headersConfig;
+    }
+
+    public void setHeadersConfig(Consumer<Map<String, String>> headersConfig) {
+        this.headersConfig = headersConfig;
+    }
 }
diff --git a/agents-flex-llm/agents-flex-llm-openai/src/main/java/com/agentsflex/llm/openai/OpenAiLlm.java b/agents-flex-llm/agents-flex-llm-openai/src/main/java/com/agentsflex/llm/openai/OpenAiLlm.java
index 96da9dc..e3080e9 100644
--- a/agents-flex-llm/agents-flex-llm-openai/src/main/java/com/agentsflex/llm/openai/OpenAiLlm.java
+++ b/agents-flex-llm/agents-flex-llm-openai/src/main/java/com/agentsflex/llm/openai/OpenAiLlm.java
@@ -41,6 +41,7 @@ import com.alibaba.fastjson.JSONPath;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.function.Consumer;
 
 public class OpenAiLlm extends BaseLlm<OpenAiLlmConfig> {
 
@@ -72,6 +73,11 @@ public class OpenAiLlm extends BaseLlm<OpenAiLlmConfig> {
         headers.put("Content-Type", "application/json");
         headers.put("Authorization", "Bearer " + getConfig().getApiKey());
 
+        Consumer<Map<String, String>> headersConfig = config.getHeadersConfig();
+        if (headersConfig != null) {
+            headersConfig.accept(headers);
+        }
+
         String payload = OpenAiLLmUtil.promptToPayload(prompt, config, options, false);
         String endpoint = config.getEndpoint();
         String response = httpClient.post(endpoint + "/v1/chat/completions", headers, payload);