"""
LLM integration module for LangChain with support for OpenAI and Ollama.
Provides a reusable LLM client based on environment configuration.
"""

import os
from typing import Union

from langchain.llms import OpenAI, Ollama
from langchain.chat_models import ChatOpenAI, ChatOllama

# Environment variable to select provider: "openai" or "ollama"
LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai").lower()


def get_llm() -> Union[OpenAI, Ollama, ChatOpenAI, ChatOllama]:
    """
    Returns an LLM instance based on the configured provider.

    For OpenAI, uses the default OpenAI LLM (text-davinci-003 or gpt-3.5-turbo).
    For Ollama, uses the default Ollama LLM (e.g., llama2).

    Raises:
        ValueError: If an unsupported provider is specified.
    """
    if LLM_PROVIDER == "openai":
        # Use ChatOpenAI for GPT-3.5-turbo by default
        return ChatOpenAI(temperature=0.7)
    elif LLM_PROVIDER == "ollama":
        # Use ChatOllama for local models
        return ChatOllama(model="llama2", temperature=0.7)
    else:
        raise ValueError(f"Unsupported LLM provider: {LLM_PROVIDER}")