databricks-switch-plugin 0.1.0rc3__tar.gz → 0.1.1.dev2__tar.gz

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-switch-plugin
-Version: 0.1.0rc3
+Version: 0.1.1.dev2
 Summary: LLM-powered tool to convert SQL, code, and workflow files into Databricks notebooks.
 Project-URL: Documentation, https://databrickslabs.github.io/lakebridge
 Project-URL: Issues, https://github.com/databrickslabs/lakebridge/issues
@@ -57,6 +57,8 @@ Description-Content-Type: text/markdown
 # Switch
 LLM-Powered Code Conversion Plugin for Lakebridge
 
+[![codecov](https://codecov.io/gh/databrickslabs/switch/branch/main/graph/badge.svg?token=YOUR_TOKEN)](https://codecov.io/gh/databrickslabs/switch)
+
 ## Project Description
 Switch is a Lakebridge transpiler plugin that transforms SQL and other source formats into Databricks-compatible notebooks using Large Language Models (LLMs). As a core component of the Lakebridge migration platform, Switch provides automated code conversion capabilities through a multi-stage processing pipeline designed for large-scale platform migrations.
 
@@ -90,14 +92,14 @@ For advanced testing or direct control, you can deploy Switch directly to a Data
 
 **Cell 1: Install Switch package**
 
-From PyPI (stable):
+From PyPI (stable version):
 ```python
 %pip install databricks-switch-plugin
 ```
 
-Or from Test PyPI (development/preview versions):
+For prerelease versions (dev/rc):
 ```python
-%pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple databricks-switch-plugin
+%pip install --pre databricks-switch-plugin
 ```
 
 **Cell 2: Deploy Switch to workspace**

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/README.md

@@ -1,6 +1,8 @@
 # Switch
 LLM-Powered Code Conversion Plugin for Lakebridge
 
+[![codecov](https://codecov.io/gh/databrickslabs/switch/branch/main/graph/badge.svg?token=YOUR_TOKEN)](https://codecov.io/gh/databrickslabs/switch)
+
 ## Project Description
 Switch is a Lakebridge transpiler plugin that transforms SQL and other source formats into Databricks-compatible notebooks using Large Language Models (LLMs). As a core component of the Lakebridge migration platform, Switch provides automated code conversion capabilities through a multi-stage processing pipeline designed for large-scale platform migrations.
 
@@ -34,14 +36,14 @@ For advanced testing or direct control, you can deploy Switch directly to a Data
 
 **Cell 1: Install Switch package**
 
-From PyPI (stable):
+From PyPI (stable version):
 ```python
 %pip install databricks-switch-plugin
 ```
 
-Or from Test PyPI (development/preview versions):
+For prerelease versions (dev/rc):
 ```python
-%pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple databricks-switch-plugin
+%pip install --pre databricks-switch-plugin
 ```
 
 **Cell 2: Deploy Switch to workspace**

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databricks-switch-plugin"
-version = "0.1.0rc3"
+version = "0.1.1dev2"
 description = "LLM-powered tool to convert SQL, code, and workflow files into Databricks notebooks."
 license-files = { paths = ["LICENSE", "NOTICE"] }
 keywords = ["Databricks", "SQL", "Migration", "LLM", "Conversion"]
@@ -326,4 +326,20 @@ allow-global-unused-variables = true
 callbacks = ["cb_", "_cb"]
 dummy-variables-rgx = "_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_"
 ignored-argument-names = "_.*|^ignored_|^unused_"
-redefining-builtins-modules = ["six.moves", "past.builtins", "future.builtins", "builtins", "io"]
+redefining-builtins-modules = ["six.moves", "past.builtins", "future.builtins", "builtins", "io"]
+
+[tool.coverage.run]
+source = ["switch.api", "switch.notebooks.pyscripts"]
+omit = [
+    "*/tests/*",
+    "switch/notebooks/*.py",
+    "*/test_*.py"
+]
+
+[tool.coverage.report]
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "raise AssertionError",
+    "raise NotImplementedError"
+]

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/switch/__init__.py

@@ -1,3 +1,3 @@
 from importlib.metadata import version
 
-__version__ = "0.1.0rc3"
+__version__ = "0.1.1dev2"

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/switch/api/__init__.py

@@ -1,4 +1,5 @@
 """Switch API module for Lakebridge integration"""
+
 from .job_runner import SwitchJobRunner
 from .installer import SwitchInstaller
 

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/switch/api/installer.py

@@ -1,4 +1,5 @@
 """Switch installer for Databricks workspace deployment"""
+
 import logging
 import os
 from dataclasses import dataclass
@@ -16,6 +17,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class InstallResult:
     """Result of installation"""
+
     job_id: int
     job_name: str
     job_url: str
@@ -26,6 +28,7 @@ class InstallResult:
 @dataclass
 class UninstallResult:
     """Result of uninstallation"""
+
     success: bool
     message: str = ""
 
@@ -37,6 +40,7 @@ class SwitchInstaller:
     - install(): Install Switch (optionally removes previous installation)
     - uninstall(): Remove Switch installation
     """
+
     # Directory and path constants
     _WORKSPACE_BASE_DIR_NAME = ".lakebridge-switch"
     _MAIN_NOTEBOOK_PATH = "switch/notebooks/00_main"
@@ -58,10 +62,12 @@ class SwitchInstaller:
         """
         self.workspace_client = workspace_client
 
-    def install(self,
-                default_parameters: Optional[SwitchJobParameters] = None,
-                previous_job_id: Optional[int] = None,
-                previous_switch_home: Optional[str] = None) -> InstallResult:
+    def install(
+        self,
+        default_parameters: Optional[SwitchJobParameters] = None,
+        previous_job_id: Optional[int] = None,
+        previous_switch_home: Optional[str] = None,
+    ) -> InstallResult:
         """Install Switch for current user. Optionally cleans up previous installation.
 
         Args:
@@ -77,13 +83,12 @@ class SwitchInstaller:
         """
         # Clean up previous installation if specified
         if previous_job_id or previous_switch_home:
-            logger.info(f"Cleaning up previous installation: job_id={previous_job_id}, switch_home={previous_switch_home}")
-
-            uninstall_result = self.uninstall(
-                job_id=previous_job_id,
-                switch_home=previous_switch_home
+            logger.info(
+                f"Cleaning up previous installation: job_id={previous_job_id}, switch_home={previous_switch_home}"
             )
 
+            uninstall_result = self.uninstall(job_id=previous_job_id, switch_home=previous_switch_home)
+
             if uninstall_result.success:
                 logger.info("Successfully cleaned up previous installation")
             else:
@@ -104,16 +109,10 @@ class SwitchInstaller:
         logger.info(f"Successfully installed Switch: job_id={job_id}, job_name={self.JOB_NAME}")
 
         return InstallResult(
-            job_id=job_id,
-            job_name=self.JOB_NAME,
-            job_url=job_url,
-            switch_home=switch_home,
-            created_by=created_by
+            job_id=job_id, job_name=self.JOB_NAME, job_url=job_url, switch_home=switch_home, created_by=created_by
         )
 
-    def uninstall(self,
-                  job_id: Optional[int] = None,
-                  switch_home: Optional[str] = None) -> UninstallResult:
+    def uninstall(self, job_id: Optional[int] = None, switch_home: Optional[str] = None) -> UninstallResult:
         """Uninstall Switch
 
         Args:
@@ -153,10 +152,7 @@ class SwitchInstaller:
         else:
             message = f"Failed to uninstall Switch: {'; '.join(errors)}"
 
-        return UninstallResult(
-            success=success,
-            message=message
-        )
+        return UninstallResult(success=success, message=message)
 
     def _deploy_switch_package(self) -> str:
         """Deploy Switch package to workspace for both PyPI and development installations.
@@ -190,6 +186,7 @@ class SwitchInstaller:
 
         # Get site-packages directory for PyPI installations
         import switch
+
         site_packages_dir = os.path.dirname(os.path.dirname(switch.__file__))
 
         # Define search locations: PyPI install location first, then development location
@@ -206,9 +203,7 @@ class SwitchInstaller:
         for root, _, files in os.walk(switch_package_dir):
             for file_name in files:
                 # Skip unwanted files
-                if (file_name.startswith('.') or 
-                    file_name.endswith(('.pyc', '.pyo')) or
-                    file_name == '__pycache__'):
+                if file_name.startswith('.') or file_name.endswith(('.pyc', '.pyo')) or file_name == '__pycache__':
                     continue
 
                 local_path = os.path.join(root, file_name)
@@ -225,10 +220,7 @@ class SwitchInstaller:
                 with open(local_path, 'rb') as f:
                     content = f.read()
                 self.workspace_client.workspace.upload(
-                    path=remote_path,
-                    content=content,
-                    format=ImportFormat.AUTO,
-                    overwrite=True
+                    path=remote_path, content=content, format=ImportFormat.AUTO, overwrite=True
                 )
                 logger.debug(f"Uploaded {rel_path}")
 
@@ -263,10 +255,8 @@ class SwitchInstaller:
             task_key=self._MAIN_TASK_KEY,
             max_retries=self._MAIN_TASK_MAX_RETRIES,
             notebook_task=NotebookTask(
-                notebook_path=notebook_path,
-                base_parameters=template_params,
-                source=Source.WORKSPACE
-            )
+                notebook_path=notebook_path, base_parameters=template_params, source=Source.WORKSPACE
+            ),
         )
 
         # Create the job with created_by tag
@@ -274,10 +264,7 @@ class SwitchInstaller:
         tags = {self.JOB_TAG_CREATED_BY: current_user}
 
         response = self.workspace_client.jobs.create(
-            name=self.JOB_NAME,
-            tasks=[task],
-            max_concurrent_runs=self._JOB_MAX_CONCURRENT_RUNS,
-            tags=tags
+            name=self.JOB_NAME, tasks=[task], max_concurrent_runs=self._JOB_MAX_CONCURRENT_RUNS, tags=tags
         )
         if response.job_id is None:
             raise ValueError("Failed to create job: job_id is None")
@@ -316,11 +303,7 @@ class SwitchInstaller:
                 logger.debug(f"Found {filename} at {file_path}, uploading to {dest_path}")
                 with open(file_path, 'rb') as f:
                     content = f.read()
-                self.workspace_client.workspace.upload(
-                    path=dest_path,
-                    content=content,
-                    format=ImportFormat.AUTO
-                )
+                self.workspace_client.workspace.upload(path=dest_path, content=content, format=ImportFormat.AUTO)
                 return True
 
         logger.debug(f"{filename} not found in any of the provided locations")

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/switch/api/job_parameters.py

@@ -1,4 +1,5 @@
 """Switch job parameters and related definitions"""
+
 import dataclasses
 import json
 from dataclasses import dataclass
@@ -23,27 +24,28 @@ class SwitchJobParameters:
     2. Template parameters for SwitchInstaller
     3. Default value management
     """
+
     # Basic settings
     input_dir: Optional[str] = None
     output_dir: Optional[str] = None
     result_catalog: Optional[str] = None
     result_schema: Optional[str] = None
-    builtin_prompt: Optional[BuiltinPrompt] = None
+    builtin_prompt: Optional[BuiltinPrompt | str] = None
 
     # Conversion settings
-    source_format: Optional[SourceFormat] = None
-    target_type: Optional[TargetType] = None
+    source_format: Optional[SourceFormat | str] = None
+    target_type: Optional[TargetType | str] = None
     output_extension: Optional[str] = None
 
     # Execution settings
     endpoint_name: Optional[str] = None
     concurrency: Optional[int] = None
     max_fix_attempts: Optional[int] = None
-    log_level: Optional[LogLevel] = None
+    log_level: Optional[LogLevel | str] = None
 
     # Advanced settings
     token_count_threshold: Optional[int] = None
-    comment_lang: Optional[CommentLanguage] = None
+    comment_lang: Optional[CommentLanguage | str] = None
     conversion_prompt_yaml: Optional[str] = None
     sql_output_dir: Optional[str] = None
 
@@ -99,7 +101,7 @@ class SwitchJobParameters:
         """Validate parameters
 
         Args:
-            require_all: If True, validate all required parameters. 
+            require_all: If True, validate all required parameters.
                         If False, only validate provided parameters.
 
         Raises:
@@ -118,29 +120,19 @@ class SwitchJobParameters:
 
             # Either builtin_prompt or conversion_prompt_yaml must be specified
             if not self.builtin_prompt and not self.conversion_prompt_yaml:
-                raise ValueError(
-                    "Either builtin_prompt or conversion_prompt_yaml must be specified"
-                )
+                raise ValueError("Either builtin_prompt or conversion_prompt_yaml must be specified")
 
         # Validate source format if provided
         if self.source_format and self.source_format not in SourceFormat:
-            raise ValueError(
-                f"source_format must be one of {list(SourceFormat)}, "
-                f"got '{self.source_format}'"
-            )
+            raise ValueError(f"source_format must be one of {list(SourceFormat)}, " f"got '{self.source_format}'")
 
         # Validate target type if provided
         if self.target_type and self.target_type not in TargetType:
-            raise ValueError(
-                f"target_type must be one of {list(TargetType)}, "
-                f"got '{self.target_type}'"
-            )
+            raise ValueError(f"target_type must be one of {list(TargetType)}, " f"got '{self.target_type}'")
 
         # Validate output_extension for file target type
         if self.target_type == TargetType.FILE and not self.output_extension:
-            raise ValueError(
-                "output_extension is required when target_type is 'file'"
-            )
+            raise ValueError("output_extension is required when target_type is 'file'")
 
         # Validate numeric parameters if they're not None
         if self.token_count_threshold is not None and self.token_count_threshold <= 0:
@@ -163,7 +155,13 @@ class SwitchJobParameters:
         required_fields = {"input_dir", "output_dir", "result_catalog", "result_schema"}
 
         # Define fields that should always be empty in template
-        template_empty_fields = {"builtin_prompt", "conversion_prompt_yaml", "output_extension", "request_params", "sql_output_dir"}
+        template_empty_fields = {
+            "builtin_prompt",
+            "conversion_prompt_yaml",
+            "output_extension",
+            "request_params",
+            "sql_output_dir",
+        }
 
         # Process all dataclass fields dynamically
         template_params = {}
@@ -199,25 +197,25 @@ class SwitchJobParameters:
 
     def _serialize_field_value(self, field_name: str, value: Any) -> str:
         """Serialize field value to string for job parameters
-        
+
         Args:
             field_name: Name of the dataclass field
             value: Value to serialize
-            
+
         Returns:
             String representation suitable for Databricks job parameters
         """
         if value is None:
             return ""
-        
+
         # Handle enum types
         if hasattr(value, 'value'):
             return value.value
-        
+
         # Handle request_params special case (JSON serialization)
         if field_name == 'request_params' and isinstance(value, dict):
             return json.dumps(value)
-        
+
         # Handle all other types as string
         return str(value)
 
@@ -234,9 +232,7 @@ def _load_switch_config_defaults() -> Dict[str, Any]:
     config_path = Path(__file__).parent.parent.parent / "lsp" / "config.yml"
 
     if not config_path.exists():
-        raise ValueError(
-            f"Switch configuration file not found at {config_path}"
-        )
+        raise ValueError(f"Switch configuration file not found at {config_path}")
 
     with open(config_path, 'r') as f:
         config_data = yaml.safe_load(f)

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/switch/api/job_runner.py

@@ -1,4 +1,5 @@
 """Switch job runner for SQL conversion operations"""
+
 import logging
 
 from databricks.sdk import WorkspaceClient
@@ -39,10 +40,7 @@ class SwitchJobRunner:
 
         # Run the job with parameters
         try:
-            response = self.workspace_client.jobs.run_now(
-                job_id=self.job_id,
-                notebook_params=job_params
-            )
+            response = self.workspace_client.jobs.run_now(job_id=self.job_id, notebook_params=job_params)
             run_id = response.run_id
             logger.info(f"Started job run {run_id} for job {self.job_id}")
             return run_id
@@ -69,12 +67,11 @@ class SwitchJobRunner:
         # Run the job synchronously and wait for completion
         try:
             from datetime import timedelta
+
             timeout_delta = timedelta(seconds=timeout_seconds)
 
             run = self.workspace_client.jobs.run_now_and_wait(
-                job_id=self.job_id,
-                notebook_params=job_params,
-                timeout=timeout_delta
+                job_id=self.job_id, notebook_params=job_params, timeout=timeout_delta
             )
             if run.state and run.state.life_cycle_state:
                 logger.info(f"Job run {run.run_id} completed with state: {run.state.life_cycle_state}")

{databricks_switch_plugin-0.1.0rc3 → databricks_switch_plugin-0.1.1.dev2}/switch/notebooks/pyscripts/helpers/batch_inference_helper.py

@@ -6,6 +6,7 @@ Key components:
 - BatchInferenceManager: Manages concurrent processing of multiple requests
 - BatchInferenceRequest/Response: Data structures for request/response handling
 """
+
 import asyncio
 import json
 import logging
@@ -16,9 +17,14 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import httpx
 from httpx import codes
-from tenacity import (RetryCallState, RetryError, retry,
-                      retry_if_exception_type, stop_after_attempt,
-                      wait_random_exponential)
+from tenacity import (
+    RetryCallState,
+    RetryError,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_random_exponential,
+)
 
 from ..utils.databricks_credentials import DatabricksCredentials
 from ..utils.common_utils import setup_logger
@@ -50,6 +56,7 @@ class BatchInferenceRequest:
         system_message (str): The system message to guide the model's behavior.
         few_shots (Optional[List[Dict[str, str]]]): Optional few-shot examples for the model.
     """
+
     index: int
     text: str
     system_message: str
@@ -66,6 +73,7 @@ class TokenUsage:
         completion_tokens (int): Number of tokens used for the completion/response.
         total_tokens (int): Total number of tokens used (prompt + completion).
     """
+
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
@@ -83,6 +91,7 @@ class BatchInferenceResponse:
         error (Optional[str]): Any error message, if an error occurred during processing.
         processing_time_seconds (Optional[float]): Total time in seconds spent processing this request.
     """
+
     index: int
     content: Optional[str]
     token_usage: Optional[TokenUsage]
@@ -137,17 +146,16 @@ class BatchInferenceManager:
         counter = AsyncCounter()
         start_time = time.perf_counter()
 
-        tasks = [self._process_inference(request, semaphore, counter, start_time)
-                 for request in requests]
+        tasks = [self._process_inference(request, semaphore, counter, start_time) for request in requests]
         responses = await asyncio.gather(*tasks)
         await self.client.close()
 
         self.logger.info(f"Completed batch inference for {len(requests)} requests")
         return responses
 
-    async def _process_inference(self, request: BatchInferenceRequest,
-                                 semaphore: asyncio.Semaphore, counter: 'AsyncCounter',
-                                 start_time: float) -> BatchInferenceResponse:
+    async def _process_inference(
+        self, request: BatchInferenceRequest, semaphore: asyncio.Semaphore, counter: 'AsyncCounter', start_time: float
+    ) -> BatchInferenceResponse:
         """
         Process inference for a single request.
 
@@ -171,10 +179,11 @@ class BatchInferenceManager:
                     content=content,
                     token_usage=token_usage,
                     error=None,
-                    processing_time_seconds=processing_time_seconds
+                    processing_time_seconds=processing_time_seconds,
                 )
                 self.logger.info(
-                    f"Completed inference for index {request.index} in {processing_time_seconds:.2f} seconds")
+                    f"Completed inference for index {request.index} in {processing_time_seconds:.2f} seconds"
+                )
             except Exception as e:
                 response = await self._handle_error(e, request.index)
 
@@ -226,11 +235,7 @@ class BatchInferenceManager:
             self.client._log_general_error(e, request_index, logger=self.logger)
 
         return BatchInferenceResponse(
-            index=request_index,
-            content=None,
-            token_usage=None,
-            error=error_message,
-            processing_time_seconds=None
+            index=request_index, content=None, token_usage=None, error=error_message, processing_time_seconds=None
         )
 
 
@@ -292,6 +297,7 @@ class AsyncChatClient:
             httpx.HTTPStatusError: If an HTTP error occurs that can't be resolved by retrying.
             Exception: For any other unexpected errors.
         """
+
         @retry(
             retry=retry_if_exception_type(httpx.HTTPStatusError),
             stop=lambda rs: self._get_stop_condition(rs),
@@ -320,11 +326,7 @@ class AsyncChatClient:
                     request_body = {"messages": messages}
                     if self.request_params:
                         request_body.update(self.request_params)
-                    request_args = {
-                        'url': url,
-                        'headers': headers,
-                        'json': request_body
-                    }
+                    request_args = {'url': url, 'headers': headers, 'json': request_body}
                     self._log_request_args(request_args, request.index, level=logging.INFO)
 
                     # Send the request
@@ -365,13 +367,16 @@ class AsyncChatClient:
                     # If the response indicates that the finished reason is "length",
                     # we need to continue the conversation by appending the last message
                     messages.append({"role": "assistant", "content": content})
-                    messages.append({"role": "user", "content": f"The previous response ended with: '{content[-50:]}'. "
-                                     f"Please continue exactly from this point without repeating any content."})
+                    messages.append(
+                        {
+                            "role": "user",
+                            "content": f"The previous response ended with: '{content[-50:]}'. "
+                            f"Please continue exactly from this point without repeating any content.",
+                        }
+                    )
 
                 return total_content, TokenUsage(
-                    prompt_tokens=prompt_tokens,
-                    completion_tokens=completion_tokens,
-                    total_tokens=total_tokens
+                    prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens
                 )
 
             # Handle HTTP errors (e.g., 4xx, 5xx)
@@ -427,7 +432,7 @@ class AsyncChatClient:
 
         Args:
             content: The content field from the API response, which could be either
-                    a string (standard format) or a list of structured objects 
+                    a string (standard format) or a list of structured objects
                     (Extended thinking mode with reasoning and text objects).
 
         Returns:
@@ -441,9 +446,7 @@ class AsyncChatClient:
             return content
         elif isinstance(content, list):
             return "".join(
-                content_part.get("text", "")
-                for content_part in content
-                if content_part.get("type") == "text"
+                content_part.get("text", "") for content_part in content if content_part.get("type") == "text"
             )
         raise TypeError(f"Unexpected content format: {type(content).__name__}")
 
@@ -519,7 +522,7 @@ class AsyncChatClient:
                         "required parameter",
                         "missing parameter",
                         "invalid format",
-                        "not a valid"
+                        "not a valid",
                     ]
 
                     for pattern in validation_patterns:
@@ -553,8 +556,7 @@ class AsyncChatClient:
         This is a temporary workaround for Databricks Claude endpoints
         that may return 400 instead of 429 for certain size-related limitations.
         """
-        if (error.response.status_code == codes.BAD_REQUEST and
-                "databricks-claude" in self.endpoint_name):
+        if error.response.status_code == codes.BAD_REQUEST and "databricks-claude" in self.endpoint_name:
             try:
                 response_data = error.response.json()
                 if response_data.get("message") is None:
@@ -567,7 +569,9 @@ class AsyncChatClient:
                 pass
         return False
 
-    def _log_request_args(self, request_args: Dict[str, Any], request_index: int, max_head: int = 200, max_tail=100, level=logging.INFO) -> None:
+    def _log_request_args(
+        self, request_args: Dict[str, Any], request_index: int, max_head: int = 200, max_tail=100, level=logging.INFO
+    ) -> None:
         """
         Log the request arguments for debugging purposes.
 
@@ -585,10 +589,17 @@ class AsyncChatClient:
             level,
             f"Request for index {request_index} - Request details: "
             f"Size: {body_len} chars, "
-            f"Content: {body_str[:max_head]}... [truncated] ...{body_str[-max_tail:]}"
+            f"Content: {body_str[:max_head]}... [truncated] ...{body_str[-max_tail:]}",
         )
 
-    def _log_response_details(self, response: httpx.Response, request_index: int, max_head: int = 200, max_tail: int = 100, level=logging.DEBUG) -> None:
+    def _log_response_details(
+        self,
+        response: httpx.Response,
+        request_index: int,
+        max_head: int = 200,
+        max_tail: int = 100,
+        level=logging.DEBUG,
+    ) -> None:
         """
         Log the detailed response data at DEBUG level, with truncation.
 
@@ -615,7 +626,7 @@ class AsyncChatClient:
             f"Request for index {request_index} - Response details ({content_type}): "
             f"Status: {response.status_code}, "
             f"Headers: {dict(response.headers)}, "
-            f"Body: {body_content}"
+            f"Body: {body_content}",
         )
 
     def _truncate_long_strings_in_dict(self, data, max_length: int = 2000):
@@ -635,7 +646,11 @@ class AsyncChatClient:
             return [self._truncate_long_strings_in_dict(item, max_length) for item in data]
         elif isinstance(data, str) and len(data) > max_length:
             # Truncate long string values
-            return data[:max_length // 2] + f" ... [truncated {len(data) - max_length} chars] ... " + data[-max_length // 2:]
+            return (
+                data[: max_length // 2]
+                + f" ... [truncated {len(data) - max_length} chars] ... "
+                + data[-max_length // 2 :]
+            )
         else:
             return data