RDoc-3621 Add missing embeddings_max_concurrent_batches parameter

Author: poissoncorp

ravendb/documents/operations/ai/abstract_ai_settings.py

@@ -3,8 +3,8 @@
 
 
 class AbstractAiSettings(ABC):
-    def __init__(self):
-        self.embeddings_max_concurrent_batches = None
+    def __init__(self, embeddings_max_concurrent_batches: int = None):
+        self.embeddings_max_concurrent_batches = embeddings_max_concurrent_batches
 
     @classmethod
     @abstractmethod

ravendb/documents/operations/ai/azure_open_ai_settings.py

@@ -12,8 +12,9 @@ def __init__(
         deployment_name: str = None,
         dimensions: int = None,
         temperature: float = None,
+        embeddings_max_concurrent_batches: int = None,
     ):
-        super().__init__(api_key, endpoint, model, dimensions, temperature)
+        super().__init__(api_key, endpoint, model, dimensions, temperature, embeddings_max_concurrent_batches)
         if deployment_name is None:
             raise ValueError("deployment_name cannot be None")
         self.deployment_name = deployment_name
@@ -27,6 +28,9 @@ def from_json(cls, json_dict: Dict[str, Any]) -> "AzureOpenAiSettings":
             dimensions=json_dict["Dimensions"] if "Dimensions" in json_dict else None,
             temperature=json_dict["Temperature"] if "Temperature" in json_dict else None,
             deployment_name=json_dict["DeploymentName"] if "DeploymentName" in json_dict else None,
+            embeddings_max_concurrent_batches=(
+                json_dict["EmbeddingsMaxConcurrentBatches"] if "EmbeddingsMaxConcurrentBatches" in json_dict else None
+            ),
         )
 
     def to_json(self) -> Dict[str, Any]:
@@ -37,4 +41,5 @@ def to_json(self) -> Dict[str, Any]:
             "Dimensions": self.dimensions,
             "Temperature": self.temperature,
             "DeploymentName": self.deployment_name,
+            "EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches,
         }

ravendb/documents/operations/ai/embedded_settings.py

@@ -4,12 +4,14 @@
 
 
 class EmbeddedSettings(AbstractAiSettings):
-    def __init__(self):
-        super().__init__()
+    def __init__(self, embeddings_max_concurrent_batches: int = None):
+        super().__init__(embeddings_max_concurrent_batches)
 
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> "EmbeddedSettings":
-        return cls()
+        return cls(
+            embeddings_max_concurrent_batches=json_dict.get("EmbeddingsMaxConcurrentBatches") if json_dict.get("EmbeddingsMaxConcurrentBatches") else None,
+        )
 
     def to_json(self) -> Dict[str, Any]:
         return {"EmbeddingsMaxConcurrentBatches": self.embeddings_max_concurrent_batches}

ravendb/documents/operations/ai/google_settings.py

@@ -11,9 +11,14 @@ class GoogleAiVersion(Enum):
 
 class GoogleSettings(AbstractAiSettings):
     def __init__(
-        self, model: str = None, api_key: str = None, ai_version: GoogleAiVersion = None, dimensions: int = None
+        self,
+        model: str = None,
+        api_key: str = None,
+        ai_version: GoogleAiVersion = None,
+        dimensions: int = None,
+        embeddings_max_concurrent_batches: int = None,
     ):
-        super().__init__()
+        super().__init__(embeddings_max_concurrent_batches)
         self.model = model
         self.api_key = api_key
         self.ai_version = ai_version
