diff --git a/app/config/config.py b/app/config/config.py index 1c7dd97..5bc79d4 100644 --- a/app/config/config.py +++ b/app/config/config.py @@ -67,6 +67,9 @@ class Settings(BaseSettings): # 智能路由配置 URL_NORMALIZATION_ENABLED: bool = False # 是否启用智能路由映射功能 + # 自定义 Headers + CUSTOM_HEADERS: Dict[str, str] = {} + # 模型相关配置 SEARCH_MODELS: List[str] = ["gemini-2.0-flash-exp"] IMAGE_MODELS: List[str] = ["gemini-2.0-flash-exp"] diff --git a/app/service/client/api_client.py b/app/service/client/api_client.py index 641b8d6..2b672f9 100644 --- a/app/service/client/api_client.py +++ b/app/service/client/api_client.py @@ -44,6 +44,13 @@ class GeminiApiClient(ApiClient): model = model[:-20] return model + def _prepare_headers(self) -> Dict[str, str]: + headers = {} + if settings.CUSTOM_HEADERS: + headers.update(settings.CUSTOM_HEADERS) + logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}") + return headers + async def get_models(self, api_key: str) -> Optional[Dict[str, Any]]: """获取可用的 Gemini 模型列表""" timeout = httpx.Timeout(timeout=5) @@ -56,10 +63,11 @@ class GeminiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers() async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/models?key={api_key}&pageSize=1000" try: - response = await client.get(url) + response = await client.get(url, headers=headers) response.raise_for_status() return response.json() except httpx.HTTPStatusError as e: @@ -82,9 +90,10 @@ class GeminiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers() async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/models/{model}:generateContent?key={api_key}" - response = await client.post(url, json=payload) + response = await client.post(url, json=payload, headers=headers) if response.status_code != 200: error_content = response.text raise Exception(f"API call failed with status code {response.status_code}, {error_content}") @@ -102,9 +111,10 @@ class GeminiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers() async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse&key={api_key}" - async with client.stream(method="POST", url=url, json=payload) as response: + async with client.stream(method="POST", url=url, json=payload, headers=headers) as response: if response.status_code != 200: error_content = await response.aread() error_msg = error_content.decode("utf-8") @@ -124,9 +134,10 @@ class GeminiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for counting tokens: {proxy_to_use}") + headers = self._prepare_headers() async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/models/{model}:countTokens?key={api_key}" - response = await client.post(url, json=payload) + response = await client.post(url, json=payload, headers=headers) if response.status_code != 200: error_content = response.text raise Exception(f"API call failed with status code {response.status_code}, {error_content}") @@ -140,6 +151,13 @@ class OpenaiApiClient(ApiClient): self.base_url = base_url self.timeout = timeout + def _prepare_headers(self, api_key: str) -> Dict[str, str]: + headers = {"Authorization": f"Bearer {api_key}"} + if settings.CUSTOM_HEADERS: + headers.update(settings.CUSTOM_HEADERS) + logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}") + return headers + async def get_models(self, api_key: str) -> Dict[str, Any]: timeout = httpx.Timeout(self.timeout, read=self.timeout) @@ -151,9 +169,9 @@ class OpenaiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers(api_key) async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/openai/models" - headers = {"Authorization": f"Bearer {api_key}"} response = await client.get(url, headers=headers) if response.status_code != 200: error_content = response.text @@ -171,9 +189,9 @@ class OpenaiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers(api_key) async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/openai/chat/completions" - headers = {"Authorization": f"Bearer {api_key}"} response = await client.post(url, json=payload, headers=headers) if response.status_code != 200: error_content = response.text @@ -190,9 +208,9 @@ class OpenaiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers(api_key) async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/openai/chat/completions" - headers = {"Authorization": f"Bearer {api_key}"} async with client.stream(method="POST", url=url, json=payload, headers=headers) as response: if response.status_code != 200: error_content = await response.aread() @@ -212,9 +230,9 @@ class OpenaiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers(api_key) async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/openai/embeddings" - headers = {"Authorization": f"Bearer {api_key}"} payload = { "input": input, "model": model, @@ -236,9 +254,9 @@ class OpenaiApiClient(ApiClient): proxy_to_use = random.choice(settings.PROXIES) logger.info(f"Using proxy for getting models: {proxy_to_use}") + headers = self._prepare_headers(api_key) async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client: url = f"{self.base_url}/openai/images/generations" - headers = {"Authorization": f"Bearer {api_key}"} response = await client.post(url, json=payload, headers=headers) if response.status_code != 200: error_content = response.text diff --git a/app/static/js/config_editor.js b/app/static/js/config_editor.js index 9c8de79..874cf1b 100644 --- a/app/static/js/config_editor.js +++ b/app/static/js/config_editor.js @@ -5,6 +5,9 @@ const ARRAY_INPUT_CLASS = "array-input"; const MAP_ITEM_CLASS = "map-item"; const MAP_KEY_INPUT_CLASS = "map-key-input"; const MAP_VALUE_INPUT_CLASS = "map-value-input"; +const CUSTOM_HEADER_ITEM_CLASS = "custom-header-item"; +const CUSTOM_HEADER_KEY_INPUT_CLASS = "custom-header-key-input"; +const CUSTOM_HEADER_VALUE_INPUT_CLASS = "custom-header-value-input"; const SAFETY_SETTING_ITEM_CLASS = "safety-setting-item"; const SHOW_CLASS = "show"; // For modals const API_KEY_REGEX = /AIzaSy\S{33}/g; @@ -383,6 +386,12 @@ document.addEventListener("DOMContentLoaded", function () { addSafetySettingBtn.addEventListener("click", () => addSafetySettingItem()); } + // Add Custom Header button + const addCustomHeaderBtn = document.getElementById("addCustomHeaderBtn"); + if (addCustomHeaderBtn) { + addCustomHeaderBtn.addEventListener("click", () => addCustomHeaderItem()); + } + initializeSensitiveFields(); // Initialize sensitive field handling // Vertex API Key Modal Elements and Events @@ -691,6 +700,14 @@ async function initConfig() { ) { config.THINKING_BUDGET_MAP = {}; // 默认为空对象 } + // --- 新增:处理 CUSTOM_HEADERS 默认值 --- + if ( + !config.CUSTOM_HEADERS || + typeof config.CUSTOM_HEADERS !== "object" || + config.CUSTOM_HEADERS === null + ) { + config.CUSTOM_HEADERS = {}; // 默认为空对象 + } // --- 新增:处理 SAFETY_SETTINGS 默认值 --- if (!config.SAFETY_SETTINGS || !Array.isArray(config.SAFETY_SETTINGS)) { config.SAFETY_SETTINGS = []; // 默认为空数组 @@ -756,6 +773,7 @@ async function initConfig() { VERTEX_EXPRESS_BASE_URL: "", // 确保默认值存在 THINKING_MODELS: [], THINKING_BUDGET_MAP: {}, + CUSTOM_HEADERS: {}, AUTO_DELETE_ERROR_LOGS_ENABLED: false, AUTO_DELETE_ERROR_LOGS_DAYS: 7, // 新增默认值 AUTO_DELETE_REQUEST_LOGS_ENABLED: false, // 新增默认值 @@ -854,6 +872,20 @@ function populateForm(config) { '
请在上方添加思考模型,预算将自动关联。
'; } + // Populate CUSTOM_HEADERS + const customHeadersContainer = document.getElementById("CUSTOM_HEADERS_container"); + let customHeadersAdded = false; + if (customHeadersContainer && config.CUSTOM_HEADERS && typeof config.CUSTOM_HEADERS === "object") { + for (const [key, value] of Object.entries(config.CUSTOM_HEADERS)) { + createAndAppendCustomHeaderItem(key, value); + customHeadersAdded = true; + } + } + if (!customHeadersAdded && customHeadersContainer) { + customHeadersContainer.innerHTML = + '
添加自定义请求头,例如 X-Api-Key: your-key
'; + } + // 4. Populate other array fields (excluding THINKING_MODELS) for (const [key, value] of Object.entries(config)) { if (Array.isArray(value) && key !== "THINKING_MODELS") { @@ -1562,6 +1594,60 @@ function createAndAppendBudgetMapItem(mapKey, mapValue, modelId) { container.appendChild(mapItem); } +/** + * Adds a new custom header item to the DOM. + */ +function addCustomHeaderItem() { + createAndAppendCustomHeaderItem("", ""); +} + +/** + * Creates and appends a DOM element for a custom header. + * @param {string} key - The header key. + * @param {string} value - The header value. + */ +function createAndAppendCustomHeaderItem(key, value) { + const container = document.getElementById("CUSTOM_HEADERS_container"); + if (!container) { + console.error("Cannot add custom header: CUSTOM_HEADERS_container not found!"); + return; + } + + const placeholder = container.querySelector(".text-gray-500.italic"); + if (placeholder && container.children.length === 1 && container.firstChild === placeholder) { + container.innerHTML = ""; + } + + const headerItem = document.createElement("div"); + headerItem.className = `${CUSTOM_HEADER_ITEM_CLASS} flex items-center mb-2 gap-2`; + + const keyInput = document.createElement("input"); + keyInput.type = "text"; + keyInput.value = key; + keyInput.placeholder = "Header Name"; + keyInput.className = `${CUSTOM_HEADER_KEY_INPUT_CLASS} flex-grow px-3 py-2 border border-gray-300 rounded-md focus:outline-none bg-gray-100 text-gray-500`; + + const valueInput = document.createElement("input"); + valueInput.type = "text"; + valueInput.value = value; + valueInput.placeholder = "Header Value"; + valueInput.className = `${CUSTOM_HEADER_VALUE_INPUT_CLASS} flex-grow px-3 py-2 border border-gray-300 rounded-md focus:outline-none focus:border-primary-500 focus:ring focus:ring-primary-200 focus:ring-opacity-50`; + + const removeBtn = createRemoveButton(); + removeBtn.addEventListener("click", () => { + headerItem.remove(); + if (container.children.length === 0) { + container.innerHTML = '
添加自定义请求头,例如 X-Api-Key: your-key
'; + } + }); + + headerItem.appendChild(keyInput); + headerItem.appendChild(valueInput); + headerItem.appendChild(removeBtn); + + container.appendChild(headerItem); +} + /** * Collects all data from the configuration form. * @returns {object} An object containing all configuration data. @@ -1638,6 +1724,19 @@ function collectFormData() { }); } + const customHeadersContainer = document.getElementById("CUSTOM_HEADERS_container"); + if (customHeadersContainer) { + formData["CUSTOM_HEADERS"] = {}; + const customHeaderItems = customHeadersContainer.querySelectorAll(`.${CUSTOM_HEADER_ITEM_CLASS}`); + customHeaderItems.forEach((item) => { + const keyInput = item.querySelector(`.${CUSTOM_HEADER_KEY_INPUT_CLASS}`); + const valueInput = item.querySelector(`.${CUSTOM_HEADER_VALUE_INPUT_CLASS}`); + if (keyInput && valueInput && keyInput.value.trim() !== "") { + formData["CUSTOM_HEADERS"][keyInput.value.trim()] = valueInput.value.trim(); + } + }); + } + if (safetySettingsContainer) { formData["SAFETY_SETTINGS"] = []; const settingItems = safetySettingsContainer.querySelectorAll( diff --git a/app/templates/config_editor.html b/app/templates/config_editor.html index 98aae05..5e83547 100644 --- a/app/templates/config_editor.html +++ b/app/templates/config_editor.html @@ -898,7 +898,24 @@ endblock %} {% block head_extra_styles %} /> Gemini API的基础URL - + + +
+ +
+ +
+ 添加自定义请求头,例如 X-Api-Key: your-key +
+
+
+ +
+ 在这里添加的键值对将被添加到所有出站API请求的Header中。 +
+