Spaces:
Running
Running
| """Load Gemini models using LangChain.""" | |
| import os | |
| from getpass import getpass | |
| from langchain_google_genai import ChatGoogleGenerativeAI | |
| from chemgraph.models.supported_models import supported_gemini_models | |
| from chemgraph.utils.logging_config import setup_logger | |
| logger = setup_logger(__name__) | |
| def load_gemini_model( | |
| model_name: str, | |
| temperature: float, | |
| api_key: str = None, | |
| prompt: str = None, | |
| base_url: str = None, | |
| ) -> ChatGoogleGenerativeAI: | |
| """Load an Gemini chat model into LangChain. | |
| This function loads an Gemini model and configures it for use with LangChain. | |
| It handles API key management, including prompting for the key if not provided | |
| or if the provided key is invalid. | |
| Parameters | |
| ---------- | |
| model_name : str | |
| The name of the Gemini chat model to load. See supported_gemini_models for list | |
| of supported models. | |
| temperature : float | |
| Controls the randomness of the generated text. Higher values (e.g., 0.8) | |
| make the output more random, while lower values (e.g., 0.2) make it more | |
| deterministic. | |
| api_key : str, optional | |
| The Google API key. If not provided, the function will attempt to retrieve it | |
| from the environment variable `GEMINI_API_KEY`. | |
| prompt : str, optional | |
| Custom prompt to use when requesting the API key from the user. | |
| Returns | |
| ------- | |
| ChatGoogleGenerativeAI | |
| An instance of LangChain's ChatGoogleGenerativeAI model. | |
| Raises | |
| ------ | |
| ValueError | |
| If the model name is not in the list of supported models. | |
| Exception | |
| If there is an error loading the model or if the API key is invalid. | |
| Notes | |
| ----- | |
| The function will: | |
| 1. Check for the API key in the environment variables | |
| 2. Prompt for the key if not found | |
| 3. Validate the model name against supported models | |
| 4. Attempt to load the model | |
| 5. Handle any authentication errors by prompting for a new key | |
| """ | |
| if api_key is None: | |
| api_key = os.getenv("GEMINI_API_KEY") | |
| if not api_key: | |
| logger.info("Google API key not found in environment variables.") | |
| api_key = getpass("Please enter your Google API key: ") | |
| os.environ["GEMINI_API_KEY"] = api_key | |
| if model_name not in supported_gemini_models: | |
| raise ValueError( | |
| f"Unsupported model '{model_name}'. Supported models are: {supported_gemini_models}." | |
| ) | |
| try: | |
| logger.info(f"Loading Gemini model: {model_name}") | |
| llm = ChatGoogleGenerativeAI( | |
| model=model_name, | |
| temperature=temperature, | |
| api_key=api_key, | |
| max_output_tokens=6000, | |
| ) | |
| # No guarantee that api_key is valid, authentication happens only during invocation | |
| logger.info(f"Requested model: {model_name}") | |
| logger.info("Gemini model loaded successfully") | |
| return llm | |
| except Exception as e: | |
| # Can remove this since authentication happens only during invocation | |
| if "AuthenticationError" in str(e) or "invalid_api_key" in str(e): | |
| logger.warning("Invalid Google API key.") | |
| api_key = getpass("Please enter a valid Google API key: ") | |
| os.environ["GEMINI_API_KEY"] = api_key | |
| # Retry with new API key | |
| return load_gemini_model(model_name, temperature, api_key, prompt) | |
| else: | |
| logger.error(f"Error loading Google model: {str(e)}") | |
| raise | |