@@ -22,10 +27,13 @@ def __init__(
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> "GoogleSettings":
         return cls(
-            model=json_dict["Model"],
-            api_key=json_dict["ApiKey"],
-            ai_version=GoogleAiVersion(json_dict["AiVersion"]),
-            dimensions=json_dict["Dimensions"],
+            model=json_dict["Model"] if "Model" in json_dict else None,
+            api_key=json_dict["ApiKey"] if "ApiKey" in json_dict else None,
+            ai_version=GoogleAiVersion(json_dict["AiVersion"]) if "AiVersion" in json_dict else None,
+            dimensions=json_dict["Dimensions"] if "Dimensions" in json_dict else None,
+            embeddings_max_concurrent_batches=(
+                json_dict["EmbeddingsMaxConcurrentBatches"] if "EmbeddingsMaxConcurrentBatches" in json_dict else None
+            ),
         )
 
     def to_json(self) -> Dict[str, Any]:

ravendb/documents/operations/ai/hugging_face_settings.py

@@ -4,18 +4,27 @@
 
 
 class HuggingFaceSettings(AbstractAiSettings):
-    def __init__(self, api_key: str = None, model: str = None, endpoint: str = None):
-        super().__init__()
+    def __init__(
+        self,
+        api_key: str = None,
+        model: str = None,
+        endpoint: str = None,
+        embeddings_max_concurrent_batches: int = None,
+    ):
+        super().__init__(embeddings_max_concurrent_batches)
         self.api_key = api_key
         self.model = model
         self.endpoint = endpoint
 
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> "HuggingFaceSettings":
         return cls(
-            api_key=json_dict["ApiKey"],
-            model=json_dict["Model"],
-            endpoint=json_dict["Endpoint"],
+            api_key=json_dict["ApiKey"] if "ApiKey" in json_dict else None,
+            model=json_dict["Model"] if "Model" in json_dict else None,
+            endpoint=json_dict["Endpoint"] if "Endpoint" in json_dict else None,
+            embeddings_max_concurrent_batches=(
+                json_dict["EmbeddingsMaxConcurrentBatches"] if "EmbeddingsMaxConcurrentBatches" in json_dict else None
+            ),
         )
 
     def to_json(self) -> Dict[str, Any]:

ravendb/documents/operations/ai/mistral_ai_settings.py

@@ -4,18 +4,27 @@
 
 
 class MistralAiSettings(AbstractAiSettings):
-    def __init__(self, api_key: str = None, model: str = None, endpoint: str = None):
-        super().__init__()
+    def __init__(
+        self,
+        api_key: str = None,
+        model: str = None,
+        endpoint: str = None,
+        embeddings_max_concurrent_batches: int = None,
+    ):
+        super().__init__(embeddings_max_concurrent_batches)
         self.api_key = api_key
         self.model = model
         self.endpoint = endpoint
 
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> "MistralAiSettings":
         return cls(
-            api_key=json_dict["ApiKey"],
-            model=json_dict["Model"],
-            endpoint=json_dict["Endpoint"],
+            api_key=json_dict["ApiKey"] if "ApiKey" in json_dict else None,
+            model=json_dict["Model"] if "Model" in json_dict else None,
+            endpoint=json_dict["Endpoint"] if "Endpoint" in json_dict else None,
+            embeddings_max_concurrent_batches=(
+                json_dict["EmbeddingsMaxConcurrentBatches"] if "EmbeddingsMaxConcurrentBatches" in json_dict else None
+            ),
         )
 
     def to_json(self) -> Dict[str, Any]:

ravendb/documents/operations/ai/ollama_settings.py

@@ -12,12 +12,11 @@ def __init__(
         temperature: float = None,
         embeddings_max_concurrent_batches: int = None,
     ):
-        super().__init__()
+        super().__init__(embeddings_max_concurrent_batches)
         self.uri = uri
         self.model = model
         self.think = think
         self.temperature = temperature
-        self.embeddings_max_concurrent_batches = embeddings_max_concurrent_batches
 
     @classmethod
     def from_json(cls, json_dict: Dict[str, Any]) -> "OllamaSettings":

ravendb/documents/operations/ai/open_ai_base_settings.py

@@ -11,8 +11,9 @@ def __init__(
         model: str = None,
         dimensions: int = None,
         temperature: float = None,
+        embeddings_max_concurrent_batches: int = None,
     ):
-        super().__init__()
+        super().__init__(embeddings_max_concurrent_batches)
         self.api_key = api_key
         self.endpoint = endpoint
         self.model = model

ravendb/documents/operations/ai/open_ai_settings.py

@@ -13,8 +13,9 @@ def __init__(
         project_id: str = None,
         dimensions: int = None,
         temperature: float = None,
+        embeddings_max_concurrent_batches: int = None,
     ):
-        super().__init__(api_key, endpoint, model, dimensions, temperature)
+        super().__init__(api_key, endpoint, model, dimensions, temperature, embeddings_max_concurrent_batches)
         self.organization_id = organization_id
         self.project_id = project_id
 
@@ -28,6 +29,9 @@ def from_json(cls, json_dict: Dict[str, Any]) -> "OpenAiSettings":
             temperature=json_dict["Temperature"] if "Temperature" in json_dict else None,
             organization_id=json_dict["OrganizationId"] if "OrganizationId" in json_dict else None,
             project_id=json_dict["ProjectId"] if "ProjectId" in json_dict else None,
+            embeddings_max_concurrent_batches=(
+                json_dict["EmbeddingsMaxConcurrentBatches"] if "EmbeddingsMaxConcurrentBatches" in json_dict else None
+            ),
         )
 
     def to_json(self) -> Dict[str, Any]:

ravendb/documents/operations/ai/vertex_settings.py

@@ -16,8 +16,9 @@ def __init__(
         google_credentials_json: Optional[str] = None,
         location: Optional[str] = None,
         ai_version: Optional[VertexAIVersion] = None,
+        embeddings_max_concurrent_batches: Optional[int] = None,
     ):
-        super().__init__()
+        super().__init__(embeddings_max_concurrent_batches)
         self.model = model
         self.google_credentials_json = google_credentials_json
         self.location = location
@@ -30,6 +31,7 @@ def from_json(cls, json_dict: Dict[str, Any]) -> "VertexSettings":
             google_credentials_json=json_dict.get("GoogleCredentialsJson"),
             location=json_dict.get("Location"),
             ai_version=VertexAIVersion(json_dict["AiVersion"]) if json_dict.get("AiVersion") else None,
+            embeddings_max_concurrent_batches=json_dict.get("EmbeddingsMaxConcurrentBatches") if json_dict.get("EmbeddingsMaxConcurrentBatches") else None,
         )
 
     def to_json(self) -> Dict[str, Any]: