diff --git a/src/main/java/com/unfbx/chatgpt/function/KeyRandomStrategy.java b/src/main/java/com/unfbx/chatgpt/function/KeyRandomStrategy.java index ede96d8..88f61ba 100644 --- a/src/main/java/com/unfbx/chatgpt/function/KeyRandomStrategy.java +++ b/src/main/java/com/unfbx/chatgpt/function/KeyRandomStrategy.java @@ -1,6 +1,7 @@ package com.unfbx.chatgpt.function; import cn.hutool.core.util.RandomUtil; +import okhttp3.Request; import java.util.List; @@ -13,7 +14,7 @@ public class KeyRandomStrategy implements KeyStrategyFunction, String> { @Override - public String apply(List apiKeys) { + public String apply(List apiKeys, Request request) { return RandomUtil.randomEle(apiKeys); } } diff --git a/src/main/java/com/unfbx/chatgpt/function/KeyStrategyFunction.java b/src/main/java/com/unfbx/chatgpt/function/KeyStrategyFunction.java index ae166fa..01ba4dc 100644 --- a/src/main/java/com/unfbx/chatgpt/function/KeyStrategyFunction.java +++ b/src/main/java/com/unfbx/chatgpt/function/KeyStrategyFunction.java @@ -1,5 +1,7 @@ package com.unfbx.chatgpt.function; +import okhttp3.Request; + import java.util.function.Function; /** @@ -19,6 +21,6 @@ public interface KeyStrategyFunction { * @param t the function argument * @return the function result */ - R apply(T t); + R apply(T t, Request request); } diff --git a/src/main/java/com/unfbx/chatgpt/interceptor/DefaultOpenAiAuthInterceptor.java b/src/main/java/com/unfbx/chatgpt/interceptor/DefaultOpenAiAuthInterceptor.java index b6ac618..5d8a539 100644 --- a/src/main/java/com/unfbx/chatgpt/interceptor/DefaultOpenAiAuthInterceptor.java +++ b/src/main/java/com/unfbx/chatgpt/interceptor/DefaultOpenAiAuthInterceptor.java @@ -42,7 +42,7 @@ public DefaultOpenAiAuthInterceptor(Map warringConfig) { @Override public Response intercept(Chain chain) throws IOException { Request original = chain.request(); - return chain.proceed(auth(super.getKey(), original)); + return chain.proceed(auth(super.getKey(original), original)); } /** diff --git a/src/main/java/com/unfbx/chatgpt/interceptor/DynamicKeyOpenAiAuthInterceptor.java b/src/main/java/com/unfbx/chatgpt/interceptor/DynamicKeyOpenAiAuthInterceptor.java index 0efa32d..b097b23 100644 --- a/src/main/java/com/unfbx/chatgpt/interceptor/DynamicKeyOpenAiAuthInterceptor.java +++ b/src/main/java/com/unfbx/chatgpt/interceptor/DynamicKeyOpenAiAuthInterceptor.java @@ -52,8 +52,8 @@ public DynamicKeyOpenAiAuthInterceptor(Map warringConfig) { @Override public Response intercept(Chain chain) throws IOException { - String key = getKey(); Request original = chain.request(); + String key = getKey(original); Request request = this.auth(key, original); Response response = chain.proceed(request); if (!response.isSuccessful() && response.body() != null) { diff --git a/src/main/java/com/unfbx/chatgpt/interceptor/OpenAiAuthInterceptor.java b/src/main/java/com/unfbx/chatgpt/interceptor/OpenAiAuthInterceptor.java index 3cf29ce..dd76080 100644 --- a/src/main/java/com/unfbx/chatgpt/interceptor/OpenAiAuthInterceptor.java +++ b/src/main/java/com/unfbx/chatgpt/interceptor/OpenAiAuthInterceptor.java @@ -58,12 +58,12 @@ public abstract class OpenAiAuthInterceptor implements Interceptor { * * @return key */ - public final String getKey() { + public final String getKey(Request request) { if (CollectionUtil.isEmpty(apiKey)) { this.noHaveActiveKeyWarring(); throw new BaseException(CommonError.NO_ACTIVE_API_KEYS); } - return keyStrategy.apply(apiKey); + return keyStrategy.apply(apiKey, request); } /** diff --git a/src/test/java/com/unfbx/chatgpt/FirstKeyStrategy.java b/src/test/java/com/unfbx/chatgpt/FirstKeyStrategy.java index 5bdf72f..67d50b1 100644 --- a/src/test/java/com/unfbx/chatgpt/FirstKeyStrategy.java +++ b/src/test/java/com/unfbx/chatgpt/FirstKeyStrategy.java @@ -1,6 +1,7 @@ package com.unfbx.chatgpt; import com.unfbx.chatgpt.function.KeyStrategyFunction; +import okhttp3.Request; import java.util.List; @@ -19,7 +20,7 @@ public class FirstKeyStrategy implements KeyStrategyFunction, Strin * @return */ @Override - public String apply(List keys) { + public String apply(List keys, Request request) { return keys.get(0); } }