src/chemgraph/models/openai.py

"""Load OpenAI models using LangChain."""

import os
from getpass import getpass
from langchain_openai import ChatOpenAI
from chemgraph.models.supported_models import (
    ARGO_DEFAULT_BASE_URL,
    supported_openai_models,
    supported_argo_models,
)
from chemgraph.utils.config_utils import normalize_openai_base_url
from chemgraph.utils.logging_config import setup_logger

logger = setup_logger(__name__)

# Maps user-facing ``argo:`` model names to the internal wire names
# expected by the Argo API (https://apps.inside.anl.gov/argoapi).
# When a different endpoint (e.g. ArgoProxy) is used, the ``argo:``
# prefix is stripped instead and the remainder is sent as-is.
ARGO_MODEL_MAP = {
    # GPT family
    "argo:gpt-3.5-turbo": "gpt35",
    "argo:gpt-3.5-turbo-16k": "gpt35turbo16k",
    "argo:gpt-4": "gpt4",
    "argo:gpt-4-32k": "gpt432k",
    "argo:gpt-4-turbo": "gpt4turbo",
    "argo:gpt-4o": "gpt4o",
    "argo:gpt-4o-latest": "gpt4olatest",
    "argo:gpt-4o-mini": "gpt4omini",
    "argo:gpt-4.1": "gpt41",
    "argo:gpt-4.1-mini": "gpt41mini",
    "argo:gpt-4.1-nano": "gpt41nano",
    "argo:gpt-5": "gpt5",
    "argo:gpt-5-mini": "gpt5mini",
    "argo:gpt-5-nano": "gpt5nano",
    "argo:gpt-5.1": "gpt51",
    "argo:gpt-5.2": "gpt52",
    "argo:gpt-5.4": "gpt54",

    # Reasoning / o-series
    "argo:o1-preview": "gpto1preview",
    "argo:o1-mini": "gpto1mini",
    "argo:o1": "gpto1",
    "argo:o3-mini": "gpto3mini",
    "argo:o3": "gpto3",
    "argo:o4-mini": "gpto4mini",
    # Gemini via Argo
    "argo:gemini-2.5-pro": "gemini25pro",
    "argo:gemini-2.5-flash": "gemini25flash",
    # Claude via Argo
    "argo:claude-opus-4.6": "claudeopus46",
    "argo:claude-opus-4.5": "claudeopus45",
    "argo:claude-opus-4.1": "claudeopus41",
    "argo:claude-opus-4": "claudeopus4",
    "argo:claude-haiku-4.5": "claudehaiku45",
    "argo:claude-sonnet-4.5": "claudesonnet45",
    "argo:claude-sonnet-4": "claudesonnet4",
    "argo:claude-sonnet-3.5-v2": "claudesonnet35v2",
    "argo:claude-haiku-3.5": "claudehaiku35",
}


def _normalize_argo_model(model_name: str, base_url: str) -> str:
    """Normalize an ``argo:``-prefixed model name for the target endpoint.

    * Argo API (base_url contains ``argoapi``): map to internal wire
      names via ``ARGO_MODEL_MAP`` (e.g. ``argo:gpt-4o`` -> ``gpt4o``).
    * Other endpoints (ArgoProxy, custom): strip the ``argo:`` prefix
      and send the remainder as-is (e.g. ``argo:gpt-4o`` -> ``gpt-4o``).

    Parameters
    ----------
    model_name : str
        Requested model identifier.
    base_url : str
        Endpoint base URL used to choose normalization behavior.

    Returns
    -------
    str
        Endpoint-specific model name.
    """
    if not model_name.startswith("argo:"):
        return model_name

    if base_url and "argoapi" in base_url:
        # Argo API endpoint -- use the wire-name map
        normalized = ARGO_MODEL_MAP.get(model_name)
        if normalized:
            logger.info("Normalized Argo model '%s' -> '%s'", model_name, normalized)
            return normalized
        # Fallback: strip prefix and remove punctuation
        fallback = model_name.removeprefix("argo:").replace("-", "").replace(".", "")
        logger.info(
            "Normalized Argo model '%s' -> '%s' (fallback)", model_name, fallback
        )
        return fallback
    else:
        # Non-Argo-API endpoint -- strip prefix only
        stripped = model_name.removeprefix("argo:")
        logger.info("Stripped argo: prefix '%s' -> '%s'", model_name, stripped)
        return stripped


def load_openai_model(
    model_name: str,
    temperature: float,
    api_key: str = None,
    prompt: str = None,
    base_url: str = None,
    argo_user: str = None,
) -> ChatOpenAI:
    """Load an OpenAI chat model into LangChain.

    This function loads an OpenAI 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 OpenAI chat model to load. See supported_openai_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 OpenAI API key. If not provided, the function will attempt to retrieve it
        from the environment variable `OPENAI_API_KEY`.
    prompt : str, optional
        Custom prompt to use when requesting the API key from the user.

    Returns
    -------
    ChatOpenAI
        An instance of LangChain's ChatOpenAI 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
    """

    base_url = normalize_openai_base_url(base_url)

    # Apply default Argo base URL for argo: models when none is specified.
    if model_name.startswith("argo:") and not base_url:
        base_url = ARGO_DEFAULT_BASE_URL
        logger.info("Using default Argo base URL: %s", base_url)

    if api_key is None:
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            if model_name.startswith("argo:"):
                # Argo API authenticates via the 'user' field, not an API key.
                # Use argo_user as a placeholder since ChatOpenAI requires a value.
                api_key = argo_user or os.getenv("ARGO_USER", "chemgraph")
            else:
                logger.warning("OPENAI_API_KEY not found in environment variables.")
                print("OPENAI_API_KEY not set. Please enter your OpenAI API key.")
                api_key = getpass("OpenAI API key: ")
                os.environ["OPENAI_API_KEY"] = api_key

    if model_name not in supported_openai_models and model_name not in supported_argo_models:
        raise ValueError(
            f"Unsupported model '{model_name}'. Supported models are: {supported_openai_models}."
        )

    is_argo_endpoint = bool(base_url and "argoapi" in base_url)
    argo_user = (
        argo_user or os.getenv("ARGO_USER", "chemgraph")
        if is_argo_endpoint
        else None
    )

    try:
        if base_url is not None:
            logger.info(f"Using custom base URL: {base_url}")
            model_name = _normalize_argo_model(model_name, base_url)
            llm_kwargs = dict(
                model=model_name,
                temperature=temperature,
                api_key=api_key,
                base_url=base_url,
                max_tokens=4000,
                top_p=1.0,
                frequency_penalty=0.0,
                presence_penalty=0.0,
            )
            # Argo gateways may require an explicit "user" field in payload.
            if is_argo_endpoint and argo_user:
                llm_kwargs["model_kwargs"] = {"user": argo_user}
                logger.info(
                    "Using Argo user from config/ARGO_USER/default: %s", argo_user
                )
            llm = ChatOpenAI(**llm_kwargs)
        else:
            logger.info(f"Loading OpenAI model: {model_name}")
            llm = ChatOpenAI(
                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("OpenAI 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 OpenAI API key.")
            print("The provided OpenAI API key is invalid. Please enter a valid key.")
            api_key = getpass("OpenAI API key: ")
            os.environ["OPENAI_API_KEY"] = api_key
            # Retry with new API key
            return load_openai_model(
                model_name, temperature, api_key, prompt, base_url, argo_user
            )
        else:
            logger.error(f"Error loading OpenAI model: {str(e)}")
            raise