Feat: Multiple OpenAI Keys Support (#12)

* Update README

* Multiple OpenAI Keys Support

* update pytest

README.md

@@ -47,38 +47,40 @@ Test access: https://caloi.top/v1/chat/completions is equivalent to https://api.
 # Features
 
 - [x] Supports forwarding of all OpenAI interfaces
-- [x] Supports request IP verification
-- [x] Supports streaming forwarding
-- [x] Supports default API key
+- [x] Request IP verification
+- [x] Streaming Response
+- [x] Supports default API key (cyclic call with multiple API keys)
 - [x] pip installation and deployment
 - [x] Docker deployment
 - [x] Support for multiple worker processes
-- [x] Support for specifying the forwarding routing prefix.
+- [x] Support for specifying the forwarding routing prefix
 
 # Usage
 
 > Here, the proxy address set up by the individual, https://caloi.top, is used as an example
 
 ### Using in a module
 
+
+**Python**
+
+```diff
+  import openai
++ openai.api_base = "https://caloi.top/v1"
+  openai.api_key = "sk-******"
+```
+
 **JS/TS**
 
 ```diff
   import { Configuration } from "openai";
 
   const configuration = new Configuration({
-+ basePath: "https://caloi.top",
++ basePath: "https://caloi.top/v1",
   apiKey: "sk-******",
   });
 ```
 
-**Python**
-
-```diff
-  import openai
-+ openai.api_base = "https://caloi.top/v1"
-  openai.api_key = "sk-******"
-```
 
 ### Image Generation (DALL-E):
 
@@ -155,8 +157,7 @@ Note: You can also pass in the environment variable OPENAI_API_KEY=sk-xxx as the
 
 # Service Usage
 
-Simply replace the OpenAI API address with the address of the service we set up, such as:
-
+Simply replace the OpenAI API address with the address of the service we set up, such as `Chat Completion`
 ```bash 
 https://api.openai.com/v1/chat/completions 
 ``` 
@@ -181,9 +182,9 @@ refer to the `.env` file in the project root directory
 
 | Environment Variable  | Description | Default Value |
 |-----------------|------------|:------------------------:|
-| OPENAI_API_KEY  | Default API key | None |
+| OPENAI_API_KEY  | Default API key, supports multiple default API keys separated by space. | None |
 | OPENAI_BASE_URL | Forwarding base URL | `https://api.openai.com` |
 |LOG_CHAT| Whether to log chat content | `true` |
 |ROUTE_PREFIX| Route prefix | None |
-| IP_WHITELIST    | IP whitelist | None |
-| IP_BLACKLIST    | IP blacklist | None |
+| IP_WHITELIST    | IP whitelist, separated by space. | None |
+| IP_BLACKLIST    | IP blacklist, separated by space. | None |

README_ZH.md

@@ -70,13 +70,13 @@ api的服务器上，通过该服务转发OpenAI的请求。即搭建反向代
 # Features
 
 - [x] 支持转发OpenAI所有接口
-- [x] 支持请求IP验证
-- [x] 支持流式转发
-- [x] 支持默认api key
+- [x] 支持流式响应
+- [x] 支持默认api key(多api key 循环调用)
 - [x] pip安装部署
 - [x] docker部署
 - [x] 支持多进程转发
 - [x] 支持指定转发路由前缀
+- [x] 支持请求IP验证
 
 # Usage
 
@@ -90,7 +90,7 @@ api的服务器上，通过该服务转发OpenAI的请求。即搭建反向代
   import { Configuration } from "openai";
 
   const configuration = new Configuration({
-+ basePath: "https://caloi.top",
++ basePath: "https://caloi.top/v1",
   apiKey: "sk-******",
   });
 ```
@@ -198,12 +198,12 @@ http://{ip}:{port}/v1/chat/completions
 **环境变量配置项**  
 参考项目根目录下`.env`文件
 
-| 环境变量      | 说明         |           默认值            |
-|-----------------|------------|:------------------------:|
-| OPENAI_API_KEY  | 默认api key  |            无             |
-| OPENAI_BASE_URL | 转发base url | `https://api.openai.com` |
-|LOG_CHAT| 是否记录聊天内容   |          `true`          |
-|ROUTE_PREFIX| 路由前缀     |            无             |
-| IP_WHITELIST    | ip白名单      |           无            |
-| IP_BLACKLIST    | ip黑名单      |           无            | 
+| 环境变量      | 说明                             |           默认值            |
+|-----------------|--------------------------------|:------------------------:|
+| OPENAI_API_KEY  | 默认api key，支持多个默认api key, 以空格分割 |            无             |
+| OPENAI_BASE_URL | 转发base url                     | `https://api.openai.com` |
+|LOG_CHAT| 是否记录聊天内容                       |          `true`          |
+|ROUTE_PREFIX| 路由前缀                           |            无             |
+| IP_WHITELIST    | ip白名单, 空格分开                    |           无            |
+| IP_BLACKLIST    | ip黑名单, 空格分开                    |           无            | 
 

openai_forward/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.3"
+__version__ = "0.1.4"
 
 from dotenv import load_dotenv
 

openai_forward/_base.py

@@ -4,32 +4,26 @@
 import httpx
 from starlette.background import BackgroundTask
 import os
+from itertools import cycle
 from .content.chat import parse_chat_completions, ChatSaver
+from .config import env2list
 
 
 class OpenaiBase:
-    _default_api_key = os.environ.get("OPENAI_API_KEY", "").strip()
     _base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com").strip()
     _LOG_CHAT = os.environ.get("LOG_CHAT", "False").lower() == "true"
     _ROUTE_PREFIX = os.environ.get("ROUTE_PREFIX", "").strip()
-    IP_WHITELIST = os.environ.get("IP_WHITELIST", "").strip()
-    IP_BLACKLIST = os.environ.get("IP_BLACKLIST", "").strip()
-    if IP_BLACKLIST:
-        IP_BLACKLIST = [i.strip() for i in IP_BLACKLIST.split(' ')]
-    else:
-        IP_BLACKLIST = []
-    if IP_WHITELIST:
-        IP_WHITELIST = [i.strip() for i in IP_WHITELIST.split(' ')]
-    else:
-        IP_WHITELIST = []
+    _default_api_key_list = env2list("OPENAI_API_KEY", sep=" ")
+    _cycle_api_key = cycle(_default_api_key_list)
+    IP_WHITELIST = env2list("IP_WHITELIST", sep=" ")
+    IP_BLACKLIST = env2list("IP_BLACKLIST", sep=" ")
+
     if _ROUTE_PREFIX:
         if _ROUTE_PREFIX.endswith('/'):
             _ROUTE_PREFIX = _ROUTE_PREFIX[:-1]
         if not _ROUTE_PREFIX.startswith('/'):
             _ROUTE_PREFIX = '/' + _ROUTE_PREFIX
-    stream_timeout = 20
     timeout = 30
-    non_stream_timeout = 30
     chatsaver = ChatSaver(save_interval=10)
 
     def validate_request_host(self, ip):
@@ -66,8 +60,8 @@ async def _reverse_proxy(cls, request: Request):
         auth = headers.pop("authorization", None)
         if auth and str(auth).startswith("Bearer sk-"):
             tmp_headers = {'Authorization': auth}
-        elif cls._default_api_key:
-            auth = "Bearer " + cls._default_api_key
+        elif cls._default_api_key_list:
+            auth = "Bearer " + next(cls._cycle_api_key)
             tmp_headers = {'Authorization': auth}
         else:
             tmp_headers = {}

openai_forward/config.py

@@ -4,6 +4,7 @@
 from typing import Union, List, Dict
 import sys
 import logging
+import os
 
 
 class InterceptHandler(logging.Handler):
@@ -75,3 +76,13 @@ def json_dump(data: Union[List, Dict], filepath: str, rel=False, indent_2=False,
     abs_path = relp(filepath, parents=1) if rel else filepath
     with open(abs_path, mode=mode) as f:
         f.write(orjson.dumps(data, option=orjson_option))
+
+
+def str2list(s: str, sep=' '):
+    if s:
+        return [i.strip() for i in s.split(sep) if i.strip()]
+    else:
+        return []
+
+def env2list(env_name: str, sep=" "):
+    return str2list(os.environ.get(env_name, "").strip(), sep=sep)

tests/test_api.py

@@ -1,6 +1,7 @@
 from openai_forward.openai import OpenaiBase
 from fastapi import HTTPException
 import pytest
+from itertools import cycle
 
 
 @pytest.fixture(scope="module")
@@ -9,9 +10,25 @@ def openai() -> OpenaiBase:
 
 
 class TestOpenai:
+
+    def teardown_method(self):
+        OpenaiBase.IP_BLACKLIST = []
+        OpenaiBase.IP_WHITELIST = []
+        OpenaiBase._default_api_key_list = []
+
     def test_env(self, openai: OpenaiBase):
         assert openai._base_url == "https://api.openai.com"
 
+    def test_api_keys(self, openai: OpenaiBase):
+        assert openai._default_api_key_list == []
+        openai._default_api_key_list = ["a", "b"]
+        openai._cycle_api_key = cycle(openai._default_api_key_list)
+        assert next(openai._cycle_api_key) == "a"
+        assert next(openai._cycle_api_key) == "b"
+        assert next(openai._cycle_api_key) == "a"
+        assert next(openai._cycle_api_key) == "b"
+        assert next(openai._cycle_api_key) == "a"
+
     def test_validate_ip(self, openai: OpenaiBase):
         ip1 = "1.1.1.1"
         ip2 = "2.2.2.2"

tests/test_chat_save.py

@@ -1,13 +1,20 @@
 from openai_forward.content.chat import ChatSaver
 import pytest
 import os
+from utils import rm
 
 @pytest.fixture(scope="module")
 def saver() -> ChatSaver:
     os.system("rm -rf chat_*.txt")
     return ChatSaver(save_interval=1, max_chat_size=2)
 
 class TestChatSaver:
+
+    @classmethod
+    def teardown_class(cls):
+        rm("Log/*.log")
+        rm("chat*.txt")
+
     def test_init(self, saver: ChatSaver):
         assert saver.chat_file == "chat_0.txt"
 

tests/utils.py

@@ -0,0 +1,20 @@
+from sparrow import ls
+import os
+import shutil
+
+def rm(*file_pattern: str, rel=False):
+    """Remove files or directories.
+    Example:
+    --------
+        >>> rm("*.jpg", "*.png")
+        >>> rm("*.jpg", "*.png", rel=True)
+    """
+    path_list = ls(".", *file_pattern, relp=rel, concat="extend")
+    for file in path_list:
+        if os.path.isfile(file):
+            print("remove ", file)
+            os.remove(file)
+            # os.system("rm -f " + file)
+        elif os.path.isdir(file):
+            shutil.rmtree(file, ignore_errors=True)
+            print("rm tree ", file)