feat(core): user should be able to configure default types of default…

… LLMs

packages/ragbits-core/src/ragbits/core/cli.py

@@ -9,6 +9,7 @@
 from rich import print as pprint
 
 from ragbits.core.config import core_config
+from ragbits.core.llms.base import LLMType
 from ragbits.core.prompt.prompt import Prompt
 
 
@@ -38,7 +39,7 @@ def register(app: typer.Typer) -> None:
     @prompts_app.command()
     def lab(
         file_pattern: str = core_config.prompt_path_pattern,
-        llm_factory: str | None = core_config.default_llm_factory,
+        llm_factory: str | None = core_config.default_llm_factories[LLMType.TEXT],
     ) -> None:
         """
         Launches the interactive application for listing, rendering, and testing prompts
@@ -73,15 +74,16 @@ def render(prompt_path: str, payload: str | None = None) -> None:
 
     @prompts_app.command(name="exec")
     def execute(
-        prompt_path: str, payload: str | None = None, llm_factory: str | None = core_config.default_llm_factory
+        prompt_path: str,
+        payload: str | None = None,
+        llm_factory: str | None = core_config.default_llm_factories[LLMType.TEXT],
     ) -> None:
         """
         Executes a prompt using the specified prompt class and LLM factory.
 
         Raises:
             ValueError: If `llm_factory` is not provided.
         """
-
         from ragbits.core.llms.factory import get_llm_from_factory
 
         prompt = _render(prompt_path=prompt_path, payload=payload)

packages/ragbits-core/src/ragbits/core/config.py

@@ -1,5 +1,6 @@
 from pydantic import BaseModel
 
+from ragbits.core.llms.base import LLMType
 from ragbits.core.utils._pyproject import get_config_instance
 
 
@@ -11,8 +12,12 @@ class CoreConfig(BaseModel):
     # Pattern used to search for prompt files
     prompt_path_pattern: str = "**/prompt_*.py"
 
-    # Path to a function that returns an LLM object, e.g. "my_project.llms.get_llm"
-    default_llm_factory: str | None = None
+    # Path to a functions that returns LLM objects, e.g. "my_project.llms.get_llm"
+    default_llm_factories: dict[LLMType, str | None] = {
+        LLMType.TEXT: None,
+        LLMType.VISION: None,
+        LLMType.STRUCTURED_OUTPUT: None,
+    }
 
 
 core_config = get_config_instance(CoreConfig, subproject="core")

packages/ragbits-core/src/ragbits/core/llms/base.py

@@ -1,3 +1,4 @@
+import enum
 from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import Generic, cast, overload
@@ -7,6 +8,16 @@
 from .clients.base import LLMClient, LLMClientOptions, LLMOptions
 
 
+class LLMType(enum.Enum):
+    """
+    Types of LLMs based on supported features
+    """
+
+    TEXT = "text"
+    VISION = "vision"
+    STRUCTURED_OUTPUT = "structured_output"
+
+
 class LLM(Generic[LLMClientOptions], ABC):
     """
     Abstract class for interaction with Large Language Model.

packages/ragbits-core/src/ragbits/core/llms/factory.py

@@ -1,7 +1,7 @@
 import importlib
 
 from ragbits.core.config import core_config
-from ragbits.core.llms.base import LLM
+from ragbits.core.llms.base import LLM, LLMType
 from ragbits.core.llms.litellm import LiteLLM
 
 
@@ -21,28 +21,34 @@ def get_llm_from_factory(factory_path: str) -> LLM:
     return function()
 
 
-def has_default_llm() -> bool:
+def has_default_llm(llm_type: LLMType = LLMType.TEXT) -> bool:
     """
     Check if the default LLM factory is set in the configuration.
 
     Returns:
         bool: Whether the default LLM factory is set.
     """
