feat: allow users to configure models for openai and anthropic

storyprotocol · Dec 12, 2024 · 8b5175b · 8b5175b
1 parent d0c6c55
commit 8b5175b
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/.env.example b/.env.example
@@ -5,6 +5,11 @@ DISCORD_VOICE_CHANNEL_ID= # The ID of the voice channel the bot should join (opt
 
 # AI Model API Keys
 OPENAI_API_KEY= # OpenAI API key, starting with sk-
+SMALL_OPENAI_MODEL= # Default: gpt-4o-mini
+MEDIUM_OPENAI_MODEL= # Default: gpt-4o
+LARGE_OPENAI_MODEL= # Default: gpt-4o
+EMBEDDING_OPENAI_MODEL= # Default: text-embedding-3-small
+IMAGE_OPENAI_MODEL= # Default: dall-e-3
 
 # Eternal AI's Decentralized Inference API
 ETERNALAI_URL=
@@ -107,6 +112,10 @@ LLAMALOCAL_PATH= # Default: "" which is the current directory in plugin-node/dis
 
 # API Keys
 ANTHROPIC_API_KEY= # For Claude
+SMALL_ANTHROPIC_MODEL= # Default: claude-3-haiku-20240307
+MEDIUM_ANTHROPIC_MODEL= # Default: claude-3-5-sonnet-20241022
+LARGE_ANTHROPIC_MODEL= # Default: claude-3-5-sonnet-20241022
+
 HEURIST_API_KEY=   # Get from https://heurist.ai/dev-access
 
 # Heurist Models

diff --git a/packages/core/src/models.ts b/packages/core/src/models.ts
@@ -13,11 +13,11 @@ export const models: Models = {
             temperature: 0.6,
         },
         model: {
-            [ModelClass.SMALL]: "gpt-4o-mini",
-            [ModelClass.MEDIUM]: "gpt-4o",
-            [ModelClass.LARGE]: "gpt-4o",
-            [ModelClass.EMBEDDING]: "text-embedding-3-small",
-            [ModelClass.IMAGE]: "dall-e-3",
+            [ModelClass.SMALL]: settings.SMALL_OPENAI_MODEL || "gpt-4o-mini",
+            [ModelClass.MEDIUM]: settings.MEDIUM_OPENAI_MODEL || "gpt-4o",
+            [ModelClass.LARGE]: settings.LARGE_OPENAI_MODEL || "gpt-4o",
+            [ModelClass.EMBEDDING]: settings.EMBEDDING_OPENAI_MODEL || "text-embedding-3-small",
+            [ModelClass.IMAGE]: settings.IMAGE_OPENAI_MODEL || "dall-e-3",
         },
     },
     [ModelProviderName.ETERNALAI]: {
@@ -55,9 +55,9 @@ export const models: Models = {
         },
         endpoint: "https://api.anthropic.com/v1",
         model: {
-            [ModelClass.SMALL]: "claude-3-haiku-20240307",
-            [ModelClass.MEDIUM]: "claude-3-5-sonnet-20241022",
-            [ModelClass.LARGE]: "claude-3-5-sonnet-20241022",
+            [ModelClass.SMALL]: settings.SMALL_ANTHROPIC_MODEL || "claude-3-haiku-20240307",
+            [ModelClass.MEDIUM]: settings.MEDIUM_ANTHROPIC_MODEL || "claude-3-5-sonnet-20241022",
+            [ModelClass.LARGE]: settings.LARGE_ANTHROPIC_MODEL || "claude-3-5-sonnet-20241022",
         },
     },
     [ModelProviderName.CLAUDE_VERTEX]: {