Sync embedder models with latest meili settings

Author: ellnix

meilisearch/index.py

@@ -30,9 +30,11 @@
     HuggingFaceEmbedder,
     IndexStats,
     LocalizedAttributes,
+    OllamaEmbedder,
     OpenAiEmbedder,
     Pagination,
     ProximityPrecision,
+    RestEmbedder,
     TypoTolerance,
     UserProvidedEmbedder,
 )
@@ -958,12 +960,23 @@ def get_settings(self) -> Dict[str, Any]:
         )
 
         if settings.get("embedders"):
-            embedders: dict[str, OpenAiEmbedder | HuggingFaceEmbedder | UserProvidedEmbedder] = {}
+            embedders: dict[
+                str,
+                OpenAiEmbedder
+                | HuggingFaceEmbedder
+                | OllamaEmbedder
+                | RestEmbedder
+                | UserProvidedEmbedder,
+            ] = {}
             for k, v in settings["embedders"].items():
                 if v.get("source") == "openAi":
                     embedders[k] = OpenAiEmbedder(**v)
+                elif v.get("source") == "ollama":
+                    embedders[k] = OllamaEmbedder(**v)
                 elif v.get("source") == "huggingFace":
                     embedders[k] = HuggingFaceEmbedder(**v)
+                elif v.get("source") == "rest":
+                    embedders[k] = RestEmbedder(**v)
                 else:
                     embedders[k] = UserProvidedEmbedder(**v)
 
@@ -1886,12 +1899,23 @@ def get_embedders(self) -> Embedders | None:
         if not response:
             return None
 
-        embedders: dict[str, OpenAiEmbedder | HuggingFaceEmbedder | UserProvidedEmbedder] = {}
+        embedders: dict[
+            str,
+            OpenAiEmbedder
+            | HuggingFaceEmbedder
+            | OllamaEmbedder
+            | RestEmbedder
+            | UserProvidedEmbedder,
+        ] = {}
         for k, v in response.items():
             if v.get("source") == "openAi":
                 embedders[k] = OpenAiEmbedder(**v)
+            elif v.get("source") == "ollama":
+                embedders[k] = OllamaEmbedder(**v)
             elif v.get("source") == "huggingFace":
                 embedders[k] = HuggingFaceEmbedder(**v)
+            elif v.get("source") == "rest":
+                embedders[k] = RestEmbedder(**v)
             else:
                 embedders[k] = UserProvidedEmbedder(**v)
 

meilisearch/models/index.py

@@ -54,18 +54,26 @@ class ProximityPrecision(str, Enum):
     BY_ATTRIBUTE = "byAttribute"
 
 
+class EmbedderDistribution(CamelBase):
+    mean: float
+    sigma: float
+
+
 class LocalizedAttributes(CamelBase):
     attribute_patterns: List[str]
     locales: List[str]
 
 
 class OpenAiEmbedder(CamelBase):
     source: str = "openAi"
+    url: Optional[str] = None
     model: Optional[str] = None  # Defaults to text-embedding-3-small
     dimensions: Optional[int] = None  # Uses the model default
     api_key: Optional[str] = None  # Can be provided through a CLI option or environment variable
     document_template: Optional[str] = None
     document_template_max_bytes: Optional[int] = None  # Default to 400
+    distribution: Optional[EmbedderDistribution] = None
+    binary_quantized: Optional[bool] = None
 
 
 class HuggingFaceEmbedder(CamelBase):
@@ -74,12 +82,45 @@ class HuggingFaceEmbedder(CamelBase):
     revision: Optional[str] = None
     document_template: Optional[str] = None
     document_template_max_bytes: Optional[int] = None  # Default to 400
+    distribution: Optional[EmbedderDistribution] = None
+    binary_quantized: Optional[bool] = None
+
+
+class OllamaEmbedder(CamelBase):
+    source: str = "ollama"
+    url: Optional[str] = None
+    api_key: Optional[str] = None
+    model: str
+    document_template: Optional[str] = None
+    document_template_max_bytes: Optional[int] = None  # Default to 400
+    distribution: Optional[EmbedderDistribution] = None
+    binary_quantized: Optional[bool] = None
+
+
+class RestEmbedder(CamelBase):
+    source: str = "rest"
+    url: str
+    api_key: Optional[str]  # required for protected APIs
+    document_template: Optional[str] = None
+    document_template_max_bytes: Optional[int] = None  # Default to 400
+    request: Dict[str, Any]
+    response: Dict[str, Any]
+    distribution: Optional[EmbedderDistribution] = None
+    headers: Optional[Dict[str, Any]] = None
+    binary_quantized: Optional[bool] = None
 
 
 class UserProvidedEmbedder(CamelBase):
     source: str = "userProvided"
     dimensions: int
+    distribution: Optional[EmbedderDistribution] = None
+    binary_quantized: Optional[bool] = None
 
 
 class Embedders(CamelBase):
-    embedders: Dict[str, Union[OpenAiEmbedder, HuggingFaceEmbedder, UserProvidedEmbedder]]
+    embedders: Dict[
+        str,
+        Union[
+            OpenAiEmbedder, HuggingFaceEmbedder, OllamaEmbedder, RestEmbedder, UserProvidedEmbedder
+        ],
+    ]