"""Core functionality for LLMs"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/00_llm.base.ipynb.

# %% auto 0
__all__ = ['MIN_MODEL_SIZE', 'OLLAMA_URL', 'MISTRAL_MODEL_URL', 'MISTRAL_MODEL_ID', 'MISTRAL_PROMPT_TEMPLATE', 'ZEPHYR_MODEL_URL',
           'ZEPHYR_MODEL_ID', 'ZEPHYR_PROMPT_TEMPLATE', 'LLAMA_MODEL_URL', 'LLAMA_MODEL_ID', 'LLAMA_PROMPT_TEMPLATE',
           'MODEL_URL_DICT', 'URL2NAME', 'MODEL_ID_DICT', 'LLAMA_CPP', 'TRANSFORMERS', 'ENGINE_DICT', 'PROMPT_DICT',
           'DEFAULT_MODEL', 'DEFAULT_ENGINE', 'DEFAULT_EMBEDDING_MODEL', 'DEFAULT_QA_PROMPT', 'LLM']

# %% ../../nbs/00_llm.base.ipynb 3
from ..utils import get_datadir, get_models_dir, download, format_string, DEFAULT_DB
from . import helpers
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from langchain.output_parsers import OutputFixingParser
from langchain_community.llms import LlamaCpp
from langchain_openai import ChatOpenAI
from langchain_litellm import ChatLiteLLM
from .backends import ChatGovCloudBedrock
from langchain_core.messages.ai import AIMessage
from langchain_core.documents import Document
import os
import warnings
from typing import Any, Dict, Optional, Callable, Union, List

# %% ../../nbs/00_llm.base.ipynb 4
MIN_MODEL_SIZE = 250000000
OLLAMA_URL = 'http://localhost:11434/v1'
MISTRAL_MODEL_URL = "https://huggingface.co/TheBloke/Mistral-7B-Instruct-v0.2-GGUF/resolve/main/mistral-7b-instruct-v0.2.Q4_K_M.gguf"
MISTRAL_MODEL_ID = "TheBloke/Mistral-7B-Instruct-v0.2-AWQ"
MISTRAL_PROMPT_TEMPLATE = "[INST] {prompt} [/INST]"
ZEPHYR_MODEL_URL = "https://huggingface.co/TheBloke/zephyr-7B-beta-GGUF/resolve/main/zephyr-7b-beta.Q4_K_M.gguf"
ZEPHYR_MODEL_ID = "TheBloke/zephyr-7B-beta-AWQ"
ZEPHYR_PROMPT_TEMPLATE = "<|system|>\n</s>\n<|user|>\n{prompt}</s>\n<|assistant|>"
LLAMA_MODEL_URL = "https://huggingface.co/lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/resolve/main/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf"
LLAMA_MODEL_ID = "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4"
LLAMA_PROMPT_TEMPLATE = """<|start_header_id|>system<|end_header_id|>

You are a super-intelligent helpful assistant that executes instructions.<|eot_id|><|start_header_id|>user<|end_header_id|>

