improve plugins

Signed-off-by: Agustín Ramiro Díaz <[email protected]>
yeagerai · Sep 9, 2024 · 13dca11 · 13dca11
1 parent 03664e9
commit 13dca11
Show file tree

Hide file tree

Showing 9 changed files with 68 additions and 9 deletions.
diff --git a/backend/node/create_nodes/default_providers/heuristai_dolphin-2.9-llama3-8b.json b/backend/node/create_nodes/default_providers/heuristai_dolphin-2.9-llama3-8b.json
@@ -5,5 +5,9 @@
   "config": {
     "temperature": 0.75,
     "max_tokens": 500
+  },
+  "plugin_config": {
+    "api_key_env_var": "HEURISTAIAPIKEY",
+    "url": "https://llm-gateway.heurist.xyz"
   }
 }
diff --git a/backend/node/create_nodes/default_providers/heuristai_meta-llamallama-2-70b-chat.json b/backend/node/create_nodes/default_providers/heuristai_meta-llamallama-2-70b-chat.json
@@ -5,5 +5,9 @@
   "config": {
     "temperature": 0.75,
     "max_tokens": 500
+  },
+  "plugin_config": {
+    "api_key_env_var": "HEURISTAIAPIKEY",
+    "url": "https://llm-gateway.heurist.xyz"
   }
 }
diff --git a/backend/node/create_nodes/default_providers/heuristai_mistralaimixtral-8x7b-instruct.json b/backend/node/create_nodes/default_providers/heuristai_mistralaimixtral-8x7b-instruct.json
@@ -5,5 +5,9 @@
   "config": {
     "temperature": 0.75,
     "max_tokens": 500
+  },
+  "plugin_config": {
+    "api_key_env_var": "HEURISTAIAPIKEY",
+    "url": "https://llm-gateway.heurist.xyz"
   }
 }
diff --git a/backend/node/create_nodes/default_providers/heuristai_openhermes-2-yi-34b-gptq.json b/backend/node/create_nodes/default_providers/heuristai_openhermes-2-yi-34b-gptq.json
@@ -5,5 +5,9 @@
   "config": {
     "temperature": 0.75,
     "max_tokens": 500
+  },
+  "plugin_config": {
+    "api_key_env_var": "HEURISTAIAPIKEY",
+    "url": "https://llm-gateway.heurist.xyz"
   }
 }
diff --git a/backend/node/create_nodes/default_providers/openai_gpt-4.json b/backend/node/create_nodes/default_providers/openai_gpt-4.json
@@ -2,5 +2,8 @@
   "provider": "openai",
   "plugin": "openai",
   "model": "gpt-4",
-  "config": "<empty>"
+  "config": "<empty>",
+  "plugin_config": {
+    "api_key_env_var": "OPENAIKEY"
+  }
 }
diff --git a/backend/node/create_nodes/default_providers/openai_gpt-4o-mini.json b/backend/node/create_nodes/default_providers/openai_gpt-4o-mini.json
@@ -2,5 +2,8 @@
   "provider": "openai",
   "plugin": "openai",
   "model": "gpt-4o-mini",
-  "config": "<empty>"
+  "config": "<empty>",
+  "plugin_config": {
+    "api_key_env_var": "OPENAIKEY"
+  }
 }
diff --git a/backend/node/create_nodes/default_providers/openai_gpt-4o.json b/backend/node/create_nodes/default_providers/openai_gpt-4o.json
@@ -2,5 +2,8 @@
   "provider": "openai",
   "plugin": "openai",
   "model": "gpt-4o",
-  "config": "<empty>"
+  "config": "<empty>",
+  "plugin_config": {
+    "api_key_env_var": "OPENAIKEY"
+  }
 }
diff --git a/backend/node/create_nodes/providers_schema.json b/backend/node/create_nodes/providers_schema.json
@@ -22,6 +22,9 @@
     "plugin": {
       "$comment": "plugin to be loaded by the simulator to interact with the provider",
       "enum": ["heuristai", "openai", "ollama"]
+    },
+    "plugin_config": {
+      "type": "object"
     }
   },
   "allOf": [
@@ -173,6 +176,21 @@
           "plugin": {
             "const": "heuristai"
           },
+          "plugin_config": {
+            "type": "object",
+            "additionalProperties": false,
+            "properties": {
+              "api_key_env_var": {
+                "type": "string",
+                "$comment": "Environment variable that contains the API key"
+              },
+              "url": {
+                "type": "string",
+                "$comment": "URL of the API endpoint"
+              }
+            },
+            "required": ["api_key_env_var", "url"]
+          },
           "model": {
             "enum": [
               "mistralai/mixtral-8x7b-instruct",
@@ -202,7 +220,8 @@
             },
             "required": ["temperature", "max_tokens"]
           }
-        }
+        },
+        "required": ["plugin_config"]
       }
     },
     {
@@ -216,13 +235,25 @@
           "plugin": {
             "const": "openai"
           },
+          "plugin_config": {
+            "type": "object",
+            "additionalProperties": false,
+            "properties": {
+              "api_key_env_var": {
+                "type": "string",
+                "$comment": "Environment variable that contains the API key"
+              }
+            },
+            "required": ["api_key_env_var"]
+          },
           "model": {
             "enum": ["gpt-3.5-turbo", "gpt-4", "gpt-4o", "gpt-4o-mini"]
           },
           "config": {
             "const": "<empty>"
           }
-        }
+        },
+        "required": ["plugin_config"]
       }
     }
   ],

diff --git a/backend/node/genvm/llms.py b/backend/node/genvm/llms.py
@@ -20,6 +20,8 @@
 
 load_dotenv()
 
+plugin_config_key = "plugin_config"
+
 
 async def process_streaming_buffer(buffer: str, chunk: str, regex: str) -> str:
     updated_buffer = buffer + chunk
@@ -77,7 +79,8 @@ async def call_openai(
     regex: Optional[str],
     return_streaming_channel: Optional[asyncio.Queue],
 ) -> str:
-    client = get_openai_client(os.environ.get("OPENAIKEY"))
+    api_key_env_var = model_config[plugin_config_key]["api_key_env_var"]
+    client = get_openai_client(os.environ.get(api_key_env_var))
     # TODO: OpenAI exceptions need to be caught here
     stream = get_openai_stream(client, prompt, model_config)
 
@@ -90,9 +93,9 @@ async def call_heuristai(
     regex: Optional[str],
     return_streaming_channel: Optional[asyncio.Queue],
 ) -> str:
-    client = get_openai_client(
-        os.environ.get("HEURISTAIAPIKEY"), os.environ.get("HEURISTAIURL")
-    )
+    api_key_env_var = model_config[plugin_config_key]["api_key_env_var"]
+    url = model_config[plugin_config_key]["url"]
+    client = get_openai_client(os.environ.get(api_key_env_var), os.environ.get(url))
     stream = get_openai_stream(client, prompt, model_config)
     # TODO: Get the line below working
     # return await get_openai_output(stream, regex, return_streaming_channel)