utils/model_loader.py

import os
from dotenv import load_dotenv
from typing import Optional, Any
from pydantic import BaseModel, Field
from utils.config_loader import load_config
from langchain_groq import ChatGroq
from langchain_openai import ChatOpenAI

load_dotenv()


class ConfigLoader:
    def __init__(self):
        self.config = load_config()

    def __getitem__(self, key):
        return self.config[key]

    def get(self, key, default=None):
        return self.config.get(key, default)


class ModelLoader(BaseModel):
    # Accepts "ollama", "openai", "groq" (or any openai-compatible provider name)
    model_provider: str = "ollama"
    config: Optional[ConfigLoader] = Field(default=None, exclude=True)

    def model_post_init(self, __context: Any) -> None:
        self.config = ConfigLoader()

    class Config:
        arbitrary_types_allowed = True

    def load_llm(self):
        """
        Load and return the LLM based on MODEL_PROVIDER env var or the
        model_provider argument passed at construction time.

        Priority for each setting:
          1. Environment variable (MODEL_NAME, MODEL_BASE_URL, MODEL_API_KEY …)
          2. config/config.yaml values
        """
        load_dotenv()
        provider = (os.getenv("MODEL_PROVIDER") or self.model_provider).lower()
        print(f"[ModelLoader] Loading LLM — provider: {provider}")

        if provider == "groq":
            api_key = os.getenv("GROQ_API_KEY")
            if not api_key:
                raise EnvironmentError(
                    "GROQ_API_KEY not found in environment variables. "
                    "Set it in .env or as a Hugging Face Space secret."
                )
            model_name = (
                os.getenv("MODEL_NAME")
                or self.config["llm"].get("groq", {}).get("model_name", "llama-3.1-8b-instant")
            )
            return ChatGroq(model=model_name, api_key=api_key)

        elif provider == "openai":
            api_key = os.getenv("OPENAI_API_KEY")
            if not api_key:
                raise EnvironmentError(
                    "OPENAI_API_KEY not found in environment variables."
                )
            model_name = os.getenv("MODEL_NAME") or "gpt-4o-mini"
            return ChatOpenAI(model=model_name, api_key=api_key)

        else:
            # Default: "ollama" or any OpenAI-compatible endpoint
            api_key = os.getenv("OLLAMA_API_KEY") or os.getenv("MODEL_API_KEY", "ollama")
            model_name = (
                os.getenv("MODEL_NAME")
                or self.config["llm"]["ollama"]["model_name"]
            )
            base_url = (
                os.getenv("MODEL_BASE_URL")
                or self.config["llm"]["ollama"]["base_url"]
            )
            return ChatOpenAI(model=model_name, base_url=base_url, api_key=api_key)