Spaces:
Paused
Paused
File size: 3,328 Bytes
1ec3391 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
"""
Unified API client management for W&B operations.
This module provides a consistent pattern for managing W&B API instances
with per-request API keys, following the same pattern as WeaveApiClient.
"""
import os
from typing import Optional, Dict, Any
from contextvars import ContextVar
import wandb
from wandb_mcp_server.utils import get_rich_logger
logger = get_rich_logger(__name__)
# Context variable for storing the current request's API key
api_key_context: ContextVar[Optional[str]] = ContextVar('wandb_api_key', default=None)
class WandBApiManager:
"""
Manages W&B API instances with per-request API keys.
This class follows the same pattern as WeaveApiClient, providing
a consistent interface for all W&B operations that need API access.
"""
@staticmethod
def get_api_key() -> Optional[str]:
"""
Get the API key for the current request context.
Returns:
The API key from context, environment, or None.
"""
# First try context variable (set by middleware for HTTP requests)
api_key = api_key_context.get()
# Fallback to environment variable (for STDIO or testing)
if not api_key:
api_key = os.environ.get("WANDB_API_KEY")
return api_key
@staticmethod
def get_api(api_key: Optional[str] = None) -> wandb.Api:
"""
Get a W&B API instance with the specified or current API key.
Args:
api_key: Optional API key to use. If not provided, uses context or environment.
Returns:
A configured wandb.Api instance.
Raises:
ValueError: If no API key is available.
"""
if api_key is None:
api_key = WandBApiManager.get_api_key()
if not api_key:
raise ValueError(
"No W&B API key available. Provide api_key parameter or "
"ensure WANDB_API_KEY is set in environment or request context."
)
# Create API instance with the specific key
# According to docs: https://docs.wandb.ai/ref/python/public-api/
return wandb.Api(api_key=api_key)
@staticmethod
def set_context_api_key(api_key: str) -> Any:
"""
Set the API key in the current context.
Args:
api_key: The API key to set.
Returns:
A token that can be used to reset the context.
"""
return api_key_context.set(api_key)
@staticmethod
def reset_context_api_key(token: Any) -> None:
"""
Reset the API key context.
Args:
token: The token returned from set_context_api_key.
"""
api_key_context.reset(token)
def get_wandb_api(api_key: Optional[str] = None) -> wandb.Api:
"""
Convenience function to get a W&B API instance.
This is the primary function that should be used throughout the codebase
to get a W&B API instance with proper API key handling.
Args:
api_key: Optional API key. If not provided, uses context or environment.
Returns:
A configured wandb.Api instance.
"""
return WandBApiManager.get_api(api_key)
|