From eb395b3959fac45cf05baa31998811e791fd79c1 Mon Sep 17 00:00:00 2001
From: "G. P." <georgepolimenakos@outlook.com>
Date: Fri, 20 Dec 2024 10:30:14 +0000
Subject: [PATCH] add logic to match image provider

---
 packages/core/src/generation.ts | 32 ++++++++++++++++++++++++--------
 1 file changed, 24 insertions(+), 8 deletions(-)

diff --git a/packages/core/src/generation.ts b/packages/core/src/generation.ts
index 59b620d6481..548a17843da 100644
--- a/packages/core/src/generation.ts
+++ b/packages/core/src/generation.ts
@@ -949,14 +949,30 @@ export const generateImage = async (
     });
 
     const apiKey =
-        runtime.imageModelProvider === runtime.modelProvider
-            ? runtime.token
-            : (runtime.getSetting("HEURIST_API_KEY") ??
-              runtime.getSetting("TOGETHER_API_KEY") ??
-              runtime.getSetting("FAL_API_KEY") ??
-              runtime.getSetting("OPENAI_API_KEY") ??
-              runtime.getSetting("VENICE_API_KEY"));
-
+    runtime.imageModelProvider === runtime.modelProvider
+        ? runtime.token
+        : (() => {
+            // First try to match the specific provider
+            switch (runtime.imageModelProvider) {
+                case ModelProviderName.HEURIST:
+                    return runtime.getSetting("HEURIST_API_KEY");
+                case ModelProviderName.TOGETHER:
+                    return runtime.getSetting("TOGETHER_API_KEY");
+                case ModelProviderName.FAL:
+                    return runtime.getSetting("FAL_API_KEY");
+                case ModelProviderName.OPENAI:
+                    return runtime.getSetting("OPENAI_API_KEY");
+                case ModelProviderName.VENICE:
+                    return runtime.getSetting("VENICE_API_KEY");
+                default:
+                    // If no specific match, try the fallback chain
+                    return (runtime.getSetting("HEURIST_API_KEY") ??
+                           runtime.getSetting("TOGETHER_API_KEY") ??
+                           runtime.getSetting("FAL_API_KEY") ??
+                           runtime.getSetting("OPENAI_API_KEY") ??
+                           runtime.getSetting("VENICE_API_KEY"));
+            }
+        })();
     try {
         if (runtime.imageModelProvider === ModelProviderName.HEURIST) {
             const response = await fetch(