move importing llm models inside their respective functions

Author: MAfarrag

src/llama_utils/utils/models.py

@@ -3,13 +3,44 @@
 import os
 from warnings import warn
 
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.llms.azure_openai import AzureOpenAI
-from llama_index.llms.ollama import Ollama
-
 
 def azure_open_ai(model_id: str = "gpt-4o", engine: str = "4o"):
-    """Get the Azure OpenAI model."""
+    """Get the Azure OpenAI model.
+
+    Parameters
+    ----------
+    model_id: str, optional, default is "gpt-4o"
+        The model ID.
+    engine: str, optional, default is "4o"
+        The engine.
+
+    Returns
+    -------
+    AzureOpenAI
+        The Azure OpenAI model.
+
+    Raises
+    ------
+    ImportError
+        If the `llama-index-llms-azure-openai` package is not installed.
+
+    Examples
+    --------
+    >>> from llama_utils.utils.models import azure_open_ai
+    >>> from dotenv import load_dotenv
+    >>> load_dotenv() # doctest: +SKIP
+    >>> llm = azure_open_ai() # doctest: +SKIP
+    >>> print(llm.model) # doctest: +SKIP
+    gpt-4o
+    >>> print(llm.metadata) # doctest: +SKIP
+    context_window=128000 num_output=-1 is_chat_model=True is_function_calling_model=True model_name='gpt-4o' system_role=<MessageRole.SYSTEM: 'system'>
+    """
+    try:
+        from llama_index.llms.azure_openai import AzureOpenAI
+    except ImportError:
+        raise ImportError(
+            "Please install the `llama-index-llms-azure-openai` package to use the Azure OpenAI model."
+        )
     endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
     api_key = os.environ.get("AZURE_OPENAI_API_KEY")
     api_version = os.environ.get("AZURE_OPENAI_API_VERSION")
@@ -41,14 +72,36 @@ def get_ollama_llm(model_id: str = "llama3"):
     -------
     Ollama
         The Ollama LLM.
+
+    Raises
+    ------
+    ImportError
+        If the `llama-index-llms-ollama` package is not installed.
+
+    Examples
+    --------
+    >>> from llama_utils.utils.models import get_ollama_llm
+    >>> llm = get_ollama_llm()
+    >>> print(llm.model)
+    llama3
+    >>> print(llm.base_url)
+    http://localhost:11434
+    >>> print(llm.metadata)
+    context_window=3900 num_output=256 is_chat_model=True is_function_calling_model=True model_name='llama3' system_role=<MessageRole.SYSTEM: 'system'>
     """
+    try:
+        from llama_index.llms.ollama import Ollama
+    except ImportError:
+        raise ImportError(
+            "Please install the `llama-index-llms-ollama` package to use the Ollama model."
+        )
     llm = Ollama(model=model_id, request_timeout=360.0)
     return llm
 
 
 def get_hugging_face_embedding(
     model_name: str = "BAAI/bge-base-en-v1.5", cache_folder: str = None
-) -> HuggingFaceEmbedding:
+):
     """Get the hugging face embedding model.
 
     Parameters
@@ -62,6 +115,29 @@ def get_hugging_face_embedding(
     -------
     HuggingFaceEmbedding
         The hugging face embedding model.
+
+    Raises
+    ------
+    ImportError
+        If the `llama-index-embeddings-huggingface` package is not installed.
+
+    Examples
+    --------
+    >>> from llama_utils.utils.models import get_hugging_face_embedding
+    >>> embedding = get_hugging_face_embedding()
+    >>> print(embedding.model_name)
+    BAAI/bge-base-en-v1.5
+    >>> print(embedding.max_length)
+    512
+    >>> print(embedding.embed_batch_size)
+    10
     """
+    try:
+        from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+    except ImportError:
+        raise ImportError(
+            "Please install the `llama-index-embeddings-huggingface` package to use the Hugging Face embedding model."
+        )
+
     embedding = HuggingFaceEmbedding(model_name=model_name, cache_folder=cache_folder)
     return embedding