Add key validation on all models (#1578)

## Description There has been various issues raised around incorrect errors when API keys are not set. This should address it for all models. Fixes # (issue) ## Type of change Please check the options that are relevant: - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] Model update - [ ] Infrastructure change ## Checklist - [x] My code follows Phidata's style guidelines and best practices - [x] I have performed a self-review of my code - [x] I have added docstrings and comments for complex logic - [x] My changes generate no new warnings or errors - [x] I have added cookbook examples for my new addition (if needed) - [x] I have updated requirements.txt/pyproject.toml (if needed) - [x] I have verified my changes in a clean environment --------- Co-authored-by: Dirk Brand <[email protected]>
phidatahq · Dec 16, 2024 · 6d4b9db · 6d4b9db
1 parent b27c973
commit 6d4b9db
Show file tree

Hide file tree

Showing 18 changed files with 76 additions and 38 deletions.
diff --git a/phi/llm/openai/chat.py b/phi/llm/openai/chat.py
@@ -1,3 +1,5 @@
+import os
+
 import httpx
 from typing import Optional, List, Iterator, Dict, Any, Union, Tuple
 
@@ -73,6 +75,10 @@ def get_client(self) -> OpenAIClient:
         if self.openai_client:
             return self.openai_client
 
+        self.api_key = self.api_key or os.getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            logger.error("OPENAI_API_KEY not set. Please set the OPENAI_API_KEY environment variable.")
+
         _client_params: Dict[str, Any] = {}
         if self.api_key:
             _client_params["api_key"] = self.api_key

diff --git a/phi/model/InternLM/internlm.py b/phi/model/InternLM/internlm.py
@@ -1,5 +1,6 @@
 from os import getenv
 from typing import Optional
+
 from phi.model.openai.like import OpenAILike
 
 
@@ -19,5 +20,5 @@ class InternLM(OpenAILike):
     name: str = "InternLM"
     provider: str = "InternLM"
 
-    api_key: Optional[str] = getenv("INTERNLM_API_KEY")
+    api_key: Optional[str] = getenv("INTERNLM_API_KEY", None)
     base_url: Optional[str] = "https://internlm-chat.intern-ai.org.cn/puyu/api/v1/chat/completions"
diff --git a/phi/model/anthropic/claude.py b/phi/model/anthropic/claude.py
@@ -1,4 +1,5 @@
 import json
+from os import getenv
 from dataclasses import dataclass, field
 from typing import Optional, List, Iterator, Dict, Any, Union, Tuple
 
@@ -18,9 +19,8 @@
         RawContentBlockDeltaEvent,
         ContentBlockStopEvent,
     )
-except ImportError:
-    logger.error("`anthropic` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`anthropic` not installed. Please install using `pip install anthropic`")
 
 
 @dataclass
@@ -89,6 +89,10 @@ def get_client(self) -> AnthropicClient:
         if self.client:
             return self.client
 
+        self.api_key = self.api_key or getenv("ANTHROPIC_API_KEY")
+        if not self.api_key:
+            logger.error("ANTHROPIC_API_KEY not set. Please set the ANTHROPIC_API_KEY environment variable.")
+
         _client_params: Dict[str, Any] = {}
         # Set client parameters if they are provided
         if self.api_key:

diff --git a/phi/model/azure/openai_chat.py b/phi/model/azure/openai_chat.py
@@ -1,15 +1,13 @@
 from os import getenv
 from typing import Optional, Dict, Any
-from phi.utils.log import logger
 from phi.model.openai.like import OpenAILike
 import httpx
 
 try:
     from openai import AzureOpenAI as AzureOpenAIClient
     from openai import AsyncAzureOpenAI as AsyncAzureOpenAIClient
-except ImportError:
-    logger.error("`azure openai` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`azure openai` not installed. Please install using `pip install openai`")
 
 
 class AzureOpenAIChat(OpenAILike):

diff --git a/phi/model/cohere/chat.py b/phi/model/cohere/chat.py
@@ -1,4 +1,5 @@
 import json
+from os import getenv
 from dataclasses import dataclass, field
 from typing import Optional, List, Iterator, Dict, Any, Tuple
 
@@ -28,9 +29,8 @@
     )
     from cohere.types.api_meta_tokens import ApiMetaTokens
     from cohere.types.api_meta import ApiMeta
-except ImportError:
-    logger.error("`cohere` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`cohere` not installed. Please install using `pip install cohere`")
 
 
 @dataclass
@@ -72,6 +72,11 @@ def client(self) -> CohereClient:
             return self.cohere_client
 
         _client_params: Dict[str, Any] = {}
+
+        self.api_key = self.api_key or getenv("CO_API_KEY")
+        if not self.api_key:
+            logger.error("CO_API_KEY not set. Please set the CO_API_KEY environment variable.")
+
         if self.api_key:
             _client_params["api_key"] = self.api_key
         return CohereClient(**_client_params)

