Elrond MCP

""" Centralized Google AI and Instructor client management. This module provides a unified functional interface for creating and managing Google AI clients integrated with the instructor library for structured outputs. Uses the new Google GenAI SDK (google-genai) with instructor's from_provider approach. """ import logging import os import instructor logger = logging.getLogger(__name__) # Module-level state _api_key: str | None = None _critique_client: instructor.AsyncInstructor | None = None _synthesis_client: instructor.AsyncInstructor | None = None # Default model configurations DEFAULT_CRITIQUE_MODEL = "gemini-2.5-flash" DEFAULT_SYNTHESIS_MODEL = "gemini-2.5-pro" def configure(api_key: str) -> None: """Configure Google AI with the provided API key. Args: api_key: Google AI API key Raises: ValueError: If API key is empty or None """ global _api_key if not api_key or not api_key.strip(): raise ValueError("API key cannot be empty") logger.info("Configuring Google AI client") _api_key = api_key # Set environment variable so instructor can pick it up os.environ["GEMINI_API_KEY"] = api_key def _ensure_configured() -> str: """Ensure Google AI is configured and return the API key. Returns: The configured API key Raises: ValueError: If no API key is available """ global _api_key if _api_key is None: # Try to get from environment _api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY") if not _api_key: raise ValueError( "Google AI API key is required. Set GEMINI_API_KEY environment " "variable or call configure(api_key) first." ) logger.info("Auto-configured Google AI from environment") # Ensure it's also set for instructor os.environ["GEMINI_API_KEY"] = _api_key return _api_key def get_critique_client( model_name: str = DEFAULT_CRITIQUE_MODEL, ) -> instructor.AsyncInstructor: """Get cached critique client, creating it if needed. Args: model_name: Gemini model name for critique tasks Returns: Configured async instructor client for critique tasks """ global _critique_client if _critique_client is None: _ensure_configured() logger.debug(f"Creating critique client for model: {model_name}") try: _critique_client = instructor.from_provider( f"google/{model_name}", async_client=True ) logger.debug("Successfully created critique client") except Exception as e: logger.error(f"Failed to create critique client: {e}") raise return _critique_client def get_synthesis_client( model_name: str = DEFAULT_SYNTHESIS_MODEL, ) -> instructor.AsyncInstructor: """Get cached synthesis client, creating it if needed. Args: model_name: Gemini model name for synthesis tasks Returns: Configured async instructor client for synthesis tasks """ global _synthesis_client if _synthesis_client is None: _ensure_configured() logger.debug(f"Creating synthesis client for model: {model_name}") try: _synthesis_client = instructor.from_provider( f"google/{model_name}", async_client=True ) logger.debug("Successfully created synthesis client") except Exception as e: logger.error(f"Failed to create synthesis client: {e}") raise return _synthesis_client def is_configured() -> bool: """Check if Google AI is configured. Returns: True if API key is available, False otherwise """ return ( _api_key is not None or os.getenv("GEMINI_API_KEY") is not None or os.getenv("GOOGLE_API_KEY") is not None ) def reset() -> None: """Reset all cached state. Useful for testing.""" global _api_key, _critique_client, _synthesis_client _api_key = None _critique_client = None _synthesis_client = None # Also clear from environment for clean state if "GEMINI_API_KEY" in os.environ: del os.environ["GEMINI_API_KEY"] logger.debug("Client state reset") # Convenience aliases for backward compatibility configure_gemini = configure create_critique_client = get_critique_client create_synthesis_client = get_synthesis_client get_api_key = _ensure_configured