-    return core_config.default_llm_factory is not None
+    default_factory = core_config.default_llm_factories.get(llm_type, None)
+    return default_factory is not None
 
 
-def get_default_llm() -> LLM:
+def get_default_llm(llm_type: LLMType = LLMType.TEXT) -> LLM:
     """
     Get an instance of the default LLM using the factory function
     specified in the configuration.
 
+    Args:
+        llm_type: type of the LLM to get, defaults to text
+
     Returns:
         LLM: An instance of the default LLM.
 
     Raises:
-        ValueError: If the default LLM factory is not set.
+        ValueError: If the default LLM factory is not set or expected llm type is not defined in config
     """
-    factory = core_config.default_llm_factory
+    if llm_type not in core_config.default_llm_factories:
+        raise ValueError(f"Default LLM of type {llm_type} is not defined in pyproject.toml config.")
+    factory = core_config.default_llm_factories[llm_type]
     if factory is None:
         raise ValueError("Default LLM factory is not set")
 
@@ -58,3 +64,14 @@ def simple_litellm_factory() -> LLM:
         LLM: An instance of the LiteLLM.
     """
     return LiteLLM()
+
+
+def simple_litellm_vision_factory() -> LLM:
+    """
+    A basic LLM factory that creates an LiteLLM instance with the vision enabled model,
+    default options, and assumes that the API key is set in the environment.
+
+    Returns:
+        LLM: An instance of the LiteLLM.
+    """
+    return LiteLLM(model_name="gpt-4o-mini")

packages/ragbits-core/src/ragbits/core/prompt/lab/app.py

@@ -15,6 +15,7 @@
 
 from ragbits.core.config import core_config
 from ragbits.core.llms import LLM
+from ragbits.core.llms.base import LLMType
 from ragbits.core.llms.factory import get_llm_from_factory
 from ragbits.core.prompt import Prompt
 from ragbits.core.prompt.discovery import PromptDiscovery
@@ -137,7 +138,7 @@ def get_input_type_fields(obj: BaseModel | None) -> list[dict]:
 
 def lab_app(  # pylint: disable=missing-param-doc
     file_pattern: str = core_config.prompt_path_pattern,
-    llm_factory: str | None = core_config.default_llm_factory,
+    llm_factory: str | None = core_config.default_llm_factories[LLMType.TEXT],
 ) -> None:
     """
     Launches the interactive application for listing, rendering, and testing prompts

packages/ragbits-core/src/ragbits/core/utils/_pyproject.py

@@ -1,3 +1,5 @@
+import enum
+import importlib
 from pathlib import Path
 from typing import Any, TypeVar
 
