Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(platform/featureflags): Setting up feature flagging #8718

Merged
merged 22 commits into from
Nov 27, 2024
Merged
Show file tree
Hide file tree
Changes from 19 commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
4e09cef
feature flag formatting and linting
aarushik93 Nov 19, 2024
9e50c05
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 20, 2024
edcd2ae
add tests
aarushik93 Nov 20, 2024
9aeedd7
update poetry lock
aarushik93 Nov 20, 2024
4eaed30
remove unneeded changes
aarushik93 Nov 20, 2024
69a3083
fix pyproject
aarushik93 Nov 20, 2024
0da5ce5
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 20, 2024
76dc949
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 21, 2024
33b92d6
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 21, 2024
58f7882
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 22, 2024
d1cc871
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 25, 2024
1a4c210
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 25, 2024
a8173a7
fix formatting and linting
aarushik93 Nov 26, 2024
43705fb
pydantic settings
aarushik93 Nov 26, 2024
11e7bc7
address comments and format
aarushik93 Nov 26, 2024
ad838ab
alphabetize
aarushik93 Nov 26, 2024
3c01dad
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 26, 2024
c94e4f8
fix lockfile
aarushik93 Nov 26, 2024
1af97a5
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 26, 2024
365ad18
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 27, 2024
c347a40
Merge branch 'dev' into aarushikansal/feature-flagging-backend
aarushik93 Nov 27, 2024
106b077
fix conflicts
aarushik93 Nov 27, 2024
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,14 +1,17 @@
from typing import NamedTuple
import secrets
import hashlib
import secrets
from typing import NamedTuple


class APIKeyContainer(NamedTuple):
"""Container for API key parts."""

raw: str
prefix: str
postfix: str
hash: str


class APIKeyManager:
PREFIX: str = "agpt_"
PREFIX_LENGTH: int = 8
Expand All @@ -19,9 +22,9 @@ def generate_api_key(self) -> APIKeyContainer:
raw_key = f"{self.PREFIX}{secrets.token_urlsafe(32)}"
return APIKeyContainer(
raw=raw_key,
prefix=raw_key[:self.PREFIX_LENGTH],
postfix=raw_key[-self.POSTFIX_LENGTH:],
hash=hashlib.sha256(raw_key.encode()).hexdigest()
prefix=raw_key[: self.PREFIX_LENGTH],
postfix=raw_key[-self.POSTFIX_LENGTH :],
hash=hashlib.sha256(raw_key.encode()).hexdigest(),
)

def verify_api_key(self, provided_key: str, stored_hash: str) -> bool:
Expand Down
4 changes: 2 additions & 2 deletions autogpt_platform/autogpt_libs/autogpt_libs/auth/depends.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import fastapi

from .middleware import auth_middleware
from .models import User, DEFAULT_USER_ID, DEFAULT_EMAIL
from .config import Settings
from .middleware import auth_middleware
from .models import DEFAULT_USER_ID, User


def requires_user(payload: dict = fastapi.Depends(auth_middleware)) -> User:
Expand Down
Empty file.
167 changes: 167 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,167 @@
import asyncio
import contextlib
import logging
from functools import wraps
from typing import Any, Awaitable, Callable, Dict, Optional, TypeVar, Union, cast

