diff --git a/lib/crewai-tools/src/crewai_tools/tools/crewai_platform_tools/misc.py b/lib/crewai-tools/src/crewai_tools/tools/crewai_platform_tools/misc.py index 06cf7147d..8ef68bba0 100644 --- a/lib/crewai-tools/src/crewai_tools/tools/crewai_platform_tools/misc.py +++ b/lib/crewai-tools/src/crewai_tools/tools/crewai_platform_tools/misc.py @@ -1,5 +1,7 @@ import os +from crewai.context import get_platform_integration_token as _get_context_token + def get_platform_api_base_url() -> str: """Get the platform API base URL from environment or use default.""" @@ -8,10 +10,16 @@ def get_platform_api_base_url() -> str: def get_platform_integration_token() -> str: - """Get the platform API base URL from environment or use default.""" - token = os.getenv("CREWAI_PLATFORM_INTEGRATION_TOKEN") or "" + """Get the platform integration token from the context. + Fallback to the environment variable if no token has been set in the context. + + Raises: + ValueError: If no token has been set in the context. + """ + token = _get_context_token() or os.getenv("CREWAI_PLATFORM_INTEGRATION_TOKEN") if not token: raise ValueError( - "No platform integration token found, please set the CREWAI_PLATFORM_INTEGRATION_TOKEN environment variable" + "No platform integration token found. " + "Set it via platform_integration_context() or set_platform_integration_token()." ) - return token # TODO: Use context manager to get token + return token diff --git a/lib/crewai/src/crewai/context.py b/lib/crewai/src/crewai/context.py index bf73a221c..633eeae1a 100644 --- a/lib/crewai/src/crewai/context.py +++ b/lib/crewai/src/crewai/context.py @@ -10,40 +10,40 @@ _platform_integration_token: contextvars.ContextVar[str | None] = ( ) -def set_platform_integration_token(integration_token: str) -> None: +def set_platform_integration_token(integration_token: str) -> contextvars.Token[str | None]: """Set the platform integration token in the current context. Args: integration_token: The integration token to set. """ - _platform_integration_token.set(integration_token) + return _platform_integration_token.set(integration_token) + + +def reset_platform_integration_token(token: contextvars.Token[str | None]) -> None: + """Reset the platform integration token to its previous value.""" + _platform_integration_token.reset(token) def get_platform_integration_token() -> str | None: - """Get the platform integration token from the current context or environment. - + """Get the platform integration token from the current context. Returns: The integration token if set, otherwise None. """ - token = _platform_integration_token.get() - if token is None: - token = os.getenv("CREWAI_PLATFORM_INTEGRATION_TOKEN") - return token + return _platform_integration_token.get() @contextmanager -def platform_context(integration_token: str) -> Generator[None, Any, None]: +def platform_integration_context(integration_token: str) -> Generator[None, Any, None]: """Context manager to temporarily set the platform integration token. Args: integration_token: The integration token to set within the context. """ - token = _platform_integration_token.set(integration_token) + token = set_platform_integration_token(integration_token) try: yield finally: - _platform_integration_token.reset(token) - + reset_platform_integration_token(token) _current_task_id: contextvars.ContextVar[str | None] = contextvars.ContextVar( "current_task_id", default=None