Merge pull request #173 from Aleph-Alpha/PHS-622-expose-pool-size-in-…

…constructor

feat: Client exposes HTTPAdapter pool size

aleph_alpha_client/aleph_alpha_client.py

@@ -167,6 +167,7 @@ def __init__(
         nice: bool = False,
         verify_ssl=True,
         tags: Optional[Sequence[str]] = None,
+        pool_size: int = 10,
     ) -> None:
         if host[-1] != "/":
             host += "/"
@@ -184,7 +185,11 @@ def __init__(
             allowed_methods=["POST", "GET"],
             raise_on_status=False,
         )
-        adapter = HTTPAdapter(max_retries=retry_strategy)
+        adapter = HTTPAdapter(
+            max_retries=retry_strategy,
+            pool_connections=pool_size,
+            pool_maxsize=pool_size,
+        )
         self.session = requests.Session()
         self.session.verify = verify_ssl
         self.session.headers = CaseInsensitiveDict(

aleph_alpha_client/completion.py

@@ -284,9 +284,11 @@ def from_json(json: Dict[str, Any]) -> "CompletionResponse":
             ],
             num_tokens_prompt_total=json["num_tokens_prompt_total"],
             num_tokens_generated=json["num_tokens_generated"],
-            optimized_prompt=Prompt.from_json(optimized_prompt_json)
-            if optimized_prompt_json
-            else None,
+            optimized_prompt=(
+                Prompt.from_json(optimized_prompt_json)
+                if optimized_prompt_json
+                else None
+            ),
         )
 
     def to_json(self) -> Mapping[str, Any]:

aleph_alpha_client/explanation.py

@@ -172,9 +172,9 @@ class ExplanationRequest:
     control_factor: Optional[float] = None
     control_token_overlap: Optional[ControlTokenOverlap] = None
     control_log_additive: Optional[bool] = None
-    prompt_granularity: Optional[
-        Union[PromptGranularity, str, CustomGranularity]
-    ] = None
+    prompt_granularity: Optional[Union[PromptGranularity, str, CustomGranularity]] = (
+        None
+    )
     target_granularity: Optional[TargetGranularity] = None
     postprocessing: Optional[ExplanationPostprocessing] = None
     normalize: Optional[bool] = None
@@ -357,9 +357,11 @@ def from_json(item: Dict[str, Any]) -> "TextPromptItemExplanation":
     def with_text(self, prompt: Text) -> "TextPromptItemExplanation":
         return TextPromptItemExplanation(
             scores=[
-                TextScoreWithRaw.from_text_score(score, prompt)
-                if isinstance(score, TextScore)
-                else score
+                (
+                    TextScoreWithRaw.from_text_score(score, prompt)
+                    if isinstance(score, TextScore)
+                    else score
+                )
                 for score in self.scores
             ]
         )
@@ -386,9 +388,11 @@ def from_json(item: Dict[str, Any]) -> "TargetPromptItemExplanation":
     def with_text(self, prompt: str) -> "TargetPromptItemExplanation":
         return TargetPromptItemExplanation(
             scores=[
-                TargetScoreWithRaw.from_target_score(score, prompt)
-                if isinstance(score, TargetScore)
-                else score
+                (
+                    TargetScoreWithRaw.from_target_score(score, prompt)
+                    if isinstance(score, TargetScore)
+                    else score
+                )
                 for score in self.scores
             ]
         )
@@ -461,9 +465,11 @@ def with_image_prompt_items_in_pixels(self, prompt: Prompt) -> "Explanation":
         return Explanation(
             target=self.target,
             items=[
-                item.in_pixels(prompt.items[item_index])
-                if isinstance(item, ImagePromptItemExplanation)
-                else item
+                (
+                    item.in_pixels(prompt.items[item_index])
+                    if isinstance(item, ImagePromptItemExplanation)
+                    else item
+                )
                 for item_index, item in enumerate(self.items)
             ],
         )