@@ -64,7 +66,7 @@ def get_config_instance(
     model: type[ConfigModelT], subproject: str | None = None, current_dir: Path | None = None
 ) -> ConfigModelT:
     """
-    Creates an instace of pydantic model loaded with the configuration from pyproject.toml.
+    Creates an instance of pydantic model loaded with the configuration from pyproject.toml.
 
     Args:
         model (Type[BaseModel]): The pydantic model to instantiate.
@@ -81,4 +83,15 @@ def get_config_instance(
     config = get_ragbits_config(current_dir)
     if subproject:
         config = config.get(subproject, {})
+    config["default_llm_factories"] = {_resolve_enum_member(k): v for k, v in config["default_llm_factories"].items()}
     return model(**config)
+
+
+def _resolve_enum_member(enum_string: str) -> enum.Enum:
+    module_name, class_name, member_name = enum_string.rsplit(".", 2)
+    module = importlib.import_module(module_name)
+    enum_class = getattr(module, class_name)
+    try:
+        return getattr(enum_class, member_name)
+    except AttributeError as err:
+        raise ValueError("Unsupported LLMType provided in default_llm_factories in pyproject.yaml") from err

packages/ragbits-core/tests/unit/llms/factory/test_get_default_llm.py

@@ -9,7 +9,9 @@ def test_get_default_llm(monkeypatch: pytest.MonkeyPatch) -> None:
     """
     Test the get_llm_from_factory function.
     """
-    monkeypatch.setattr(core_config, "default_llm_factory", "factory.test_get_llm_from_factory.mock_llm_factory")
+    monkeypatch.setattr(
+        core_config, "default_llm_factories", {"text": "factory.test_get_llm_from_factory.mock_llm_factory"}
+    )
 
     llm = get_default_llm()
     assert isinstance(llm, LiteLLM)

packages/ragbits-core/tests/unit/llms/factory/test_has_default_llm.py

@@ -8,7 +8,7 @@ def test_has_default_llm(monkeypatch: pytest.MonkeyPatch) -> None:
     """
     Test the has_default_llm function when the default LLM factory is not set.
     """
-    monkeypatch.setattr(core_config, "default_llm_factory", None)
+    monkeypatch.setattr(core_config, "default_llm_factories", {})
 
     assert has_default_llm() is False
 
@@ -17,6 +17,6 @@ def test_has_default_llm_false(monkeypatch: pytest.MonkeyPatch) -> None:
     """
     Test the has_default_llm function when the default LLM factory is set.
     """
-    monkeypatch.setattr(core_config, "default_llm_factory", "my_project.llms.get_llm")
+    monkeypatch.setattr(core_config, "default_llm_factories", {"text": "my_project.llms.get_llm"})
 
     assert has_default_llm() is True

packages/ragbits-document-search/src/ragbits/document_search/ingestion/providers/unstructured/images.py

@@ -5,8 +5,8 @@
 from unstructured.documents.elements import Element as UnstructuredElement
 from unstructured.documents.elements import ElementType
 
-from ragbits.core.llms.base import LLM
-from ragbits.core.llms.litellm import LiteLLM
+from ragbits.core.llms.base import LLM, LLMType
+from ragbits.core.llms.factory import get_default_llm
 from ragbits.document_search.documents.document import DocumentMeta, DocumentType
 from ragbits.document_search.documents.element import Element, ImageElement
 from ragbits.document_search.ingestion.providers.unstructured.default import UnstructuredDefaultProvider
@@ -17,8 +17,6 @@
     to_text_element,
 )
 
-DEFAULT_LLM_IMAGE_SUMMARIZATION_MODEL = "gpt-4o-mini"
-
 
 class UnstructuredImageProvider(UnstructuredDefaultProvider):
     """
@@ -53,7 +51,8 @@ def __init__(
             llm: llm to use
         """
         super().__init__(partition_kwargs, chunking_kwargs, api_key, api_server, use_api)
-        self.image_summarizer = ImageDescriber(llm or LiteLLM(DEFAULT_LLM_IMAGE_SUMMARIZATION_MODEL))
+        self.image_describer: ImageDescriber | None = None
+        self._llm = llm
 
     async def _chunk_and_convert(
         self, elements: list[UnstructuredElement], document_meta: DocumentMeta, document_path: Path
@@ -79,7 +78,10 @@ async def _to_image_element(
         )
 
         img_bytes = crop_and_convert_to_bytes(image, top_x, top_y, bottom_x, bottom_y)
-        image_description = await self.image_summarizer.get_image_description(img_bytes)
+        if self.image_describer is None:
+            llm_to_use = self._llm if self._llm is not None else get_default_llm(LLMType.VISION)
+            self.image_describer = ImageDescriber(llm_to_use)
+        image_description = await self.image_describer.get_image_description(img_bytes)
         return ImageElement(
             description=image_description,
             ocr_extracted_text=element.text,

pyproject.toml

@@ -183,3 +183,8 @@ known-third-party = [
     "requests", "scipy", "setuptools", "shapely", "skimage", "sklearn", "streamlit",
     "torch", "torchvision", "tqdm", "typer"
 ]
+
+[tool.ragbits.core.default_llm_factories]
+"ragbits.core.llms.base.LLMType.TEXT" = "ragbits.core.llms.factory.simple_litellm_factory"
+"ragbits.core.llms.base.LLMType.VISION" = "ragbits.core.llms.factory.simple_litellm_vision_factory"
+"ragbits.core.llms.base.LLMType.STRUCTURED_OUTPUT" = "ragbits.core.llms.factory.simple_litellm_vision_factory"