Skip to content

Commit

Permalink
Merge pull request #55 from dougollerenshaw/add_config_manager
Browse files Browse the repository at this point in the history
Added ConfigManager to handle keys
  • Loading branch information
dougollerenshaw authored Oct 15, 2024
2 parents f6ca0a3 + c96c68d commit da8b9dc
Show file tree
Hide file tree
Showing 4 changed files with 77 additions and 180 deletions.
43 changes: 5 additions & 38 deletions codeaide/utils/api_utils.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -16,6 +15,7 @@
from codeaide.utils.logging_config import get_logger

logger = get_logger()
config_manager = ConfigManager()


class MissingAPIKeyException(Exception):
Expand All @@ -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() == "":
Expand All @@ -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)}")
Expand Down
71 changes: 71 additions & 0 deletions codeaide/utils/config_manager.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
import os
import platform
import sys
from pathlib import Path
from decouple import Config, RepositoryEnv, UndefinedValueError


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"
self._ensure_env_file()

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 _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

return keyring.get_password(
self.keyring_service, f"{provider.upper()}_API_KEY"
)
else:
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:
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")
3 changes: 1 addition & 2 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -20,4 +20,3 @@ autoflake
openai-whisper
sounddevice
scipy
ffmpeg-python
140 changes: 0 additions & 140 deletions tests/utils/test_api_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
check_api_connection,
parse_response,
send_api_request,
get_api_client,
)
from codeaide.utils.constants import (
SYSTEM_PROMPT,
Expand Down Expand Up @@ -48,145 +47,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.
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.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("anthropic.Anthropic")
def test_get_api_client_success(self, mock_anthropic, mock_auto_config):
"""
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_auto_config (MagicMock): A mock object for the AutoConfig class.
The test performs the following steps:
1. Mocks the AutoConfig 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_client = Mock()
mock_anthropic.return_value = mock_client

client = get_api_client(provider="anthropic")
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:
"""
A test class for the send_api_request function.
Expand Down

0 comments on commit da8b9dc

Please sign in to comment.