kig_core/llm_interface.py

import os
import time
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_openai import ChatOpenAI
from langchain_core.language_models.chat_models import BaseChatModel
from .config import settings
import logging

logger = logging.getLogger(__name__)

# Store initialized models to avoid re-creating them repeatedly
_llm_cache = {}

def get_llm(model_name: str) -> BaseChatModel:
    """
    Returns an initialized LangChain chat model based on the provided name.
    Caches initialized models.
    """
    global _llm_cache
    if model_name in _llm_cache:
        return _llm_cache[model_name]

    logger.info(f"Initializing LLM: {model_name}")

    if model_name.startswith("gemini"):
        if not settings.gemini_api_key:
            raise ValueError("GEMINI_API_KEY is not configured.")
        try:
            # Use GOOGLE_API_KEY environment variable set in config.py
            llm = ChatGoogleGenerativeAI(model=model_name)
            _llm_cache[model_name] = llm
            logger.info(f"Initialized Google Generative AI model: {model_name}")
            return llm
        except Exception as e:
            logger.error(f"Failed to initialize Gemini model '{model_name}': {e}", exc_info=True)
            raise RuntimeError(f"Could not initialize Gemini model: {e}") from e

    elif model_name.startswith("gpt"):
        if not settings.openai_api_key:
            raise ValueError("OPENAI_API_KEY is not configured.")
        try:
            # Base URL can be added here if using a proxy
            # base_url = "https://your-proxy-if-needed/"
            llm = ChatOpenAI(model=model_name, api_key=settings.openai_api_key) # Base URL optional
            _llm_cache[model_name] = llm
            logger.info(f"Initialized OpenAI model: {model_name}")
            return llm
        except Exception as e:
            logger.error(f"Failed to initialize OpenAI model '{model_name}': {e}", exc_info=True)
            raise RuntimeError(f"Could not initialize OpenAI model: {e}") from e

    # Add other model providers (Anthropic, Groq, etc.) here if needed

    else:
        logger.error(f"Unsupported model provider for model name: {model_name}")
        raise ValueError(f"Model '{model_name}' is not supported or configuration is missing.")

def invoke_llm(var,parameters):
    try:
        return var.invoke(parameters)
    except Exception as e:
        print(f"Error during .invoke : {e} \nwaiting 60 seconds")
        time.sleep(60)
        print("Waiting is finished")
        return var.invoke(parameters)
        
# Example usage (could be called from other modules)
# main_llm = get_llm(settings.main_llm_model)
# eval_llm = get_llm(settings.eval_llm_model)