File size: 2,712 Bytes
1bcef92 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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) |