backend/core/llm.py

"""backend/core/llm.py — LLM factory supporting Gemini and Grok."""
from __future__ import annotations


def get_llm(role: str, temperature: float = 0.1):
    """
    Return the right LLM based on LLM_PROVIDER in .env.
    role: "planner" | "executor" | "critic" | "memory"
    """
    from .config import get_settings
    settings = get_settings()

    max_tokens = settings.max_output_tokens

    if settings.llm_provider == "groq":
        from langchain_groq import ChatGroq

        model_map = {
            "planner": settings.groq_planner_model,
            "executor": settings.groq_executor_model,
            "critic": settings.groq_critic_model,
            "memory": settings.groq_memory_model,
        }
        return ChatGroq(
            model=model_map.get(role, "llama-3.3-70b-versatile"),
            temperature=temperature,
            api_key=settings.groq_api_key,
            max_tokens=max_tokens,
        )

    else:  # gemini (default)
        from langchain_google_genai import ChatGoogleGenerativeAI

        model_map = {
            "planner": settings.planner_model,
            "executor": settings.executor_model,
            "critic": settings.critic_model,
            "memory": settings.memory_model,
        }
        return ChatGoogleGenerativeAI(
            model=model_map.get(role, settings.planner_model),
            temperature=temperature,
            google_api_key=settings.google_api_key,
            max_output_tokens=max_tokens,
        )