{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
"""
MODEL_URL_DICT = {'mistral' : MISTRAL_MODEL_URL, 'zephyr':ZEPHYR_MODEL_URL, 'llama' : LLAMA_MODEL_URL}
URL2NAME = dict([(v,k) for k,v in MODEL_URL_DICT.items()])
MODEL_ID_DICT = {'mistral': MISTRAL_MODEL_ID, 'zephyr' : ZEPHYR_MODEL_ID, 'llama': LLAMA_MODEL_ID}
LLAMA_CPP = 'llama.cpp'
TRANSFORMERS = 'transformers'
ENGINE_DICT = {LLAMA_CPP: MODEL_URL_DICT, TRANSFORMERS : MODEL_ID_DICT}
PROMPT_DICT = {'mistral': MISTRAL_PROMPT_TEMPLATE, 'zephyr' : ZEPHYR_PROMPT_TEMPLATE, 'llama' : LLAMA_PROMPT_TEMPLATE} 
DEFAULT_MODEL = 'zephyr'
DEFAULT_ENGINE = LLAMA_CPP
DEFAULT_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_QA_PROMPT = """"Use the following pieces of context delimited by three backticks to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.

```{context}```

Question: {question}
Helpful Answer:"""


class LLM:
    def __init__(
        self,
        model_url:Optional[str] = None,
        model_id:Optional[str] = None,
        default_model:str  = DEFAULT_MODEL,
        default_engine:str = DEFAULT_ENGINE,
        prompt_template: Optional[str] = None,
        model_download_path: Optional[str] = None,
        vectordb_path: Optional[str] = None,
        store_type:str='dense',
        vectorstore=None,
        max_tokens: int = 512,
        n_gpu_layers: Optional[int] = None,
        n_ctx: int = 3900,
        n_batch: int = 1024,
        stop:list=[],
        mute_stream: bool = False,
        callbacks=[],
        embedding_model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        embedding_model_kwargs: Optional[dict] = None,
        embedding_encode_kwargs: dict = {"normalize_embeddings": False},
        rag_num_source_docs: int = 4,
        rag_score_threshold: float = 0.0,
        check_model_download:bool=True,
        confirm: bool = True,
        verbose: bool = True,
        **kwargs,
    ):
        """
        LLM Constructor.  Extra `kwargs` (e.g., temperature) are fed directly the LangChain LLM or Transformer Pipeline.

        **Args:**

        - *model_url*: URL to `.GGUF` model (or the filename if already been downloaded to `model_download_path`).
                       To use an OpenAI-compatible REST API (e.g., vLLM, OpenLLM, Ollama), supply the URL (e.g., `http://localhost:8080/v1`).
                       To use a cloud-based provider, replace URL with provider/model: `openai/<name_of_model>` (e.g., `openai/gpt-3.5-turbo`).
                       Any LiteLLM-supported provider is supported.
                       To use Ollama, use `ollama/<model_name` (e.g., `ollama/llama3.2`).
                       If None, use the model indicated by `default_model`.
        - *model_id*: Name of or path to Hugging Face model (e.g., in SafeTensor format). Hugging Face Transformers is used for LLM generation instead of **llama-cpp-python**. Mutually-exclusive with `model_url` and `default_model`. The `n_gpu_layers` and `model_download_path` parameters are ignored if `model_id` is supplied.
        - *default_model*: One of {'mistral', 'zephyr', 'llama'}, where mistral is Mistral-Instruct-7B-v0.2, zephyr is Zephyr-7B-beta, and llama is Llama-3.1-8B.
        - *default_engine*: The engine used to run the `default_model`. One of {'llama.cpp', 'transformers'}.
        - *prompt_template*: Optional prompt template (must have a variable named "prompt"). Prompt templates are not typically needed when using the `model_id` parameter, as transformers sets it automatically.
        - *model_download_path*: Path to download model. Default is `onprem_data` in user's home directory.
        - *vectordb_path*: Path to vector database (created if it doesn't exist).
                           Default is `onprem_data/vectordb` in user's home directory.
        - *store_type*: One of `dense` for the default dense vector database (i.e., chroma) or
                        `sparse` for the sparse vector store (i.e., a keyword search engine).  
                        (Documents stored in sparse vector databases are converted to dense vectors at inference time 
                        when used with `LLM.ask`.)
        - *vectorstore*: an onprem.ingest.stores.base.VectorStore instance.
        - *max_tokens*: The maximum number of tokens to generate.
        - *n_gpu_layers*: Number of layers to be loaded into gpu memory. Default is `None`. Only used for llama-cpp backend.
        - *n_ctx*: Token context window. Only used for llama-cpp backend.
                   For Ollama backend, explicitly supply `num_ctx` instead which is passed to LiteLLM.
                   Hugging Face Transformers backend (i.e., when using the model_id parameter)
                   sets context window automatically.
        - *n_batch*: Number of tokens to process in parallel. Only used for llama-cpp backend.
        - *stop*: a list of strings to stop generation when encountered (applied to all calls to `LLM.prompt`)
        - *mute_stream*: Mute ChatGPT-like token stream output during generation
        - *callbacks*: Callbacks to supply model
        - *embedding_model_name*: name of sentence-transformers model. Used for `LLM.ingest` and `LLM.ask`.
        - *embedding_model_kwargs*: arguments to embedding model (e.g., `{device':'cpu'}`). If None, uses GPU if available.
        - *embedding_encode_kwargs*: arguments to encode method of
                                     embedding model (e.g., `{'normalize_embeddings': False}`).
        - *rag_num_source_docs*: The maximum number of documents retrieved and fed to `LLM.ask` and `LLM.chat` to generate answers.
        - *rag_score_threshold*: Minimum similarity score for source to be considered by `LLM.ask` and `LLM.chat`.
        - *confirm*: whether or not to confirm with user before downloading a model
        - *verbose*: Verbosity
        """
        self.model_id = None
        self.model_url = None

        if model_url in URL2NAME:
            default_model = URL2NAME[model_url]
            model_url = None

        if model_url and model_id:
            raise ValueError('The parameters model_url and model_id are mutually-exclusive.')
        elif model_id:
            self.model_id = model_id
            self.model_name = os.path.basename(model_id)
        elif model_url:
            self.model_url = model_url.split("?")[0]
            self.model_name = os.path.basename(model_url)

            # handle Ollama-style URL
            provider_model = self.process_service(self.model_url)
            if provider_model and provider_model.split("/")[0] == 'ollama':
                # point ollama to ollama_chat as convenience
                self.model_url = self.model_url.replace('ollama', 'ollama_chat')

            # override model name if explicitly provided
            # for local APIs, this replaces "v1" with actual model name
            self.model_name = kwargs.get('model', self.model_name)
        else: # neither supplied so use defaults
            url_or_id = ENGINE_DICT[default_engine][default_model]
            self.model_name = os.path.basename(url_or_id)
            if default_engine == LLAMA_CPP:
                self.model_url = url_or_id
                self.model_id = None
                prompt_template = PROMPT_DICT[default_model] if not prompt_template else prompt_template
            else:
                self.model_url = None
                self.model_id = url_or_id
        self.model_download_path = model_download_path or get_models_dir()

        if self.is_llamacpp():
            try:
                from llama_cpp import Llama
            except ImportError:
                raise ValueError('To run local LLMs, the llama-cpp-python package is required. ' +\
                                 'You can visit https://python.langchain.com/docs/integrations/llms/llamacpp ' +\
                                 'and follow the instructions for your operating system.')
        if self.is_llamacpp() and not os.path.isfile(os.path.join(self.model_download_path, self.model_name)):
            self.download_model(
                self.model_url,
                model_download_path=self.model_download_path,
                confirm=confirm,
            )
        self.prompt_template = prompt_template
        self.vectordb_path = vectordb_path or os.path.join(get_datadir(), DEFAULT_DB)
        self.store_type = store_type  # one of 'dense', 'sparse', or 'dual'
        self.vectorstore = vectorstore   # set automatically if missing
        self.llm = None
        self.qa = None
        self.chatbot = None
        self.n_gpu_layers = n_gpu_layers
        self.max_tokens = max_tokens
        self.n_ctx = n_ctx
        self.n_batch = n_batch
        self.stop = stop
        self.mute_stream = mute_stream
        self.callbacks = [] if mute_stream else [StreamingStdOutCallbackHandler()]
        if callbacks:
            self.callbacks.extend(callbacks)
        self.embedding_model_name = embedding_model_name
        self.embedding_model_kwargs = embedding_model_kwargs
        self.embedding_encode_kwargs = embedding_encode_kwargs
        self.rag_num_source_docs = rag_num_source_docs
        self.rag_score_threshold = rag_score_threshold
        self.check_model_download = check_model_download
        self.verbose = verbose
        self.extra_kwargs = kwargs
        self._rag_pipeline = None


        # explicitly set offload_kqv
        # reference: https://github.com/abetlen/llama-cpp-python/issues/999#issuecomment-1858041458
        # commented out now: no longer needed as of llama-cpp-python==0.2.75
        #self.offload_kqv = True if n_gpu_layers is not None and n_gpu_layers > 0 else False

        # load LLM
        self.load_llm()

    def set_store_type(self, store_type:str):
        """
        Change store type. Not that this will reload the default vectorstore on selected store type.
        """
        if store_type not in ['dense', 'sparse', 'dual']:
            raise ValueError('store_type must be one of {"dense", "sparse", "dual"}')
        self.store_type = store_type
        self.vectorstore = None
        self.load_vectorstore()
        return

    def get_store_type(self):
        return self.store_type

    def is_sparse_store(self):
        return self.store_type == 'sparse'

    def is_dense_store(self):
        return self.store_type == 'dense'

    def is_dual_store(self):
        return self.store_type == 'dual'
        
    def is_local_api(self):
        basename = os.path.basename(self.model_url) if self.model_url else None
        return self.model_url and self.model_url.lower().startswith('http') and not basename.lower().endswith('.gguf') and not basename.lower().endswith('.bin')


    def process_service(self, model_url):
        """
        If model_url represents an LLM service (e.g, openai), return it
        in LiteLLM-style syntax.
        """
        if model_url:
            if (":" not in model_url or model_url.startswith('ollama')) and model_url.count("/") == 1:
                # conclude that this a standard LiteLLM path (e.g., openai/gpt-4o)
                # or Ollama model which contain ":"
                return model_url
            
            # Special handling for govcloud-bedrock that supports both :// and / separators
            if model_url.startswith('govcloud-bedrock'):
                if '://' in model_url:
                    return model_url.split('://', 1)[1]
                elif '/' in model_url:
                    # For govcloud-bedrock/model-id format, get everything after first /
                    return model_url.split('/', 1)[1]
                else:
                    return os.path.basename(model_url)
            
            provider = model_url.split(":")[0]
            if provider.startswith('http'):
                return None
            
            model = os.path.basename(self.model_url)
            return model if provider == model else f'{provider}/{model}'
        else:
            return None


    def is_local(self):
        return (self.model_url and self.model_url.lower().endswith('.gguf')) or self.model_id


    def is_llamacpp(self):
        return self.is_local() and not self.is_hf()

    def is_hf(self):
        return self.model_id is not None

    def is_openai_model(self):
        return self.model_url and self.model_url.lower().startswith('openai')

    def is_govcloud_bedrock(self):
        return self.model_url and self.model_url.lower().startswith('govcloud-bedrock')

    def update_max_tokens(self, value:int=512):
        """
        Update `max_tokens` (maximum length of generation).
        """
        llm = self.load_llm()
        llm.max_tokens = value


    def update_stop(self, value:list=[]):
        """
        Update `max_tokens` (maximum length of generation).
        """
        llm = self.load_llm()
        llm.stop = value


    @classmethod
    def download_model(
        cls,
        model_url:Optional[str] = None,
        default_model:str=DEFAULT_MODEL,
        model_download_path: Optional[str] = None,
        confirm: bool = True,
        ssl_verify: bool = True,
    ):
        """
        Download an LLM in GGML format supported by [lLama.cpp](https://github.com/ggerganov/llama.cpp).

        **Args:**

        - *model_url*: URL of model. If None, then use default_model.
        - *default_model*: One of {'mistral', 'zephyr', 'llama'}, where mistral is Mistral-Instruct-7B-v0.2, zephyr is Zephyr-7B-beta, and llama is Llama-3.1-8B.
        - *model_download_path*: Path to download model. Default is `onprem_data` in user's home directory.
        - *confirm*: whether or not to confirm with user before downloading
        - *ssl_verify*: If True, SSL certificates are verified.
                        You can set to False if corporate firewall gives you problems.
        """
        model_url = MODEL_URL_DICT[default_model] if not model_url else model_url
        if 'https://huggingface.co' in model_url and 'resolve' not in model_url:
            warnings.warn('\n\nThe supplied URL may not be pointing to the actual GGUF model file.  Please check it.\n\n')
        datadir = model_download_path or get_models_dir()
        model_name = os.path.basename(model_url)
        filename = os.path.join(datadir, model_name)
        confirm_msg = f"\nYou are about to download the LLM {model_name} to the {datadir} folder. Are you sure?"
        if os.path.isfile(filename):
            confirm_msg = f"There is already a file {model_name} in {datadir}.\n Do you want to still download it?"

        shall = True
        if confirm:
            shall = input("%s (y/N) " % confirm_msg).lower() == "y"
        if shall:
            download(model_url, filename, verify=ssl_verify)
        else:
            warnings.warn(
                f'{model_name} was not downloaded because "Y" was not selected.'
            )
        return


    def load_vectorstore(self, custom_vectorstore=None, reset=False):
        """
        Get `VectorStore` instance.
        Use the vectorstore's methods directly instead of accessing the underlying database.
        Supply `custom_vectorstore` to use your own `VectorStore` instance
        (i.e., subclass `DenseStore` or `SparseStore`).
        Supply `reset=True` to reload the default vectorstore.
        """
        from onprem.ingest.stores import DenseStore, SparseStore, DualStore

        if reset:
            self.vectorstore = None
        if custom_vectorstore:
            self.vectorstore = custom_vectorstore

        if not self.vectorstore:

            # store vector stores within subfolders under `vectordb_path`
            if self.is_sparse_store():
                store_path = os.path.join(self.vectordb_path, 'sparse')
                # create vector store
                self.vectorstore = SparseStore.create(
                    embedding_model_name=self.embedding_model_name,
                    embedding_model_kwargs=self.embedding_model_kwargs,
                    embedding_encode_kwargs=self.embedding_encode_kwargs,
                    persist_location=store_path,
                )
            elif self.is_dense_store():
                store_path = os.path.join(self.vectordb_path, 'dense')

                # create vector store
                self.vectorstore = DenseStore.create(
                    embedding_model_name=self.embedding_model_name,
                    embedding_model_kwargs=self.embedding_model_kwargs,
                    embedding_encode_kwargs=self.embedding_encode_kwargs,
                    persist_location=store_path,
                )
            else:

                # Create paths for both stores
                dense_path = os.path.join(self.vectordb_path, 'dense')
                sparse_path = os.path.join(self.vectordb_path, 'sparse')

                # create vector store
                self.vectorstore = DualStore(
                    embedding_model_name=self.embedding_model_name,
                    embedding_model_kwargs=self.embedding_model_kwargs,
                    embedding_encode_kwargs=self.embedding_encode_kwargs,
                    sparse_persist_location=sparse_path,
                    dense_persist_location=dense_path,
                )
        return self.vectorstore

    def ingest(
        self,
        source_directory: str, # path to folder containing documents
        chunk_size: Optional[int] = None, # text is split to this many characters by `langchain.text_splitter.RecursiveCharacterTextSplitter`. If None, uses onprem.ingest.base.DEFAULT_CHUNK_SIZE.
        chunk_overlap: Optional[int] = None, # character overlap between chunks in `langchain.text_splitter.RecursiveCharacterTextSplitter`. If None, uses onprem.ingest.base.DEFAULT_CHUNK_OVERLAP.
        ignore_fn:Optional[Callable] = None, # callable that accepts the file path and returns True for ignored files
        batch_size:int=1000, # batch size used when processing documents(e.g, creating embeddings).
        **kwargs, # Extra kwargs fed to downstream functions, `load_single_document` and/or `load_documents`
    ):
        """
        Ingests all documents in `source_folder` into vector database.
        Previously-ingested documents are ignored.
        Extra kwargs fed to `load_single_document`, `load_documents`, and/or `chunk_documents`.
        """
        # Import defaults from ingest module
        from ..ingest.base import DEFAULT_CHUNK_SIZE, DEFAULT_CHUNK_OVERLAP
        
        # Use defaults if None provided
        chunk_size = chunk_size if chunk_size is not None else DEFAULT_CHUNK_SIZE
        chunk_overlap = chunk_overlap if chunk_overlap is not None else DEFAULT_CHUNK_OVERLAP
        
        vectorstore = self.load_vectorstore()
        return vectorstore.ingest(
            source_directory,
            chunk_size=chunk_size, chunk_overlap=chunk_overlap, ignore_fn=ignore_fn,
            llm=kwargs['llm'] if 'llm' in kwargs else self,
            batch_size=batch_size,
            **kwargs
        )

    def check_model(self, silent:bool=False):
        """
        Returns the path to the model
        """
        if not self.is_llamacpp():
            return None
        datadir = self.model_download_path
        model_path = os.path.join(datadir, self.model_name)
        if not os.path.isfile(model_path):
            if silent: return None
            raise ValueError(
                f"The LLM model {self.model_name} does not appear to have been downloaded. "
                + "Execute the download_model() method to download it."
            )
        return model_path

    def _can_accept_parameter(self, llm_class, param_name):
        """
        Test if a parameter can be accepted directly by the LLM class
        without being moved to model_kwargs.
        """
        try:
            # Create a minimal test instance to see if the parameter is accepted
            test_kwargs = {param_name: 'test_value'}
            
            # Add minimal required parameters for common LLM types
            if hasattr(llm_class, 'model_fields') and 'openai_api_key' in llm_class.model_fields:
                test_kwargs['openai_api_key'] = 'dummy'
            
            # Suppress warnings during this test
            import warnings
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                # Try to instantiate - we don't care if it fails for other reasons,
                # only if it fails due to unexpected keyword argument
                llm_class(**test_kwargs)
            
            return True
            
        except TypeError as e:
            # If it's specifically about unexpected keyword argument, parameter not accepted
            if 'unexpected keyword argument' in str(e):
                return False
            # Other TypeErrors (missing required args, etc.) mean parameter is accepted
            return True
        except Exception:
            # Any other exception means the parameter was accepted but something else failed
            return True

    def _prepare_llm_kwargs(self, llm_class):
        """
        Inspect the LLM class signature and separate direct parameters from model_kwargs.
        Parameters that are in the class __init__ signature stay as direct parameters.
        Parameters not in the signature get moved to model_kwargs.
        """
        import inspect
        
        # Get the class constructor signature
        direct_param_names = set()
        
        # Check if it's a Pydantic model (most LangChain classes are)
        if hasattr(llm_class, 'model_fields'):
            # Pydantic v2
            direct_param_names = set(llm_class.model_fields.keys())
        elif hasattr(llm_class, '__fields__'):
            # Pydantic v1
            direct_param_names = set(llm_class.__fields__.keys())
        else:
            # Try traditional signature inspection
            try:
                sig = inspect.signature(llm_class.__init__)
                direct_param_names = set(sig.parameters.keys()) - {'self', 'args', 'kwargs'}
            except Exception as e:
                direct_param_names = set()
        
        # Start with a copy of extra_kwargs
        kwargs = self.extra_kwargs.copy()
        
        # If model_kwargs doesn't exist, create it
        if 'model_kwargs' not in kwargs:
            kwargs['model_kwargs'] = {}
        
        # Move parameters not in the class signature to model_kwargs
        # But first check if they can be accepted directly via compatibility layer
        for param_name in list(kwargs.keys()):
            if param_name != 'model_kwargs' and param_name not in direct_param_names:
                # Check if parameter can be accepted directly despite not being in signature
                # Only do compatibility check for local APIs
                if (self.is_local_api() or self.is_openai_model()) and\
                    self._can_accept_parameter(llm_class, param_name):
                    # Keep as direct parameter - compatibility layer can handle it
                    continue
                else:
                    # Move to model_kwargs
                    kwargs['model_kwargs'][param_name] = kwargs.pop(param_name)
        
        return kwargs

    def load_llm(self):
        """
        Loads the LLM from the model path.
        """


        if not self.llm and self.is_openai_model():
            kwargs = self._prepare_llm_kwargs(ChatOpenAI)
            self.llm = ChatOpenAI(model_name=self.model_name, 
                                  callbacks=self.callbacks, 
                                  streaming=not self.mute_stream,
                                  max_tokens=self.max_tokens,
                                  **kwargs)
        elif not self.llm and self.is_govcloud_bedrock():
            kwargs = self._prepare_llm_kwargs(ChatGovCloudBedrock)
            model_id = self.process_service(self.model_url)
            self.llm = ChatGovCloudBedrock(
                model_id=model_id,
                callbacks=self.callbacks,
                streaming=not self.mute_stream,
                max_tokens=self.max_tokens,
                **kwargs)
        elif not self.llm and\
                self.process_service(self.model_url) and\
                not self.check_model(silent=True):
            kwargs = self._prepare_llm_kwargs(ChatLiteLLM)
            self.llm = ChatLiteLLM(model=self.process_service(self.model_url),
                                  callbacks=self.callbacks,
                                  streaming=not self.mute_stream,
                                  max_tokens=self.max_tokens,
                                  **kwargs)
        elif not self.llm and self.is_local_api():
            kwargs = self._prepare_llm_kwargs(ChatOpenAI)
            self.llm = ChatOpenAI(base_url=self.model_url,
                                  #model_name=self.model_name, 
                                  callbacks=self.callbacks, 
                                  streaming=not self.mute_stream,
                                  max_tokens=self.max_tokens,
                                  **kwargs)
        elif not self.llm and self.is_hf():
            from transformers import pipeline, TextStreamer, AutoTokenizer
            from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline
            tokenizer = self.extra_kwargs['tokenizer'] if 'tokenizer' in self.extra_kwargs\
                        else AutoTokenizer.from_pretrained(self.model_id)

            if 'tokenizer' in self.extra_kwargs:
                del self.extra_kwargs['tokenizer']


            streamer = TextStreamer(tokenizer)


            pipe = pipeline('text-generation',
                              self.model_id,
                              tokenizer=tokenizer,
                              streamer=streamer if not self.mute_stream else None,
                              max_new_tokens = self.max_tokens,
                              return_full_text=False,
                              do_sample=True if\
                                     self.extra_kwargs.get('temperature', 0.8)>0.0 else False ,
                              **self.extra_kwargs)

            model = pipe.model
            if not model.generation_config.pad_token_id:
                tokenid = model.generation_config.eos_token_id
                model.generation_config.pad_token_id = tokenid[0] if isinstance(tokenid, list) else tokenid

            hfpipe = HuggingFacePipeline(pipeline=pipe)
            self.llm = ChatHuggingFace(llm=hfpipe, model_id=self.model_id, tokenizer=tokenizer)

            # Set generation_config.pad_token_id
            model = self.llm.llm.pipeline.model
            if not model.generation_config.pad_token_id:
                tokenid = model.generation_config.eos_token_id
                model.generation_config.pad_token_id = tokenid[0] if isinstance(tokenid, list) else tokenid

        elif not self.llm:
            model_path = self.check_model()
            if self.check_model_download and os.path.getsize(model_path) < MIN_MODEL_SIZE:
                raise ValueError(f'The model file ({model_path} is less than {MIN_MODEL_SIZE} bytes. ' +\
                                 'It may not have been fully downloaded. Please delete the file and start again. ')
            kwargs = self._prepare_llm_kwargs(LlamaCpp)
            self.llm = LlamaCpp(
                model_path=model_path,
                max_tokens=self.max_tokens,
                n_batch=self.n_batch,
                callbacks=self.callbacks,
                verbose=self.verbose,
                n_gpu_layers=self.n_gpu_layers,
                n_ctx=self.n_ctx,
                **kwargs)

        return self.llm


    def _format_image_prompt(self, prompt:str, image_path_or_url:str):
        """
        Correctly format image prompt
        """
        from langchain_core.messages import HumanMessage
        import base64
        if not image_path_or_url.startswith('http'):
            with open(image_path_or_url, "rb") as f:
                image_data = base64.b64encode(f.read()).decode('utf-8')
            image_path_or_url = f"data:image/jpeg;base64,{image_data}"

        message = HumanMessage(
            content=[
                {"type": "text", "text": prompt},
                {
                    "type": "image_url",
                    "image_url": {"url": image_path_or_url},
                },
            ],
        )
        prompt = [message]
        return prompt


    def _format_pydantic_prompt(self, prompt, pydantic_model):
        """
        Correctly format prompt for Pydantic model
        """
        parser = PydanticOutputParser(pydantic_object=pydantic_model)
        prompt_obj = PromptTemplate(
                     template="Answer the user query.\n{format_instructions}\n{prompt}\n",
                     input_variables=["prompt"],
                     partial_variables={"format_instructions": parser.get_format_instructions()},)
        return (prompt_obj.invoke({'prompt': prompt}).text, parser)


    def pydantic_prompt(self, 
                        prompt:str,
                        pydantic_model=None,
                        attempt_fix:bool= False,
                        fix_llm=None,
                        stop:list=[],
                        **kwargs):
        """
        Accept a prompt as string and Pydantic model describing the desired output.
        Output will be a Pydantic object in the requested format.


        **Args:**

        - *prompt*: The prompt to supply to the model.
                    Either a string or OpenAI-style list of dictionaries
                    representing messages (e.g., "human", "system").
        - *pydantic_model*: A Pydanatic model (sublass of `pydantic.BaseModel` that describes the desired output format.
                           Output will be a desired Pydantic object.
                           If `put_format=None`, then output is a string.
        - *attempt_fix*: Use an LLM call in attempt to correct malformed or incomplete outputs
        - *fix_llm*:  LLM to use for fixing (e.g., `langchain_openai.ChatOpenAI()`). If `None`, then existing `LLM.llm` used.
        - *stop*: a list of strings to stop generation when encountered. 
                  This value will override the `stop` parameter supplied to `LLM` constructor.
        """
        # setup up prompt for output parsing
        prompt, output_parser = self._format_pydantic_prompt(prompt, pydantic_model)

        # generate output
        output = self.prompt(prompt, stop=stop, **kwargs)

        # set parser
        fix_llm = fix_llm if fix_llm else self.llm
        parser = OutputFixingParser.from_llm(parser=output_parser, llm=fix_llm)\
                if attempt_fix else output_parser

        # parse output into Pydantic class
        try:
            return parser.parse(output)
        except Exception as e:
            try:
                from onprem.llm.helpers import extract_json
                json_string = extract_json(output)
                if json_string:
                    return parser.parse(json_string)
                raise Exception(str(e))
            except Exception as e:
                print()
                print()
                warnings.warn(f'LLM output was malformed or incomplete, so returning raw string output: {str(e)}')
                print()
                return output


    def _prompt_internal(self,
                        prompt: Union[str, List[Dict]],
                        image_path_or_url: Optional[str] = None,
                        prompt_template: Optional[str] = None,
                        stop: list = [],
                        truncate_prompt: bool = False,
                        truncate_strategy: str = 'start',
                        invoke_fn: Optional[Callable] = None,
                        **kwargs):
        """
        Internal prompt method that handles both sync and async via invoke_fn callback.
        """
        # Extract response_format from kwargs if present
        response_format = kwargs.pop('response_format', None)
        
        # Load LLM once
        llm = self.load_llm()
        
        # Handle structured output BEFORE any other processing
        if response_format:
            if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
                # This is vLLM native format - pass it directly to the API call
                kwargs['response_format'] = response_format
                response_format = None  # Don't trigger LangChain handling
            else:
                if hasattr(llm, 'with_structured_output'):
                    try:
                        # Try native structured output first
                        llm = llm.with_structured_output(response_format)
                        # Continue with normal processing using the wrapped LLM
                    except NotImplementedError:
                        # Fall back to pydantic_prompt immediately, before any template processing
                        return self.pydantic_prompt(prompt if isinstance(prompt, str) else prompt,
                                                  pydantic_model=response_format)
                else:
                    # No native support, use pydantic_prompt immediately
                    return self.pydantic_prompt(prompt if isinstance(prompt, str) else prompt,
                                              pydantic_model=response_format)

        # prompt is a list of dictionaries representing messages
        if isinstance(prompt, list):
            if self.is_llamacpp():
                # LangChain's LlamaCpp does not provide access to create_chat_completion,
                # so access it directly (with streaming disabled)
                res = llm.client.create_chat_completion(prompt)
                result = res['choices'][0]['message']['content']
            else:
                try:
                    res = invoke_fn(llm, prompt, stop=stop, **kwargs)
                except Exception as e: # stop param fails with GPT-4o vision prompts
                    res = invoke_fn(llm, prompt, **kwargs)
                result = res.content if isinstance(res, AIMessage) else res
        # prompt is string
        else:
            if image_path_or_url:
                prompt = self._format_image_prompt(prompt, image_path_or_url)
                res = invoke_fn(llm, prompt, **kwargs) # including stop causes errors in gpt-4o
                result = res.content if isinstance(res, AIMessage) else res
            else:
                # set prompt template
                prompt_template = self.prompt_template if prompt_template is None else prompt_template

                # truncate prompt
                if truncate_prompt and self.is_hf():
                    prompt = helpers.truncate_prompt(llm.llm.pipeline,
                                                     prompt,
                                                     max_gen_tokens=self.max_tokens,
                                                     truncate_from=truncate_strategy,
                                                     prompt_template=prompt_template)
                elif truncate_prompt and self.is_llamacpp():
                    prompt = helpers.truncate_prompt(self.llm.client,
                                                     prompt,
                                                     max_gen_tokens=self.max_tokens,
                                                     truncate_from=truncate_strategy,
                                                     prompt_template=prompt_template)

                # apply prompt template
                if prompt_template:
                    prompt = format_string(prompt_template, prompt=prompt)

                # set stop characters
                stop = stop if stop else self.stop

                # handle hf models
                if self.is_hf():
                    tokenizer = llm.llm.pipeline.tokenizer
                    # FIX for #113/#114
                    prompt = [{'role':'user', 'content':prompt}] if tokenizer.chat_template else prompt
                    # Call HF pipeline directly instead of `invoke`
                    # since LangChain is not passing along stop_strings
                    # parameter to pipeline
                    if 'max_tokens' in kwargs:
                        kwargs['max_new_tokens'] = kwargs['max_tokens']
                        del kwargs['max_tokens']
                    result = llm.llm.pipeline(prompt,
                                           stop_strings=stop if stop else None,
                                           tokenizer=tokenizer,
                                           **kwargs)[0]['generated_text']

                # handle other models (e.g., llama_cpp, LLMs served through APIs)
                else:
                    res = invoke_fn(llm, prompt, stop=stop, **kwargs)
                    result = res.content if isinstance(res, AIMessage) else res

        
        return result

    def prompt(self,
               prompt: Union[str, List[Dict]],
               image_path_or_url: Optional[str] = None,
               prompt_template: Optional[str] = None,
               stop: list = [],
               truncate_prompt: bool = False,
               truncate_strategy: str = 'start',
               **kwargs):
        """
        Send prompt to LLM to generate a response.
        Extra keyword arguments are sent directly to the model invocation.

        **Args:**

        - *prompt*: The prompt to supply to the model.
                    Either a string or OpenAI-style list of dictionaries
                    representing messages (e.g., "human", "system").
        - *image_path_or_url*: Path or URL to an image file
        - *prompt_template*: Optional prompt template (must have a variable named "prompt").
                             This value will override any `prompt_template` value supplied 
                             to `LLM` constructor.
        - *stop*: a list of strings to stop generation when encountered. 
                  This value will override the `stop` parameter supplied to `LLM` constructor.
        - *truncate_prompt*: Truncate long string prompts. Only applies to `llama-cpp-python` and `transformers` LLMs.
        - *truncate_strategy*: Either 'first' (keep latest) or 'last` (keep earliest). Ignored if `truncate_prompt=False`.
        - *response_format*: A Pydantic model class for structured output. When provided, the LLM will 
                            return a Pydantic object instead of a string. Uses native structured output 
                            (OpenAI, Azure, etc.) when available, otherwise falls back to the `pydantic_prompt` 
                            method which uses prompt-based parsing. Invoke `pydantic_prompt` directly for
                            more control over prompt-based parsing.

        """
        return self._prompt_internal(
            prompt=prompt,
            image_path_or_url=image_path_or_url,
            prompt_template=prompt_template,
            stop=stop,
            truncate_prompt=truncate_prompt,
            truncate_strategy=truncate_strategy,
            invoke_fn=self._invoke_llm,
            **kwargs
        )

    async def aprompt(self,
                      prompt: Union[str, List[Dict]],
                      image_path_or_url: Optional[str] = None,
                      prompt_template: Optional[str] = None,
                      stop: list = [],
                      truncate_prompt: bool = False,
                      truncate_strategy: str = 'start',
                      **kwargs):
        """
        Async version of prompt method. For cloud/API models, uses native async.
        For local models (llama.cpp/transformers), runs in thread pool to avoid blocking.

        **Args:**

        - *prompt*: The prompt to supply to the model.
                    Either a string or OpenAI-style list of dictionaries
                    representing messages (e.g., "human", "system").
        - *image_path_or_url*: Path or URL to an image file
        - *prompt_template*: Optional prompt template (must have a variable named "prompt").
                             This value will override any `prompt_template` value supplied 
                             to `LLM` constructor.
        - *stop*: a list of strings to stop generation when encountered. 
                  This value will override the `stop` parameter supplied to `LLM` constructor.
        - *truncate_prompt*: Truncate long string prompts. Only applies to `llama-cpp-python` and `transformers` LLMs.
        - *truncate_strategy*: Either 'first' (keep latest) or 'last` (keep earliest). Ignored if `truncate_prompt=False`.
        - *response_format*: A Pydantic model class for structured output. When provided, the LLM will 
                            return a Pydantic object instead of a string. Uses native structured output 
                            (OpenAI, Azure, etc.) when available, otherwise falls back to the `pydantic_prompt` 
                            method which uses prompt-based parsing.
        """
        # For local models, use thread executor to avoid blocking
        if self.is_hf() or self.is_llamacpp():
            import asyncio
            loop = asyncio.get_event_loop()
            return await loop.run_in_executor(
                None, 
                self.prompt,
                prompt, image_path_or_url, prompt_template, stop,
                truncate_prompt, truncate_strategy,
                **kwargs
            )
        else:
            # Cloud models - use async callback
            return await self._prompt_internal(
                prompt=prompt,
                image_path_or_url=image_path_or_url,
                prompt_template=prompt_template,
                stop=stop,
                truncate_prompt=truncate_prompt,
                truncate_strategy=truncate_strategy,
                invoke_fn=self._ainvoke_llm,
                **kwargs
            )

    def _invoke_llm(self, llm, prompt, stop=[], **kwargs):
        """Sync LLM invocation"""
        return llm.invoke(prompt, stop=stop, **kwargs)

    async def _ainvoke_llm(self, llm, prompt, stop=[], **kwargs):
        """Async LLM invocation"""
        return await llm.ainvoke(prompt, stop=stop, **kwargs)



    def load_chatbot(self):
        """
        Prepares chat functionality with message-based conversation history.
        """
        if self.chatbot is None:
            llm = self.load_llm()

            # Store chat history as list of messages
            # Structure: {'llm': model, 'history': [messages]}
            self.chatbot = {
                'llm': llm,
                'history': []
            }

        return self.chatbot


    def query(self, *args, **kwargs):
        """
        Perform a semantic search of vectorstore.
        """
        rag_pipeline = self.load_rag_pipeline()
        return rag_pipeline.semantic_search(*args, **kwargs)


    def semantic_search(self, *args, **kwargs):
        """
        Perform a semantic search of the vector DB. Alias for `LLM.query`.
        This method delegates to RAGPipeline. See RAGPipeline.semantic_search for parameter details.
        """
        rag_pipeline = self.load_rag_pipeline()
        return rag_pipeline.semantic_search(*args, **kwargs)




    def load_rag_pipeline(self):
        """
        Load RAG pipeline instance.
        """
        if self._rag_pipeline is None:
            from ..pipelines.rag import RAGPipeline
            self._rag_pipeline = RAGPipeline(self)
        return self._rag_pipeline

    def ask(self, question: str, **kwargs):
        """
        Answer a question based on source documents fed to the `LLM.ingest` method.
        This method delegates to RAGPipeline. See RAGPipeline.ask for parameter details.
        Extra keyword arguments are sent directly to `LLM.prompt`.
        Returns a dictionary with keys: `answer`, `source_documents`, `question`
        """
        rag_pipeline = self.load_rag_pipeline()
        return rag_pipeline.ask(question=question, **kwargs)


    def chat(self, prompt: str, prompt_template=None, **kwargs):
        """
        Chat with LLM.

        **Args:**

        - *prompt*: a question you want to ask

        """
        from langchain_core.messages import HumanMessage

        # Apply prompt template if provided
        prompt_template = self.prompt_template if prompt_template is None else prompt_template
        if prompt_template:
            prompt = format_string(prompt_template, prompt=prompt)

        # Load chatbot (creates history if needed)
        chatbot = self.load_chatbot()

        # Add user message to history
        chatbot['history'].append(HumanMessage(content=prompt))

        # Invoke LLM with full conversation history
        response = chatbot['llm'].invoke(chatbot['history'], **kwargs)

        # Add AI response to history
        chatbot['history'].append(response)

        # Return response content
        return response.content if hasattr(response, 'content') else str(response)

