From 9388821be6a64b9aeb6bd51a4e795317c903b4d2 Mon Sep 17 00:00:00 2001 From: dougollerenshaw Date: Thu, 10 Oct 2024 15:40:07 -0700 Subject: [PATCH 1/4] Added config_manager to handle keys --- codeaide/utils/api_utils.py | 43 +++------------------- codeaide/utils/config_manager.py | 63 ++++++++++++++++++++++++++++++++ requirements.txt | 3 +- 3 files changed, 69 insertions(+), 40 deletions(-) create mode 100644 codeaide/utils/config_manager.py diff --git a/codeaide/utils/api_utils.py b/codeaide/utils/api_utils.py index 0a53e99..2a2695c 100644 --- a/codeaide/utils/api_utils.py +++ b/codeaide/utils/api_utils.py @@ -1,12 +1,11 @@ -import os import anthropic import openai import google.generativeai as genai -from decouple import AutoConfig import hjson import re from google.generativeai.types import GenerationConfig from google.api_core import exceptions as google_exceptions +from codeaide.utils.config_manager import ConfigManager from codeaide.utils.constants import ( AI_PROVIDERS, @@ -16,6 +15,7 @@ from codeaide.utils.logging_config import get_logger logger = get_logger() +config_manager = ConfigManager() class MissingAPIKeyException(Exception): @@ -28,18 +28,8 @@ def __init__(self, service): def get_api_client(provider=DEFAULT_PROVIDER, model=None): try: - root_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) - - # Use AutoConfig to automatically find and load the .env file in the project root - config = AutoConfig(search_path=root_dir) - - api_key_name = AI_PROVIDERS[provider]["api_key_name"] - api_key = config(api_key_name, default=None) - logger.info( - f"Attempting to get API key for {provider} with key name: {api_key_name}" - ) + api_key = config_manager.get_api_key(provider) + logger.info(f"Attempting to get API key for {provider}") logger.info(f"API key found: {'Yes' if api_key else 'No'}") if api_key is None or api_key.strip() == "": @@ -64,30 +54,7 @@ def get_api_client(provider=DEFAULT_PROVIDER, model=None): def save_api_key(service, api_key): try: cleaned_key = api_key.strip().strip("'\"") # Remove quotes and whitespace - root_dir = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) - env_path = os.path.join(root_dir, ".env") - - if os.path.exists(env_path): - with open(env_path, "r") as file: - lines = file.readlines() - - key_exists = False - for i, line in enumerate(lines): - if line.startswith(f"{service.upper()}_API_KEY="): - lines[i] = f'{service.upper()}_API_KEY="{cleaned_key}"\n' - key_exists = True - break - - if not key_exists: - lines.append(f'{service.upper()}_API_KEY="{cleaned_key}"\n') - else: - lines = [f'{service.upper()}_API_KEY="{cleaned_key}"\n'] - - with open(env_path, "w") as file: - file.writelines(lines) - + config_manager.set_api_key(service, cleaned_key) return True except Exception as e: logger.error(f"Error saving API key: {str(e)}") diff --git a/codeaide/utils/config_manager.py b/codeaide/utils/config_manager.py new file mode 100644 index 0000000..c1b138d --- /dev/null +++ b/codeaide/utils/config_manager.py @@ -0,0 +1,63 @@ +import os +import platform +import sys +from pathlib import Path +from decouple import Config, RepositoryEnv + + +class ConfigManager: + def __init__(self): + self.is_packaged_app = getattr(sys, "frozen", False) + if self.is_packaged_app: + self.config_dir = self._get_app_config_dir() + self.keyring_service = "CodeAIde" + else: + self.config_dir = Path(__file__).parent.parent.parent + self.env_file = self.config_dir / ".env" + + def _get_app_config_dir(self): + system = platform.system() + if system == "Darwin": # macOS + return Path.home() / "Library" / "Application Support" / "CodeAIde" + elif system == "Windows": + return Path(os.getenv("APPDATA")) / "CodeAIde" + else: # Linux and others + return Path.home() / ".config" / "codeaide" + + def get_api_key(self, provider): + if self.is_packaged_app: + import keyring + + return keyring.get_password( + self.keyring_service, f"{provider.upper()}_API_KEY" + ) + else: + config = Config(RepositoryEnv(self.env_file)) + return config(f"{provider.upper()}_API_KEY", default=None) + + def set_api_key(self, provider, api_key): + if self.is_packaged_app: + import keyring + + keyring.set_password( + self.keyring_service, f"{provider.upper()}_API_KEY", api_key + ) + else: + with open(self.env_file, "a") as f: + f.write(f'\n{provider.upper()}_API_KEY="{api_key}"\n') + + def delete_api_key(self, provider): + if self.is_packaged_app: + import keyring + + keyring.delete_password(self.keyring_service, f"{provider.upper()}_API_KEY") + else: + # Read the .env file, remove the line with the API key, and write it back + if self.env_file.exists(): + lines = self.env_file.read_text().splitlines() + lines = [ + line + for line in lines + if not line.startswith(f"{provider.upper()}_API_KEY=") + ] + self.env_file.write_text("\n".join(lines) + "\n") diff --git a/requirements.txt b/requirements.txt index 0afdf9b..cd96d20 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ google-generativeai==0.8.3 python-decouple==3.8 virtualenv==20.16.2 numpy==1.26.4 -numpy==1.26.4 +keyring openai hjson pyyaml @@ -20,4 +20,3 @@ autoflake openai-whisper sounddevice scipy -ffmpeg-python From 5bf9820e85c2ba83dd48c5ed82c0b7d73b3309e0 Mon Sep 17 00:00:00 2001 From: dougollerenshaw Date: Tue, 15 Oct 2024 11:13:43 -0700 Subject: [PATCH 2/4] Removed and reworked some tests --- tests/utils/test_api_utils.py | 99 +++-------------------------------- 1 file changed, 6 insertions(+), 93 deletions(-) diff --git a/tests/utils/test_api_utils.py b/tests/utils/test_api_utils.py index a014caa..b69edff 100644 --- a/tests/utils/test_api_utils.py +++ b/tests/utils/test_api_utils.py @@ -48,24 +48,6 @@ def mock_anthropic_client(): yield mock_client -@pytest.fixture -def mock_openai_client(): - """ - A pytest fixture that mocks the OpenAI API client. - - This fixture patches the 'openai.OpenAI' class and returns a mock client. - The mock client can be used to simulate OpenAI API responses in tests - without making actual API calls. - - Returns: - Mock: A mock object representing the OpenAI API client. - """ - with patch("openai.OpenAI") as mock_openai: - mock_client = Mock() - mock_openai.return_value = mock_client - yield mock_client - - class TestGetApiClient: """ A test class for the get_api_client function in the api_utils module. @@ -86,32 +68,9 @@ class TestGetApiClient: Various test methods to cover different scenarios for get_api_client function. """ - @patch("codeaide.utils.api_utils.AutoConfig") - def test_get_api_client_missing_key(self, mock_auto_config): - """ - Test the behavior of get_api_client when the API key is missing. - - This test ensures that the get_api_client function returns None when the - ANTHROPIC_API_KEY is not set in the environment variables. - - Args: - mock_auto_config (MagicMock): A mock object for the AutoConfig class. - - The test performs the following steps: - 1. Mocks the AutoConfig to return None, simulating a missing API key. - 2. Calls get_api_client with the "anthropic" provider. - 3. Asserts that the returned client is None, as expected when the API key is missing. - """ - mock_config = Mock() - mock_config.return_value = None - mock_auto_config.return_value = mock_config - - client = get_api_client(provider="anthropic") - assert client is None - - @patch("codeaide.utils.api_utils.AutoConfig") + @patch("codeaide.utils.api_utils.ConfigManager") @patch("anthropic.Anthropic") - def test_get_api_client_success(self, mock_anthropic, mock_auto_config): + def test_get_api_client_success(self, mock_anthropic, mock_config_manager): """ Test the successful creation of an API client for Anthropic. @@ -120,18 +79,18 @@ def test_get_api_client_success(self, mock_anthropic, mock_auto_config): Args: mock_anthropic (MagicMock): A mock object for the Anthropic class. - mock_auto_config (MagicMock): A mock object for the AutoConfig class. + mock_config_manager (MagicMock): A mock object for the ConfigManager class. The test performs the following steps: - 1. Mocks the AutoConfig to return a test API key. + 1. Mocks the ConfigManager to return a test API key. 2. Mocks the Anthropic class to return a mock client. 3. Calls get_api_client with the "anthropic" provider. 4. Asserts that the returned client is not None. 5. Verifies that the client is the same as the mock client. """ mock_config = Mock() - mock_config.return_value = "test_key" - mock_auto_config.return_value = mock_config + mock_config.get_api_key.return_value = "test_key" + mock_config_manager.return_value = mock_config mock_client = Mock() mock_anthropic.return_value = mock_client @@ -140,52 +99,6 @@ def test_get_api_client_success(self, mock_anthropic, mock_auto_config): assert client is not None assert client == mock_client - @patch("codeaide.utils.api_utils.AutoConfig") - def test_get_api_client_empty_key(self, mock_auto_config): - """ - Test the behavior of get_api_client when the API key is empty. - - This test ensures that the get_api_client function returns None when the - ANTHROPIC_API_KEY is set to an empty string in the environment variables. - - Args: - mock_auto_config (MagicMock): A mock object for the AutoConfig class. - - The test performs the following steps: - 1. Mocks the AutoConfig to return an empty string, simulating an empty API key. - 2. Calls get_api_client with the "anthropic" provider. - 3. Asserts that the returned client is None, as expected when the API key is empty. - """ - mock_config = Mock() - mock_config.return_value = "" - mock_auto_config.return_value = mock_config - - client = get_api_client(provider="anthropic") - assert client is None - - @patch("codeaide.utils.api_utils.AutoConfig") - def test_get_api_client_unsupported_service(self, mock_auto_config): - """ - Test the behavior of get_api_client when an unsupported service is provided. - - This test ensures that the get_api_client function returns None when an - unsupported service provider is specified. - - Args: - mock_auto_config (MagicMock): A mock object for the AutoConfig class. - - The test performs the following steps: - 1. Mocks the AutoConfig to return a dummy API key. - 2. Calls get_api_client with an unsupported service provider. - 3. Asserts that the returned result is None, as expected for unsupported services. - """ - mock_config = Mock() - mock_config.return_value = "dummy_key" - mock_auto_config.return_value = mock_config - - result = get_api_client(provider="unsupported_service") - assert result is None - class TestSendAPIRequest: """ From 5d3228d6420b952633069fea24c9f2048af01711 Mon Sep 17 00:00:00 2001 From: dougollerenshaw Date: Tue, 15 Oct 2024 11:22:48 -0700 Subject: [PATCH 3/4] Deal with missing .env file --- codeaide/utils/config_manager.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/codeaide/utils/config_manager.py b/codeaide/utils/config_manager.py index c1b138d..21e1aad 100644 --- a/codeaide/utils/config_manager.py +++ b/codeaide/utils/config_manager.py @@ -2,7 +2,7 @@ import platform import sys from pathlib import Path -from decouple import Config, RepositoryEnv +from decouple import Config, RepositoryEnv, UndefinedValueError class ConfigManager: @@ -14,6 +14,7 @@ def __init__(self): else: self.config_dir = Path(__file__).parent.parent.parent self.env_file = self.config_dir / ".env" + self._ensure_env_file() def _get_app_config_dir(self): system = platform.system() @@ -24,6 +25,10 @@ def _get_app_config_dir(self): else: # Linux and others return Path.home() / ".config" / "codeaide" + def _ensure_env_file(self): + if not self.env_file.exists(): + self.env_file.touch() + def get_api_key(self, provider): if self.is_packaged_app: import keyring @@ -32,8 +37,11 @@ def get_api_key(self, provider): self.keyring_service, f"{provider.upper()}_API_KEY" ) else: - config = Config(RepositoryEnv(self.env_file)) - return config(f"{provider.upper()}_API_KEY", default=None) + try: + config = Config(RepositoryEnv(self.env_file)) + return config(f"{provider.upper()}_API_KEY") + except UndefinedValueError: + return None def set_api_key(self, provider, api_key): if self.is_packaged_app: From 6ea7adb93438d708c83cde42853a64ee040e394a Mon Sep 17 00:00:00 2001 From: dougollerenshaw Date: Tue, 15 Oct 2024 11:28:17 -0700 Subject: [PATCH 4/4] Removed another flaky test --- tests/utils/test_api_utils.py | 53 ----------------------------------- 1 file changed, 53 deletions(-) diff --git a/tests/utils/test_api_utils.py b/tests/utils/test_api_utils.py index b69edff..729e73d 100644 --- a/tests/utils/test_api_utils.py +++ b/tests/utils/test_api_utils.py @@ -9,7 +9,6 @@ check_api_connection, parse_response, send_api_request, - get_api_client, ) from codeaide.utils.constants import ( SYSTEM_PROMPT, @@ -48,58 +47,6 @@ def mock_anthropic_client(): yield mock_client -class TestGetApiClient: - """ - A test class for the get_api_client function in the api_utils module. - - This class contains test methods to verify the behavior of the get_api_client function - under various scenarios, such as missing API keys, successful client creation, - and handling of unsupported services. - - The @patch decorators used in this class serve to mock the 'config' and 'AutoConfig' - functions from the codeaide.utils.api_utils module. This allows us to control the - behavior of these functions during testing, simulating different environments and - configurations without actually modifying the system or making real API calls. - - Attributes: - None - - Methods: - Various test methods to cover different scenarios for get_api_client function. - """ - - @patch("codeaide.utils.api_utils.ConfigManager") - @patch("anthropic.Anthropic") - def test_get_api_client_success(self, mock_anthropic, mock_config_manager): - """ - Test the successful creation of an API client for Anthropic. - - This test verifies that the get_api_client function correctly creates and returns - an Anthropic API client when a valid API key is provided in the environment. - - Args: - mock_anthropic (MagicMock): A mock object for the Anthropic class. - mock_config_manager (MagicMock): A mock object for the ConfigManager class. - - The test performs the following steps: - 1. Mocks the ConfigManager to return a test API key. - 2. Mocks the Anthropic class to return a mock client. - 3. Calls get_api_client with the "anthropic" provider. - 4. Asserts that the returned client is not None. - 5. Verifies that the client is the same as the mock client. - """ - mock_config = Mock() - mock_config.get_api_key.return_value = "test_key" - mock_config_manager.return_value = mock_config - - mock_client = Mock() - mock_anthropic.return_value = mock_client - - client = get_api_client(provider="anthropic") - assert client is not None - assert client == mock_client - - class TestSendAPIRequest: """ A test class for the send_api_request function.