refactor: make litellm and litellmembedder consistent (#463)

Author: mhordynski

docs/api_reference/core/embeddings.md

@@ -5,3 +5,7 @@
 ::: ragbits.core.embeddings.local.LocalEmbedder
 
 ::: ragbits.core.embeddings.litellm.LiteLLMEmbedder
+
+::: ragbits.core.embeddings.fastembed.FastEmbedEmbedder
+
+::: ragbits.core.embeddings.fastembed.FastEmbedSparseEmbedder

docs/how-to/llms/use_local_llms.md

@@ -71,7 +71,7 @@ from ragbits.core.prompt.base import SimplePrompt
 
 
 async def main() -> None:
-    llm = LiteLLM(model_name="openai/local", api_key="<api_key>", base_url="http://127.0.0.1:8080")
+    llm = LiteLLM(model_name="openai/local", api_key="<api_key>", api_base="http://127.0.0.1:8080")
     prompt = SimplePrompt("Tell me a joke about software developers.")
     response = await llm.generate(prompt)
     print(response)
@@ -99,7 +99,7 @@ from ragbits.core.prompt.base import SimplePrompt
 
 
 async def main() -> None:
-    llm = LiteLLM(model_name="hosted_vllm/<model_name>", base_url="http://127.0.0.1:8000/v1")
+    llm = LiteLLM(model_name="hosted_vllm/<model_name>", api_base="http://127.0.0.1:8000/v1")
     prompt = SimplePrompt("Tell me a joke about software developers.")
     response = await llm.generate(prompt)
     print(response)
@@ -123,7 +123,7 @@ from ragbits.core.embeddings.litellm import LiteLLMEmbedder
 
 
 async def main() -> None:
-    embedder = LiteLLMEmbedder(model="hosted_vllm/<model_name>", api_base="http://127.0.0.1:8000/v1")
+    embedder = LiteLLMEmbedder(model_name="hosted_vllm/<model_name>", api_base="http://127.0.0.1:8000/v1")
     embeddings = await embedder.embed_text(["Hello"])
     print(len(embeddings[0]))
 

examples/document-search/basic.py

@@ -69,7 +69,7 @@ async def main() -> None:
     Run the example.
     """
     embedder = LiteLLMEmbedder(
-        model="text-embedding-3-small",
+        model_name="text-embedding-3-small",
     )
     vector_store = InMemoryVectorStore(embedder=embedder)
     document_search = DocumentSearch(

packages/ragbits-core/CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 - Make the score in VectorStoreResult consistent (always bigger is better)
 - Add router option to LiteLLMEmbedder (#440)
+- Make LLM / Embedder APIs consistent (#463)
 - New methods in Prompt class for appending conversation history (#480)
 - Fix: make unflatten_dict symmetric to flatten_dict (#461)
 - Cost and capabilities config for custom litellm models (#481)

packages/ragbits-core/src/ragbits/core/embeddings/fastembed.py

@@ -57,7 +57,7 @@ async def embed_text(self, data: list[str], options: EmbedderOptionsT | None = N
         """
         merged_options = (self.default_options | options) if options else self.default_options
         with trace(
-            data=data, model_name=self.model_name, model=repr(self._model), options=merged_options.dict()
+            data=data, model_name=self.model_name, model_obj=repr(self._model), options=merged_options.dict()
         ) as outputs:
             embeddings = [[float(x) for x in result] for result in self._model.embed(data, **merged_options.dict())]
             outputs.embeddings = embeddings
@@ -104,7 +104,7 @@ async def embed_text(self, data: list[str], options: EmbedderOptionsT | None = N
         """
         merged_options = (self.default_options | options) if options else self.default_options
         with trace(
-            data=data, model_name=self.model_name, model=repr(self._model), options=merged_options.dict()
+            data=data, model_name=self.model_name, model_obj=repr(self._model), options=merged_options.dict()
         ) as outputs:
             outputs.embeddings = [
                 SparseVector(values=[float(x) for x in result.values], indices=[int(x) for x in result.indices])

packages/ragbits-core/src/ragbits/core/embeddings/litellm.py

@@ -36,9 +36,11 @@ class LiteLLMEmbedder(Embedder[LiteLLMEmbedderOptions]):
 
     def __init__(
         self,
-        model: str = "text-embedding-3-small",
+        model_name: str = "text-embedding-3-small",
         default_options: LiteLLMEmbedderOptions | None = None,
+        *,
         api_base: str | None = None,
+        base_url: str | None = None,  # Alias for api_base
         api_key: str | None = None,
         api_version: str | None = None,
         router: litellm.Router | None = None,
@@ -47,19 +49,21 @@ def __init__(
         Constructs the LiteLLMEmbeddingClient.
 
         Args:
-            model: Name of the [LiteLLM supported model](https://docs.litellm.ai/docs/embedding/supported_embedding)\
+            model_name: Name of the [LiteLLM supported model](https://docs.litellm.ai/docs/embedding/supported_embedding)\
                 to be used. Default is "text-embedding-3-small".
             default_options: Default options to pass to the LiteLLM API.
             api_base: The API endpoint you want to call the model with.
-            api_key: API key to be used. API key to be used. If not specified, an environment variable will be used,
+            base_url: Alias for api_base. If both are provided, api_base takes precedence.
+            api_key: API key to be used. If not specified, an environment variable will be used,
                 for more information, follow the instructions for your specific vendor in the\
                 [LiteLLM documentation](https://docs.litellm.ai/docs/embedding/supported_embedding).
             api_version: The API version for the call.
             router: Router to be used to [route requests](https://docs.litellm.ai/docs/routing) to different models.
         """
         super().__init__(default_options=default_options)
-        self.model = model
-        self.api_base = api_base
+
+        self.model_name = model_name
+        self.api_base = api_base or base_url
         self.api_key = api_key
         self.api_version = api_version
         self.router = router
@@ -85,7 +89,7 @@ async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | No
 
         with trace(
             data=data,
-            model=self.model,
+            model=self.model_name,
             api_base=self.api_base,
             api_version=self.api_version,
             options=merged_options.dict(),
@@ -94,7 +98,7 @@ async def embed_text(self, data: list[str], options: LiteLLMEmbedderOptions | No
                 entrypoint = self.router or litellm
                 response = await entrypoint.aembedding(
                     input=data,
-                    model=self.model,
+                    model=self.model_name,
                     api_base=self.api_base,
                     api_key=self.api_key,
                     api_version=self.api_version,
@@ -132,4 +136,9 @@ def from_config(cls, config: dict[str, Any]) -> Self:
         if "router" in config:
             router = litellm.router.Router(model_list=config["router"])
             config["router"] = router
+
+        # Map base_url to api_base if present
+        if "base_url" in config and "api_base" not in config:
+            config["api_base"] = config.pop("base_url")
+
         return super().from_config(config)

packages/ragbits-core/src/ragbits/core/embeddings/local.py

@@ -73,7 +73,7 @@ async def embed_text(self, data: list[str], options: LocalEmbedderOptions | None
         with trace(
             data=data,
             model_name=self.model_name,
-            model=repr(self.model),
+            model_obj=repr(self.model),
             tokenizer=repr(self.tokenizer),
             device=self.device,
             options=merged_options.dict(),

packages/ragbits-core/src/ragbits/core/embeddings/vertex_multimodal.py

@@ -30,8 +30,9 @@ class VertexAIMultimodelEmbedder(Embedder[LiteLLMEmbedderOptions]):
 
     def __init__(
         self,
-        model: str = "multimodalembedding",
+        model_name: str = "multimodalembedding",
         api_base: str | None = None,
+        base_url: str | None = None,  # Alias for api_base
         api_key: str | None = None,
         concurency: int = 10,
         default_options: LiteLLMEmbedderOptions | None = None,
@@ -40,8 +41,9 @@ def __init__(
         Constructs the embedding client for multimodal VertexAI models.
 
         Args:
-            model: One of the VertexAI multimodal models to be used. Default is "multimodalembedding".
+            model_name: One of the VertexAI multimodal models to be used. Default is "multimodalembedding".
             api_base: The API endpoint you want to call the model with.
+            base_url: Alias for api_base. If both are provided, api_base takes precedence.
             api_key: API key to be used. If not specified, an environment variable will be used.
             concurency: The number of concurrent requests to make to the API.
             default_options: Additional options to pass to the API.
@@ -54,17 +56,18 @@ def __init__(
             raise ImportError("You need to install the 'litellm' extra requirements to use LiteLLM embeddings models")
 
         super().__init__(default_options=default_options)
-        if model.startswith(self.VERTEX_AI_PREFIX):
-            model = model[len(self.VERTEX_AI_PREFIX) :]
 
-        self.model = model
-        self.api_base = api_base
+        if model_name.startswith(self.VERTEX_AI_PREFIX):
+            model_name = model_name[len(self.VERTEX_AI_PREFIX) :]
+
+        self.model_name = model_name
+        self.api_base = api_base or base_url
         self.api_key = api_key
         self.concurency = concurency
 
         supported_models = VertexMultimodalEmbedding().SUPPORTED_MULTIMODAL_EMBEDDING_MODELS
-        if model not in supported_models:
-            raise ValueError(f"Model {model} is not supported by VertexAI multimodal embeddings")
+        if model_name not in supported_models:
+            raise ValueError(f"Model {model_name} is not supported by VertexAI multimodal embeddings")
 
     async def _embed(self, data: list[dict], options: LiteLLMEmbedderOptions | None = None) -> list[dict]:
         """
@@ -86,7 +89,7 @@ async def _embed(self, data: list[dict], options: LiteLLMEmbedderOptions | None
         merged_options = (self.default_options | options) if options else self.default_options
         with trace(
             data=data,
-            model=self.model,
+            model=self.model_name,
             api_base=self.api_base,
             options=merged_options.dict(),
         ) as outputs:
@@ -123,7 +126,7 @@ async def _call_litellm(
         async with semaphore:
             response = await litellm.aembedding(
                 input=[instance],
-                model=f"{self.VERTEX_AI_PREFIX}{self.model}",
+                model=f"{self.VERTEX_AI_PREFIX}{self.model_name}",
                 api_base=self.api_base,
                 api_key=self.api_key,
                 **options.dict(),

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

@@ -52,7 +52,8 @@ def __init__(
         model_name: str = "gpt-3.5-turbo",
         default_options: LiteLLMOptions | None = None,
         *,
-        base_url: str | None = None,
+        api_base: str | None = None,
+        base_url: str | None = None,  # Alias for api_base
         api_key: str | None = None,
         api_version: str | None = None,
         use_structured_output: bool = False,
@@ -66,7 +67,8 @@ def __init__(
             model_name: Name of the [LiteLLM supported model](https://docs.litellm.ai/docs/providers) to be used.\
                 Default is "gpt-3.5-turbo".
             default_options: Default options to be used.
-            base_url: Base URL of the LLM API.
+            api_base: Base URL of the LLM API.
+            base_url: Alias for api_base. If both are provided, api_base takes precedence.
             api_key: API key to be used. API key to be used. If not specified, an environment variable will be used,
                 for more information, follow the instructions for your specific vendor in the\
                 [LiteLLM documentation](https://docs.litellm.ai/docs/providers).
@@ -81,7 +83,7 @@ def __init__(
                 for more information.
         """
         super().__init__(model_name, default_options)
-        self.base_url = base_url
+        self.api_base = api_base or base_url
         self.api_key = api_key
         self.api_version = api_version
         self.use_structured_output = use_structured_output
@@ -187,7 +189,7 @@ async def _call_streaming(
         with trace(
             messages=prompt.chat,
             model=self.model_name,
-            base_url=self.base_url,
+            base_url=self.api_base,
             api_version=self.api_version,
             response_format=response_format,
             options=options.dict(),
@@ -222,7 +224,7 @@ async def _get_litellm_response(
             response = await entrypoint.acompletion(
                 messages=conversation,
                 model=self.model_name,
-                base_url=self.base_url,
+                base_url=self.api_base,
                 api_key=self.api_key,
                 api_version=self.api_version,
                 response_format=response_format,
@@ -250,6 +252,13 @@ def _get_response_format(
                 response_format = {"type": "json_object"}
         return response_format
 
+    @property
+    def base_url(self) -> str | None:
+        """
+        Returns the base URL of the LLM API. Alias for `api_base`.
+        """
+        return self.api_base
+
     @classmethod
     def from_config(cls, config: dict[str, Any]) -> Self:
         """
@@ -264,13 +273,18 @@ def from_config(cls, config: dict[str, Any]) -> Self:
         if "router" in config:
             router = litellm.router.Router(model_list=config["router"])
             config["router"] = router
+
+        # Map base_url to api_base if present
+        if "base_url" in config and "api_base" not in config:
+            config["api_base"] = config.pop("base_url")
+
         return super().from_config(config)
 
     def __reduce__(self) -> tuple[Callable, tuple]:
         config = {
             "model_name": self.model_name,
             "default_options": self.default_options.dict(),
-            "base_url": self.base_url,
+            "api_base": self.api_base,
             "api_key": self.api_key,
             "api_version": self.api_version,
             "use_structured_output": self.use_structured_output,

packages/ragbits-core/tests/unit/embeddings/test_from_config.py

@@ -12,7 +12,7 @@ def test_subclass_from_config_litellm():
         {
             "type": "ragbits.core.embeddings.litellm:LiteLLMEmbedder",
             "config": {
-                "model": "some_model",
+                "model_name": "some_model",
                 "default_options": {
                     "option1": "value1",
                     "option2": "value2",
@@ -22,7 +22,7 @@ def test_subclass_from_config_litellm():
     )
     embedder: Embedder = Embedder.subclass_from_config(config)
     assert isinstance(embedder, LiteLLMEmbedder)
-    assert embedder.model == "some_model"
+    assert embedder.model_name == "some_model"
     assert embedder.default_options == LiteLLMEmbedderOptions(
         dimensions=NOT_GIVEN,
         timeout=NOT_GIVEN,
@@ -66,7 +66,7 @@ def test_from_config_with_router():
     config = ObjectConstructionConfig(
         type="ragbits.core.embeddings.litellm:LiteLLMEmbedder",
         config={
-            "model": "text-embedding-3-small",
+            "model_name": "text-embedding-3-small",
             "api_key": "test_api_key",
             "router": [
                 {
@@ -91,7 +91,7 @@ def test_from_config_with_router():
     embedder: Embedder = Embedder.subclass_from_config(config)
     assert isinstance(embedder, LiteLLMEmbedder)
     assert embedder.api_base is None
-    assert embedder.model == "text-embedding-3-small"
+    assert embedder.model_name == "text-embedding-3-small"
     assert embedder.api_key == "test_api_key"
     assert isinstance(embedder.router, litellm.router.Router)
     assert len(embedder.router.model_list) == 2

packages/ragbits-core/tests/unit/llms/test_from_config.py

@@ -65,7 +65,7 @@ def test_from_config_with_router():
 
     llm: LLM = LLM.subclass_from_config(config)
     assert isinstance(llm, LiteLLM)
-    assert llm.base_url is None
+    assert llm.api_base is None
     assert llm.model_name == "gpt-4-turbo"
     assert llm.api_key == "test_api_key"
     assert isinstance(llm.router, litellm.router.Router)