🔑 Key Rotator (#633)

Author: awtkns

platform/reworkd_platform/tests/test_api_utils.py

@@ -0,0 +1,48 @@
+from collections import Counter
+
+import pytest
+
+from reworkd_platform.web.api.agent.api_utils import rotate_keys
+
+ITERATIONS = 10000
+
+
+def test_rotate_keys():
+    pk = "primary_key"
+    sk = "secondary_key"
+
+    results = []
+    for _ in range(ITERATIONS):
+        key = rotate_keys(pk, sk)
+        assert key in [pk, sk]
+        results.append(key)
+
+    counter = Counter(results)
+    assert 0.65 < counter[pk] / ITERATIONS < 0.75
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "gpt-4",
+        "gpt-4-0314",
+        "gpt-4-turbo",
+    ],
+)
+def test_rotate_keys_gpt_4(model):
+    pk = "gpt-3-primary_key"
+    sk = "gpt-4-primary_key"
+
+    results = []
+    for _ in range(ITERATIONS):
+        key = rotate_keys(pk, sk, model=model)
+        assert key in [pk, sk]
+        results.append(key)
+
+    counter = Counter(results)
+    assert 0.65 < counter[sk] / ITERATIONS < 0.75
+
+
+def test_rotate_keys_no_secondary():
+    pk = "primary_key"
+    assert rotate_keys(pk, None) == pk

platform/reworkd_platform/web/api/agent/api_utils.py

@@ -0,0 +1,19 @@
+import random
+from typing import Optional
+
+PRIMARY_KEY_RATE = 0.7
+SECONDARY_KEY_RATE = 1 - PRIMARY_KEY_RATE
+WEIGHTS = [PRIMARY_KEY_RATE, SECONDARY_KEY_RATE]
+
+
+def rotate_keys(
+    primary_key: str, secondary_key: Optional[str], model: str = "gpt-3.5-turbo"
+) -> str:
+    if not secondary_key:
+        return primary_key
+
+    keys = [primary_key, secondary_key]
+    if "gpt-4" in model:
+        keys.reverse()
+
+    return random.choices(keys, WEIGHTS)[0]

platform/reworkd_platform/web/api/agent/model_settings.py

@@ -1,11 +1,11 @@
-from random import randint
 from typing import Optional
 
 import openai
 from langchain.chat_models import ChatOpenAI
 from pydantic import BaseModel
 
 from reworkd_platform.settings import settings
+from reworkd_platform.web.api.agent.api_utils import rotate_keys
 
 
 class ModelSettings(BaseModel):
@@ -16,13 +16,6 @@ class ModelSettings(BaseModel):
     language: Optional[str] = "English"
 
 
-def get_server_side_key() -> str:
-    keys = [
-        key.strip() for key in (settings.openai_api_key or "").split(",") if key.strip()
-    ]
-    return keys[randint(0, len(keys) - 1)] if keys else ""
-
-
 GPT_35_TURBO = "gpt-3.5-turbo"
 
 openai.api_base = settings.openai_api_base
@@ -33,7 +26,10 @@ def create_model(
 ) -> ChatOpenAI:
     return ChatOpenAI(
         client=None,  # Meta private value but mypy will complain its missing
-        openai_api_key=get_server_side_key(),
+        openai_api_key=rotate_keys(
+            primary_key=settings.openai_api_key,
+            secondary_key=settings.secondary_openai_api_key,
+        ),
         temperature=model_settings.customTemperature
         if model_settings and model_settings.customTemperature is not None
         else 0.9,

platform/reworkd_platform/web/api/agent/tools/image.py

@@ -3,9 +3,9 @@
 from fastapi.responses import StreamingResponse as FastAPIStreamingResponse
 
 from reworkd_platform.settings import settings
+from reworkd_platform.web.api.agent.api_utils import rotate_keys
 from reworkd_platform.web.api.agent.model_settings import (
     ModelSettings,
-    get_server_side_key,
 )
 from reworkd_platform.web.api.agent.tools.stream_mock import stream_string
 from reworkd_platform.web.api.agent.tools.tool import Tool
@@ -27,7 +27,10 @@ async def get_replicate_image(input_str: str) -> str:
 
 # Use AI to generate an Image based on a prompt
 async def get_open_ai_image(input_str: str) -> str:
-    api_key = get_server_side_key()
+    api_key = rotate_keys(
+        primary_key=settings.openai_api_key,
+        secondary_key=settings.secondary_openai_api_key,
+    )
 
     response = openai.Image.create(
         api_key=api_key,