Skip to content

Commit

Permalink
feature flag formatting and linting
Browse files Browse the repository at this point in the history
  • Loading branch information
aarushik93 committed Nov 20, 2024
1 parent c1c3345 commit 4e09cef
Show file tree
Hide file tree
Showing 7 changed files with 375 additions and 105 deletions.
Empty file.
150 changes: 150 additions & 0 deletions autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,150 @@
import asyncio
import contextlib
import logging
from functools import wraps
from typing import Any, Callable, Dict, Optional, TypeVar, cast

import ldclient
from fastapi import HTTPException, Request
from ldclient import Context, LDClient
from ldclient.config import Config
from typing_extensions import Concatenate, ParamSpec
from .config import settings

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG)

P = ParamSpec("P")
T = TypeVar("T")


def get_client() -> LDClient:
"""Get the LaunchDarkly client singleton."""
return ldclient.get()


def initialize_launchdarkly() -> None:
sdk_key = settings.SDK_KEY
logger.debug(f"Initializing LaunchDarkly with SDK key: {'present' if sdk_key else 'missing'}")

if not sdk_key:
logger.warning("LaunchDarkly SDK key not configured")
return

config = Config(sdk_key)
ldclient.set_config(config)

if ldclient.get().is_initialized():
logger.info("LaunchDarkly client initialized successfully")
else:
logger.error("LaunchDarkly client failed to initialize")

def shutdown_launchdarkly() -> None:
"""Shutdown the LaunchDarkly client."""
if ldclient.get().is_initialized():
ldclient.get().close()
logger.info("LaunchDarkly client closed successfully")

def create_context(user_id: str, additional_attributes: Optional[Dict[str, Any]] = None) -> Context:
"""Create LaunchDarkly context with optional additional attributes."""
builder = Context.builder(user_id).kind("user")
if additional_attributes:
for key, value in additional_attributes.items():
builder.set(key, value)
return builder.build()


def feature_flag(
flag_key: str,
default: bool = False,
unauthorized_response: Any = None,
):
"""
Decorator for feature flag protected endpoints.
"""
def decorator(
func: Callable[P, T]
) -> Callable[P, T]:
@wraps(func)
async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
try:
# Get user_id from kwargs since FastAPI injects it
user_id = kwargs.get('user_id')
logger.debug(f"Checking flag '{flag_key}' for user '{user_id}'")

if not get_client().is_initialized():
logger.warning(f"LaunchDarkly not initialized, using default={default}")
is_enabled = default
else:
context = create_context(user_id)
logger.debug(f"Created context for user: {context}")
is_enabled = get_client().variation(flag_key, context, default)
logger.debug(f"Flag '{flag_key}' evaluation result: {is_enabled}")

if not is_enabled:
logger.debug(f"Feature '{flag_key}' disabled, returning unauthorized response")
if unauthorized_response is not None:
return cast(T, unauthorized_response)
raise HTTPException(status_code=404, detail="Feature not available")

logger.debug(f"Feature '{flag_key}' enabled, executing function")
return await func(*args, **kwargs)
except Exception as e:
logger.error(f"Error evaluating feature flag {flag_key}: {e}", exc_info=True)
return cast(T, unauthorized_response) if unauthorized_response else None

@wraps(func)
def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
try:
user_id = kwargs.get('user_id')
if not get_client().is_initialized():
logger.warning(f"LaunchDarkly not initialized, using default={default}")
is_enabled = default
else:
context = create_context(user_id)
is_enabled = get_client().variation(flag_key, context, default)

if not is_enabled:
if unauthorized_response is not None:
return cast(T, unauthorized_response)
raise HTTPException(status_code=404, detail="Feature not available")

return func(*args, **kwargs)
except Exception as e:
logger.error(f"Error evaluating feature flag {flag_key}: {e}")
return cast(T, unauthorized_response) if unauthorized_response else None

return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper

return decorator


@contextlib.contextmanager
def mock_flag_variation(flag_key: str, return_value: Any):
"""Context manager for testing feature flags."""
original_variation = get_client().variation
get_client().variation = lambda key, context, default: return_value if key == flag_key else original_variation(key,
context,
default)
try:
yield
finally:
get_client().variation = original_variation