diff --git a/phi/model/deepseek/deepseek.py b/phi/model/deepseek/deepseek.py
@@ -20,5 +20,5 @@ class DeepSeekChat(OpenAILike):
     name: str = "DeepSeekChat"
     provider: str = "DeepSeek"
 
-    api_key: Optional[str] = getenv("DEEPSEEK_API_KEY")
+    api_key: Optional[str] = getenv("DEEPSEEK_API_KEY", None)
     base_url: str = "https://api.deepseek.com"
diff --git a/phi/model/fireworks/fireworks.py b/phi/model/fireworks/fireworks.py
@@ -22,7 +22,7 @@ class Fireworks(OpenAILike):
     name: str = "Fireworks: " + id
     provider: str = "Fireworks"
 
-    api_key: Optional[str] = getenv("FIREWORKS_API_KEY")
+    api_key: Optional[str] = getenv("FIREWORKS_API_KEY", None)
     base_url: str = "https://api.fireworks.ai/inference/v1"
 
     def invoke_stream(self, messages: List[Message]) -> Iterator[ChatCompletionChunk]:

diff --git a/phi/model/google/gemini.py b/phi/model/google/gemini.py
@@ -1,3 +1,4 @@
+from os import getenv
 import time
 import json
 from pathlib import Path
@@ -24,8 +25,7 @@
     )
     from google.protobuf.struct_pb2 import Struct
 except (ModuleNotFoundError, ImportError):
-    logger.error("`google-generativeai` not installed. Please install it using `pip install google-generativeai`")
-    raise
+    raise ImportError("`google-generativeai` not installed. Please install it using `pip install google-generativeai`")
 
 
 @dataclass
@@ -103,9 +103,12 @@ def get_client(self) -> GenerativeModel:
             return self.client
 
         client_params: Dict[str, Any] = {}
-        # Set client parameters if they are provided
-        if self.api_key:
-            client_params["api_key"] = self.api_key
+
+        self.api_key = self.api_key or getenv("GOOGLE_API_KEY")
+        if not self.api_key:
+            logger.error("GOOGLE_API_KEY not set. Please set the GOOGLE_API_KEY environment variable.")
+        client_params["api_key"] = self.api_key
+
         if self.client_params:
             client_params.update(self.client_params)
         genai.configure(**client_params)

diff --git a/phi/model/google/gemini_openai.py b/phi/model/google/gemini_openai.py
@@ -1,5 +1,7 @@
 from os import getenv
 from typing import Optional
+
+
 from phi.model.openai.like import OpenAILike
 
 
@@ -19,5 +21,5 @@ class GeminiOpenAIChat(OpenAILike):
     name: str = "Gemini"
     provider: str = "Google"
 
-    api_key: Optional[str] = getenv("GOOGLE_API_KEY")
+    api_key: Optional[str] = getenv("GOOGLE_API_KEY", None)
     base_url: Optional[str] = "https://generativelanguage.googleapis.com/v1beta/"
diff --git a/phi/model/groq/groq.py b/phi/model/groq/groq.py
@@ -1,3 +1,4 @@
+from os import getenv
 from dataclasses import dataclass, field
 from typing import Optional, List, Iterator, Dict, Any, Union
 
@@ -16,9 +17,8 @@
     from groq.types.chat import ChatCompletion, ChatCompletionMessage
     from groq.types.chat.chat_completion_chunk import ChatCompletionChunk, ChoiceDeltaToolCall, ChoiceDelta
     from groq.types.completion_usage import CompletionUsage
-except ImportError:
-    logger.error("`groq` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`groq` not installed. Please install using `pip install groq`")
 
 
 @dataclass
@@ -104,6 +104,10 @@ class Groq(Model):
     async_client: Optional[AsyncGroqClient] = None
 
     def get_client_params(self) -> Dict[str, Any]:
+        self.api_key = self.api_key or getenv("GROQ_API_KEY")
+        if not self.api_key:
+            logger.error("GROQ_API_KEY not set. Please set the GROQ_API_KEY environment variable.")
+
         client_params: Dict[str, Any] = {}
         if self.api_key:
             client_params["api_key"] = self.api_key

diff --git a/phi/model/huggingface/hf.py b/phi/model/huggingface/hf.py
@@ -1,3 +1,4 @@
+from os import getenv
 from dataclasses import dataclass, field
 from typing import Optional, List, Iterator, Dict, Any, Union
 
@@ -23,9 +24,8 @@
         ChatCompletionOutputMessage,
         ChatCompletionOutputUsage,
     )
-except ImportError:
-    logger.error("`huggingface_hub` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`huggingface_hub` not installed. Please install using `pip install huggingface_hub`")
 
 
 @dataclass
