Spaces:
Running
Running
| """Load GROQ models using LangChain.""" | |
| import os | |
| from getpass import getpass | |
| from langchain_groq import ChatGroq | |
| from chemgraph.utils.logging_config import setup_logger | |
| logger = setup_logger(__name__) | |
| def load_groq_model( | |
| model_name: str, | |
| temperature: float, | |
| api_key: str = None, | |
| prompt: str = None, | |
| base_url: str = None, | |
| ) -> ChatGroq: | |
| """Load a GROQ chat model into LangChain. | |
| This function loads a GROQ 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 GROQ chat model to load. See supported_groq_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 GROQ API key. If not provided, the function will attempt to retrieve it | |
| from the environment variable `GROQ_API_KEY`. | |
| prompt : str, optional | |
| Custom prompt to use when requesting the API key from the user. | |
| base_url : str, optional | |
| Custom base URL for the GROQ API (currently unused but included for consistency). | |
| Returns | |
| ------- | |
| ChatGroq | |
| An instance of LangChain's ChatGroq 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 | |
| """ | |
| # Strip the "groq:" routing prefix before sending to the API. | |
| if model_name.startswith("groq:"): | |
| model_name = model_name.removeprefix("groq:") | |
| if api_key is None: | |
| api_key = os.getenv("GROQ_API_KEY") | |
| if not api_key: | |
| logger.info("GROQ API key not found in environment variables.") | |
| api_key = getpass("Please enter your GROQ API key: ") | |
| os.environ["GROQ_API_KEY"] = api_key | |
| try: | |
| logger.info(f"Loading GROQ model: {model_name}") | |
| llm = ChatGroq( | |
| model=model_name, | |
| temperature=temperature, | |
| api_key=api_key, | |
| max_tokens=6000, | |
| ) | |
| # No guarantee that api_key is valid, authentication happens only during invocation | |
| logger.info(f"Requested model: {model_name}") | |
| logger.info("GROQ 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 GROQ API key.") | |
| api_key = getpass("Please enter a valid GROQ API key: ") | |
| os.environ["GROQ_API_KEY"] = api_key | |
| # Retry with new API key | |
| return load_groq_model(model_name, temperature, api_key, prompt) | |
| else: | |
| logger.error(f"Error loading GROQ model: {str(e)}") | |
| raise | |