import ldclient
from fastapi import HTTPException
from ldclient import Context, LDClient
from ldclient.config import Config
from typing_extensions import 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.launch_darkly_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(str(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,
) -> Callable[
[Callable[P, Union[T, Awaitable[T]]]], Callable[P, Union[T, Awaitable[T]]]
]:
"""
Decorator for feature flag protected endpoints.
"""

def decorator(
func: Callable[P, Union[T, Awaitable[T]]]
) -> Callable[P, Union[T, Awaitable[T]]]:
@wraps(func)
async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
try:
user_id = kwargs.get("user_id")
if not user_id:
raise ValueError("user_id is required")

if not get_client().is_initialized():
logger.warning(
f"LaunchDarkly not initialized, using default={default}"
)
is_enabled = default
else:
context = create_context(str(user_id))
is_enabled = get_client().variation(flag_key, context, default)

if not is_enabled:
raise HTTPException(status_code=404, detail="Feature not available")

result = func(*args, **kwargs)
if asyncio.iscoroutine(result):
return await result
return cast(T, result)
except Exception as e:
logger.error(f"Error evaluating feature flag {flag_key}: {e}")
raise

@wraps(func)
def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
try:
user_id = kwargs.get("user_id")
if not user_id:
raise ValueError("user_id is required")

if not get_client().is_initialized():
logger.warning(
f"LaunchDarkly not initialized, using default={default}"
)
is_enabled = default
else:
context = create_context(str(user_id))
is_enabled = get_client().variation(flag_key, context, default)

if not is_enabled:
raise HTTPException(status_code=404, detail="Feature not available")

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

return cast(
Callable[P, Union[T, Awaitable[T]]],
async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper,
)

return decorator


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


def beta_feature(
flag_key: Optional[str] = None,
unauthorized_response: Any = {"message": "Not available in beta"},
) -> Callable[
[Callable[P, Union[T, Awaitable[T]]]], Callable[P, Union[T, Awaitable[T]]]
]:
"""Decorator for beta features."""
actual_key = f"beta-{flag_key}" if flag_key else "beta"
return feature_flag(actual_key, False)


@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
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
import pytest
from autogpt_libs.feature_flag.client import feature_flag, mock_flag_variation
from ldclient import LDClient


@pytest.fixture
def ld_client(mocker):
client = mocker.Mock(spec=LDClient)
mocker.patch("ldclient.get", return_value=client)
client.is_initialized.return_value = True
return client


@pytest.mark.asyncio
async def test_feature_flag_enabled(ld_client):
ld_client.variation.return_value = True

@feature_flag("test-flag")
async def test_function(user_id: str):
return "success"

result = test_function(user_id="test-user")
assert result == "success"
ld_client.variation.assert_called_once()


@pytest.mark.asyncio
async def test_feature_flag_unauthorized_response(ld_client):
ld_client.variation.return_value = False

@feature_flag("test-flag")
async def test_function(user_id: str):
return "success"

result = test_function(user_id="test-user")
assert result == {"error": "disabled"}


def test_mock_flag_variation(ld_client):
with mock_flag_variation("test-flag", True):
assert ld_client.variation("test-flag", None, False)

with mock_flag_variation("test-flag", False):
assert ld_client.variation("test-flag", None, False)
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class Settings(BaseSettings):
launch_darkly_sdk_key: str = Field(
default="",
description="The Launch Darkly SDK key",
validation_alias="LAUNCH_DARKLY_SDK_KEY"
)

model_config = SettingsConfigDict(case_sensitive=True, extra="ignore")


SETTINGS = Settings()
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

from pydantic import Field, field_validator
from pydantic_settings import BaseSettings, SettingsConfigDict

from .filters import BelowLevelFilter
from .formatters import AGPTFormatter, StructuredLoggingFormatter

Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from .store import SupabaseIntegrationCredentialsStore
from .types import Credentials, APIKeyCredentials, OAuth2Credentials
from .types import APIKeyCredentials, Credentials, OAuth2Credentials

__all__ = [
"SupabaseIntegrationCredentialsStore",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,14 @@
from pydantic import SecretStr

if TYPE_CHECKING:
from redis import Redis
from backend.executor.database import DatabaseManager
from redis import Redis

from autogpt_libs.utils.cache import thread_cached
from autogpt_libs.utils.synchronize import RedisKeyedMutex

from backend.util.settings import Settings

from .types import (
APIKeyCredentials,
Credentials,
Expand All @@ -19,8 +21,6 @@
UserIntegrations,
)

from backend.util.settings import Settings

settings = Settings()

revid_credentials = APIKeyCredentials(
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from typing import Callable, TypeVar, ParamSpec
import threading
from typing import Callable, ParamSpec, TypeVar

P = ParamSpec("P")
R = TypeVar("R")
Expand Down
Loading