@@ -129,6 +129,10 @@ class HuggingFaceChat(Model):
     async_client: Optional[AsyncInferenceClient] = None
 
     def get_client_params(self) -> Dict[str, Any]:
+        self.api_key = self.api_key or getenv("HF_TOKEN")
+        if not self.api_key:
+            logger.error("HF_TOKEN not set. Please set the HF_TOKEN environment variable.")
+
         _client_params: Dict[str, Any] = {}
         if self.api_key is not None:
             _client_params["api_key"] = self.api_key

diff --git a/phi/model/mistral/mistral.py b/phi/model/mistral/mistral.py
@@ -1,3 +1,4 @@
+from os import getenv
 from dataclasses import dataclass, field
 from typing import Optional, List, Iterator, Dict, Any, Union
 
@@ -14,9 +15,8 @@
     from mistralai.models.chatcompletionresponse import ChatCompletionResponse
     from mistralai.models.deltamessage import DeltaMessage
     from mistralai.types.basemodel import Unset
-except ImportError:
-    logger.error("`mistralai` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`mistralai` not installed. Please install using `pip install mistralai`")
 
 MistralMessage = Union[models.UserMessage, models.AssistantMessage, models.SystemMessage, models.ToolMessage]
 
@@ -91,6 +91,10 @@ def client(self) -> Mistral:
         if self.mistral_client:
             return self.mistral_client
 
+        self.api_key = self.api_key or getenv("MISTRAL_API_KEY")
+        if not self.api_key:
+            logger.error("MISTRAL_API_KEY not set. Please set the MISTRAL_API_KEY environment variable.")
+
         _client_params: Dict[str, Any] = {}
         if self.api_key:
             _client_params["api_key"] = self.api_key

diff --git a/phi/model/nvidia/nvidia.py b/phi/model/nvidia/nvidia.py
@@ -20,5 +20,5 @@ class Nvidia(OpenAILike):
     name: str = "Nvidia"
     provider: str = "Nvidia " + id
 
-    api_key: Optional[str] = getenv("NVIDIA_API_KEY")
+    api_key: Optional[str] = getenv("NVIDIA_API_KEY", None)
     base_url: str = "https://integrate.api.nvidia.com/v1"
diff --git a/phi/model/ollama/chat.py b/phi/model/ollama/chat.py
@@ -14,9 +14,8 @@
 
 try:
     from ollama import Client as OllamaClient, AsyncClient as AsyncOllamaClient
-except ImportError:
-    logger.error("`ollama` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`ollama` not installed. Please install using `pip install ollama`")
 
 
 @dataclass

diff --git a/phi/model/openai/chat.py b/phi/model/openai/chat.py
@@ -1,3 +1,4 @@
+from os import getenv
 from dataclasses import dataclass, field
 from typing import Optional, List, Iterator, Dict, Any, Union
 
@@ -23,9 +24,8 @@
         ChoiceDeltaToolCall,
     )
     from openai.types.chat.chat_completion_message import ChatCompletionMessage
-except ImportError:
-    logger.error("`openai` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError("`openai` not installed. Please install using `pip install openai`")
 
 
 @dataclass
@@ -118,6 +118,11 @@ class OpenAIChat(Model):
 
     def get_client_params(self) -> Dict[str, Any]:
         client_params: Dict[str, Any] = {}
+
+        self.api_key = self.api_key or getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            logger.error("OPENAI_API_KEY not set. Please set the OPENAI_API_KEY environment variable.")
+
         if self.api_key is not None:
             client_params["api_key"] = self.api_key
         if self.organization is not None:

diff --git a/phi/model/vertexai/gemini.py b/phi/model/vertexai/gemini.py
@@ -21,9 +21,10 @@
         Content,
         Part,
     )
-except ImportError:
-    logger.error("`google-cloud-aiplatform` not installed")
-    raise
+except (ModuleNotFoundError, ImportError):
+    raise ImportError(
+        "`google-cloud-aiplatform` not installed. Please install using `pip install google-cloud-aiplatform`"
+    )
 
 
 @dataclass

diff --git a/phi/model/xai/xai.py b/phi/model/xai/xai.py
@@ -1,5 +1,6 @@
 from os import getenv
 from typing import Optional
+
 from phi.model.openai.like import OpenAILike
 
 

diff --git a/phi/workspace/config.py b/phi/workspace/config.py
@@ -19,6 +19,7 @@ def get_workspace_objects_from_file(resource_file: Path) -> dict:
     """Returns workspace objects from the resource file"""
     from phi.aws.resources import AwsResources
     from phi.docker.resources import DockerResources
+
     try:
         python_objects = get_python_objects_from_module(resource_file)
         # logger.debug(f"python_objects: {python_objects}")