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/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]: