Simplify LLM config (#120)

autogluon · Nov 13, 2024 · 96062ca · 96062ca
1 parent 64fb669
commit 96062ca
Show file tree

Hide file tree

Showing 6 changed files with 5 additions and 15 deletions.
diff --git a/src/autogluon_assistant/configs/best_quality.yaml b/src/autogluon_assistant/configs/best_quality.yaml
@@ -8,10 +8,8 @@ save_artifacts:
 feature_transformers:
   - _target_: autogluon_assistant.transformer.CAAFETransformer
     eval_model: lightgbm
-    llm_provider: bedrock
-    llm_model: "anthropic.claude-3-5-sonnet-20241022-v2:0"
-    # llm_provider: openai
-    # llm_model: gpt-4o-2024-08-06
+    llm_provider: ${llm.provider}
+    llm_model: ${llm.model}
     num_iterations: 5
     optimization_metric: roc
   - _target_: autogluon_assistant.transformer.OpenFETransformer
@@ -30,7 +28,6 @@ llm:
   provider: bedrock
   model: "anthropic.claude-3-5-sonnet-20241022-v2:0"
   # provider: openai
-  # api_key_location: OPENAI_API_KEY
   # model: gpt-4o-2024-08-06
   max_tokens: 512
   proxy_url: null

diff --git a/src/autogluon_assistant/configs/high_quality.yaml b/src/autogluon_assistant/configs/high_quality.yaml
@@ -17,7 +17,6 @@ llm:
   provider: bedrock
   model: "anthropic.claude-3-5-sonnet-20241022-v2:0"
   # provider: openai
-  # api_key_location: OPENAI_API_KEY
   # model: gpt-4o-2024-08-06
   max_tokens: 512
   proxy_url: null

diff --git a/src/autogluon_assistant/configs/medium_quality.yaml b/src/autogluon_assistant/configs/medium_quality.yaml
@@ -17,7 +17,6 @@ llm:
   provider: bedrock
   model: "anthropic.claude-3-5-sonnet-20241022-v2:0"
   # provider: openai
-  # api_key_location: OPENAI_API_KEY
   # model: gpt-4o-2024-08-06
   max_tokens: 512
   proxy_url: null

diff --git a/src/autogluon_assistant/llm/llm.py b/src/autogluon_assistant/llm/llm.py
@@ -140,10 +140,10 @@ def get_valid_providers(cls):
 
     @staticmethod
     def _get_openai_chat_model(config: DictConfig) -> AssistantChatOpenAI:
-        if config.api_key_location in os.environ:
-            api_key = os.environ[config.api_key_location]
+        if "OPENAI_API_KEY" in os.environ:
+            api_key = os.environ["OPENAI_API_KEY"]
         else:
-            raise Exception("OpenAI API env variable not set")
+            raise Exception("OpenAI API env variable OPENAI_API_KEY not set")
 
         logger.info(f"AGA is using model {config.model} from OpenAI to assist you with the task.")
 

diff --git a/src/autogluon_assistant/ui/constants.py b/src/autogluon_assistant/ui/constants.py
@@ -43,9 +43,6 @@
 # Provider configuration
 PROVIDER_MAPPING = {"Claude 3.5 with Amazon Bedrock": "bedrock", "GPT 4o": "openai"}
 
-# TODO: Remove model specific mappings
-API_KEY_LOCATION = {"GPT 4o": "OPENAI_API_KEY"}
-
 INITIAL_STAGE = {
     "Task Understanding": [],
     "Feature Generation": [],

diff --git a/src/autogluon_assistant/ui/pages/task.py b/src/autogluon_assistant/ui/pages/task.py
@@ -8,7 +8,6 @@
 import requests
 import streamlit as st
 from constants import (
-    API_KEY_LOCATION,
     BASE_DATA_DIR,
     CAPTIONS,
     DATASET_OPTIONS,
@@ -50,7 +49,6 @@ def update_config_overrides():
     if st.session_state.llm:
         config_overrides.append(f"llm.model={LLM_MAPPING[st.session_state.llm]}")
         config_overrides.append(f"llm.provider={PROVIDER_MAPPING[st.session_state.llm]}")
-        config_overrides.append(f"llm.api_key_location={API_KEY_LOCATION[st.session_state.llm]}")
 
     if not st.session_state.feature_generation:
         config_overrides.append("feature_transformers=[]")