def percentage_rollout(
flag_key: str,
default: bool = False,
unauthorized_response: Any = None,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
"""Decorator for percentage-based rollouts."""
return feature_flag(flag_key, default, unauthorized_response)


def beta_feature(
flag_key: str | None = None,
unauthorized_response: Any = {"message": "Not available in beta"}
) -> Callable[[Callable[P, T]], Callable[P, T]]:
"""Decorator for beta features."""
actual_key = f"beta-{flag_key}" if flag_key else "beta"
return feature_flag(actual_key, False, unauthorized_response)
12 changes: 12 additions & 0 deletions autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
import os

from dotenv import load_dotenv

load_dotenv()


class Settings:
SDK_KEY: str = os.getenv("LAUNCH_DARKLY_SDK_KEY", "")


settings = Settings()
8 changes: 8 additions & 0 deletions autogpt_platform/backend/backend/server/rest_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,10 @@
import fastapi.responses
import starlette.middleware.cors
import uvicorn
from autogpt_libs.feature_flag.client import (
initialize_launchdarkly,
shutdown_launchdarkly,
)

import backend.data.block
import backend.data.db
Expand All @@ -18,14 +22,18 @@
settings = backend.util.settings.Settings()
logger = logging.getLogger(__name__)

logging.getLogger("autogpt_libs").setLevel(logging.DEBUG)


@contextlib.asynccontextmanager
async def lifespan_context(app: fastapi.FastAPI):
await backend.data.db.connect()
await backend.data.block.initialize_blocks()
await backend.data.user.migrate_and_encrypt_user_integrations()
await backend.data.graph.fix_llm_provider_credentials()
initialize_launchdarkly()
yield
shutdown_launchdarkly()
await backend.data.db.disconnect()


Expand Down
45 changes: 38 additions & 7 deletions autogpt_platform/backend/backend/server/routers/v1.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
import asyncio
import logging
from collections import defaultdict
from typing import Annotated, Any, Dict, List
from typing import Annotated, Any, Dict, List, Union

import pydantic
from autogpt_libs.auth.middleware import auth_middleware
from autogpt_libs.feature_flag.client import feature_flag
from autogpt_libs.utils.cache import thread_cached
from fastapi import APIRouter, Body, Depends, HTTPException
from typing_extensions import Optional, TypedDict
Expand Down Expand Up @@ -554,10 +555,15 @@ async def update_configuration(

@v1_router.post(
"/api-keys",
response_model=CreateAPIKeyResponse,
response_model=Union[List[CreateAPIKeyResponse], Dict[str, str]],
tags=["api-keys"],
dependencies=[Depends(auth_middleware)],
)
@feature_flag(
"api-keys-enabled",
default=False,
unauthorized_response={"message": "API keys are not available"},
)
async def create_api_key(
request: CreateAPIKeyRequest, user_id: Annotated[str, Depends(get_user_id)]
) -> CreateAPIKeyResponse:
Expand All @@ -577,10 +583,15 @@ async def create_api_key(

@v1_router.get(
"/api-keys",
response_model=List[APIKeyWithoutHash],
response_model=Union[List[APIKeyWithoutHash], Dict[str, str]],
tags=["api-keys"],
dependencies=[Depends(auth_middleware)],
)
@feature_flag(
"api-keys-enabled",
default=False,
unauthorized_response={"message": "API keys are not available"},
)
async def get_api_keys(
user_id: Annotated[str, Depends(get_user_id)]
) -> List[APIKeyWithoutHash]:
Expand All @@ -594,10 +605,15 @@ async def get_api_keys(

@v1_router.get(
"/api-keys/{key_id}",
response_model=APIKeyWithoutHash,
response_model=Union[List[APIKeyWithoutHash], Dict[str, str]],
tags=["api-keys"],
dependencies=[Depends(auth_middleware)],
)
@feature_flag(
"api-keys-enabled",
default=False,
unauthorized_response={"message": "API keys are not available"},
)
async def get_api_key(
key_id: str, user_id: Annotated[str, Depends(get_user_id)]
) -> APIKeyWithoutHash:
Expand All @@ -614,10 +630,15 @@ async def get_api_key(

@v1_router.delete(
"/api-keys/{key_id}",
response_model=APIKeyWithoutHash,
response_model=Union[List[APIKeyWithoutHash], Dict[str, str]],
tags=["api-keys"],
dependencies=[Depends(auth_middleware)],
)
@feature_flag(
"api-keys-enabled",
default=False,
unauthorized_response={"message": "API keys are not available"},
)
async def delete_api_key(
key_id: str, user_id: Annotated[str, Depends(get_user_id)]
) -> Optional[APIKeyWithoutHash]:
Expand All @@ -635,10 +656,15 @@ async def delete_api_key(

@v1_router.post(
"/api-keys/{key_id}/suspend",
response_model=APIKeyWithoutHash,
response_model=Union[List[APIKeyWithoutHash], Dict[str, str]],
tags=["api-keys"],
dependencies=[Depends(auth_middleware)],
)
@feature_flag(
"api-keys-enabled",
default=False,
unauthorized_response={"message": "API keys are not available"},
)
async def suspend_key(
key_id: str, user_id: Annotated[str, Depends(get_user_id)]
) -> Optional[APIKeyWithoutHash]:
Expand All @@ -656,10 +682,15 @@ async def suspend_key(

@v1_router.put(
"/api-keys/{key_id}/permissions",
response_model=APIKeyWithoutHash,
response_model=Union[List[APIKeyWithoutHash], Dict[str, str]],
tags=["api-keys"],
dependencies=[Depends(auth_middleware)],
)
@feature_flag(
"api-keys-enabled",
default=False,
unauthorized_response={"message": "API keys are not available"},
)
async def update_permissions(
key_id: str,
request: UpdatePermissionsRequest,
Expand Down
Loading

0 comments on commit 4e09cef

Please sign in to comment.