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

* feature flag formatting and linting

* add tests

* update poetry lock

* remove unneeded changes

* fix pyproject

* fix formatting and linting

* pydantic settings

* address comments and format

* alphabetize

* fix lockfile

* fix conflicts

autogpt_platform/autogpt_libs/autogpt_libs/api_key/key_manager.py

@@ -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
@@ -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:

autogpt_platform/autogpt_libs/autogpt_libs/auth/depends.py

@@ -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:

autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/client.py

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

autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/client_test.py

@@ -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)

autogpt_platform/autogpt_libs/autogpt_libs/feature_flag/config.py

@@ -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()

autogpt_platform/autogpt_libs/autogpt_libs/logging/config.py

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

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/__init__.py

@@ -1,5 +1,5 @@
 from .store import SupabaseIntegrationCredentialsStore
-from .types import Credentials, APIKeyCredentials, OAuth2Credentials
+from .types import APIKeyCredentials, Credentials, OAuth2Credentials
 
 __all__ = [
     "SupabaseIntegrationCredentialsStore",

autogpt_platform/autogpt_libs/autogpt_libs/supabase_integration_credentials_store/store.py

@@ -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,
@@ -19,8 +21,6 @@
     UserIntegrations,
 )
 
-from backend.util.settings import Settings
-
 settings = Settings()
 
 revid_credentials = APIKeyCredentials(

autogpt_platform/autogpt_libs/autogpt_libs/utils/cache.py

@@ -1,5 +1,5 @@
-from typing import Callable, TypeVar, ParamSpec
 import threading
+from typing import Callable, ParamSpec, TypeVar
 
 P = ParamSpec("P")
 R = TypeVar("R")