refactor

Author: deadbits

cascade.py

@@ -48,8 +48,8 @@ def load_config(config_file: str) -> tuple[Config, str]:
     )
 
     llm_wrappers = {
-        "llm1": LLMFactory.create(config.llm1.provider, config.llm1.model),
-        "llm2": LLMFactory.create(config.llm2.provider, config.llm2.model)
+        "llm1": LLMFactory.create(config.llm1),
+        "llm2": LLMFactory.create(config.llm2)
     }
 
     state_manager = StateManager(config)

cascade/llm/anthropic.py

@@ -12,20 +12,15 @@
 class AnthropicWrapper(BaseLLMWrapper):
     """Wrapper for the Anthropic API."""
 
-    def __init__(self, model: str):
+    def __init__(self, model: str, api_key: str):
         """Initialize the Anthropic wrapper.
 
         Args:
             model: The model identifier to use
+            api_key: The Anthropic API key
         """
         super().__init__(model)
-        self.name = "anthropic"
-        if os.environ.get("ANTHROPIC_API_KEY"):
-            self.client = anthropic.Anthropic(
-                api_key=os.environ.get("ANTHROPIC_API_KEY")
-            )
-        else:
-            raise ValueError("Anthropic API key not found")
+        self.client = anthropic.Anthropic(api_key=api_key)
 
     def generate_stream(
         self, messages, system_prompt=None

cascade/llm/base.py

@@ -17,7 +17,7 @@ def __init__(self, model: str):
             model: The model identifier to use
         """
         self.model = model
-        self.name = None
+        self.name = self.__class__.__name__.replace("Wrapper", "").lower()
 
     @abstractmethod
     def generate_stream(

cascade/llm/factory.py

@@ -1,6 +1,7 @@
 """Create LLM wrappers."""
 
 from typing import Dict, Type
+import os
 from cascade.models import Provider
 from cascade.llm.base import BaseLLMWrapper
 from cascade.llm.anthropic import AnthropicWrapper
@@ -18,13 +19,19 @@ class LLMFactory:
     }
 
     @classmethod
-    def create(cls, provider: Provider, model: str) -> BaseLLMWrapper:
-        """Create an LLM wrapper instance."""
+    def create(cls, config) -> BaseLLMWrapper:
+        """Create an LLM wrapper instance from an LLMConfig."""
+        provider = config.provider
         if provider not in cls._registry:
             raise ValueError(f"Unknown provider: {provider}")
 
         wrapper_class = cls._registry[provider]
-        return wrapper_class(model=model)
+        if provider in [Provider.OPENAI, Provider.ANTHROPIC]:
+            config.require_api_key()
+            api_key = os.environ[provider.upper() + "_API_KEY"]
+            return wrapper_class(model=config.model, api_key=api_key)
+        else:
+            return wrapper_class(model=config.model)
 
     @classmethod
     def register(cls, provider: Provider, wrapper_class: Type[BaseLLMWrapper]) -> None:

cascade/llm/ollama.py

@@ -18,7 +18,6 @@ def __init__(self, model: str):
             model: The model identifier to use
         """
         super().__init__(model)
-        self.name = "ollama"
 
     def generate_stream(
         self, messages, system_prompt=None

cascade/llm/openai.py

@@ -18,18 +18,15 @@
 class OpenAIWrapper(BaseLLMWrapper):
     """Wrapper for OpenAI API"""
 
-    def __init__(self, model: str):
+    def __init__(self, model: str, api_key: str):
         """Initialize the OpenAI wrapper.
 
         Args:
             model: The model identifier to use
+            api_key: The OpenAI API key
         """
         super().__init__(model)
-        self.name = "openai"
-        if os.environ.get("OPENAI_API_KEY"):
-            self.client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
-        else:
-            raise ValueError("OpenAI API key not found in environment variables")
+        self.client = OpenAI(api_key=api_key)
 
     def generate_stream(
         self, messages, system_prompt=None

cascade/models.py

@@ -3,6 +3,7 @@
 from typing import List, Dict, Optional
 from enum import Enum
 from pydantic import BaseModel, model_validator, field_validator
+import os
 
 
 class Provider(str, Enum):
@@ -40,8 +41,10 @@ def validate_connection(cls, v: str) -> str:
         provider, model = v.split(":", 1)
         if not provider or not model:
             raise ValueError("Both provider and model must be specified")
-        if provider not in ["anthropic", "openai", "ollama"]:
-            raise ValueError("Provider must be one of: anthropic, openai, ollama")
+        try:
+            Provider(provider)
+        except ValueError:
+            raise ValueError(f"Provider must be one of: {[p.value for p in Provider]}")
         return v
 
     @property
@@ -54,6 +57,16 @@ def model(self) -> str:
         """Get the model from the connection string."""
         return self.connection.split(":", 1)[1]
 
+    def require_api_key(self):
+        """Ensure the required API key is present in the environment for providers that need it."""
+        env_vars = {
+            "openai": "OPENAI_API_KEY",
+            "anthropic": "ANTHROPIC_API_KEY",
+        }
+        env_var = env_vars.get(self.provider)
+        if env_var and not os.environ.get(env_var):
+            raise ValueError(f"{env_var} must be set in the environment for provider '{self.provider}'")
+
 
 class Config(BaseModel):
     """YAML Config."""