Refactor user API key matching

JuhaKalevi · Mar 15, 2024 · 7f2f6cd · 7f2f6cd
1 parent a905579
commit 7f2f6cd
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/src/openai_models.py b/src/openai_models.py
@@ -8,14 +8,12 @@
 
 async def chat_completion(kwargs):
   api_key = environ.get('OPENAI_API_KEY')
-  for message in reversed(kwargs['messages']):
-    if message['role'] == 'user':
-      user_api_key_path = f"../api_keys/{message.get('name')}"
-      if path.isfile(user_api_key_path):
-        with open(user_api_key_path, 'r', encoding='ascii') as api_key_file:
-          api_key = api_key_file.read().strip()
-        print(f"Using {message['name']}'s API key.")
-        break
+  triggering_message = [message for message in kwargs['messages'] if message['role'] == 'user'][-1]
+  user_api_key_path = f"../api_keys/{triggering_message.get('name')}"
+  if path.isfile(user_api_key_path):
+    with open(user_api_key_path, 'r', encoding='ascii') as api_key_file:
+      api_key = api_key_file.read().strip()
+    print(f"Using {triggering_message['name']}'s API key.")
   async for part in await AsyncOpenAI(api_key=api_key, base_url=environ.get('OPENAI_API_URL_OVERRIDE')).chat.completions.create(**kwargs, stream=True):
     yield part.choices[0].delta