Repository: stanfordnlp/dspy

Issue:
Increase test coverage for JsonAdapter
While we have unit tests for `JsonAdapter` in `tests/adapters/test_json_adapter.py`, it's only covering a few cases. We should add test cases to cover:

- Different input field and output field types, including pydantic models and nested types.
- The request sent to LM contains the signature information as expected. 
- On failure like parsing errors, an explicit error is raised.

Hints from issue comments:
@Y-1huadb: Hi @chenmoneygithub, 
I’ve submitted a [PR](https://github.com/stanfordnlp/dspy/pull/8131) that attempts to address this issue. When you have a moment, could you please take a look? Let me know if there’s anything I can improve or clarify. Thanks so much for your time!

Solution patch:
diff --git a/dspy/adapters/json_adapter.py b/dspy/adapters/json_adapter.py
index 4c6c868611..14c5230eb4 100644
--- a/dspy/adapters/json_adapter.py
+++ b/dspy/adapters/json_adapter.py
@@ -76,18 +76,19 @@ def format_field_structure(self, signature: Type[Signature]) -> str:
         parts = []
         parts.append("All interactions will be structured in the following way, with the appropriate values filled in.")
 
-        def format_signature_fields_for_instructions(fields: Dict[str, FieldInfo]):
+        def format_signature_fields_for_instructions(fields: Dict[str, FieldInfo], role: str):
             return self.format_field_with_value(
                 fields_with_values={
                     FieldInfoWithName(name=field_name, info=field_info): translate_field_type(field_name, field_info)
                     for field_name, field_info in fields.items()
                 },
+                role=role,
             )
 
         parts.append("Inputs will have the following structure:")
-        parts.append(format_signature_fields_for_instructions(signature.input_fields))
+        parts.append(format_signature_fields_for_instructions(signature.input_fields, role="user"))
         parts.append("Outputs will be a JSON object with the following fields.")
-        parts.append(format_signature_fields_for_instructions(signature.output_fields))
+        parts.append(format_signature_fields_for_instructions(signature.output_fields, role="assistant"))
         return "\n\n".join(parts).strip()
 
     def user_message_output_requirements(self, signature: Type[Signature]) -> str:
diff --git a/dspy/clients/__init__.py b/dspy/clients/__init__.py
index 3eeba1e2ca..52be745abb 100644
--- a/dspy/clients/__init__.py
+++ b/dspy/clients/__init__.py
@@ -1,15 +1,17 @@
-from dspy.clients.lm import LM
-from dspy.clients.provider import Provider, TrainingJob
-from dspy.clients.base_lm import BaseLM, inspect_history
-from dspy.clients.embedding import Embedder
-import litellm
+import logging
 import os
 from pathlib import Path
-from dspy.clients.cache import Cache
-import logging
 from typing import Optional
+
+import litellm
 from litellm.caching.caching import Cache as LitellmCache
 
+from dspy.clients.base_lm import BaseLM, inspect_history
+from dspy.clients.cache import Cache
+from dspy.clients.embedding import Embedder
+from dspy.clients.lm import LM
+from dspy.clients.provider import Provider, TrainingJob
+
 logger = logging.getLogger(__name__)
 
 DISK_CACHE_DIR = os.environ.get("DSPY_CACHEDIR") or os.path.join(Path.home(), ".dspy_cache")
diff --git a/dspy/clients/base_lm.py b/dspy/clients/base_lm.py
index 4c7d15cc3b..9080d4cbe4 100644
--- a/dspy/clients/base_lm.py
+++ b/dspy/clients/base_lm.py
@@ -1,6 +1,5 @@
 import datetime
 import uuid
-from abc import ABC
 
 from dspy.dsp.utils import settings
 from dspy.utils.callback import with_callbacks
@@ -9,7 +8,7 @@
 GLOBAL_HISTORY = []
 
 
-class BaseLM(ABC):
+class BaseLM:
     """Base class for handling LLM calls.
 
     Most users can directly use the `dspy.LM` class, which is a subclass of `BaseLM`. Users can also implement their
@@ -178,7 +177,10 @@ def _inspect_history(history, n: int = 1):
                             image_str = ""
                             if "base64" in c["image_url"].get("url", ""):
                                 len_base64 = len(c["image_url"]["url"].split("base64,")[1])
-                                image_str = f"<{c['image_url']['url'].split('base64,')[0]}base64,<IMAGE BASE 64 ENCODED({str(len_base64)})>"
+                                image_str = (
+                                    f"<{c['image_url']['url'].split('base64,')[0]}base64,"
+                                    f"<IMAGE BASE 64 ENCODED({len_base64!s})>"
+                                )
                             else:
                                 image_str = f"<image_url: {c['image_url']['url']}>"
                             print(_blue(image_str.strip()))
diff --git a/dspy/clients/cache.py b/dspy/clients/cache.py
index 9dff881363..d3bad5084c 100644
--- a/dspy/clients/cache.py
+++ b/dspy/clients/cache.py
@@ -83,7 +83,7 @@ def transform_value(value):
                 try:
                     # For regular functions, we can get the source code
                     return f"<callable_source:{inspect.getsource(value)}>"
-                except (TypeError, OSError, IOError):
+                except (TypeError, OSError):
                     # For lambda functions or other callables where source isn't available,
                     # use a string representation
                     return f"<callable:{value.__name__ if hasattr(value, '__name__') else 'lambda'}>"
@@ -170,7 +170,7 @@ def load_memory_cache(self, filepath: str) -> None:
 
 def request_cache(
     cache_arg_name: Optional[str] = None,
-    ignored_args_for_cache_key: Optional[list[str]] = ["api_key", "api_base", "base_url"],
+    ignored_args_for_cache_key: Optional[list[str]] = None,
     enable_memory_cache: bool = True,
     *,  # everything after this is keyword-only
     maxsize: Optional[int] = None,  # legacy / no-op
@@ -185,7 +185,7 @@ def request_cache(
         enable_memory_cache: Whether to enable in-memory cache at call time. If False, the memory cache will not be
             written to on new data.
     """
-
+    ignored_args_for_cache_key = ignored_args_for_cache_key or ["api_key", "api_base", "base_url"]
     # Deprecation notice
     if maxsize is not None:
         logger.warning(
diff --git a/dspy/clients/lm.py b/dspy/clients/lm.py
index f77bc40ecb..2aad71aaf3 100644
--- a/dspy/clients/lm.py
+++ b/dspy/clients/lm.py
@@ -7,7 +7,6 @@
 import litellm
 from anyio.streams.memory import MemoryObjectSendStream
 from asyncer import syncify
-from litellm import RetryPolicy
 
 import dspy
 from dspy.clients.cache import request_cache
@@ -85,9 +84,9 @@ def __init__(
 
         if model_pattern:
             # Handle OpenAI reasoning models (o1, o3)
-            assert max_tokens >= 20_000 and temperature == 1.0, (
-                "OpenAI's reasoning models require passing temperature=1.0 and max_tokens >= 20_000 to `dspy.LM(...)`"
-            )
+            assert (
+                max_tokens >= 20_000 and temperature == 1.0
+            ), "OpenAI's reasoning models require passing temperature=1.0 and max_tokens >= 20_000 to `dspy.LM(...)`"
             self.kwargs = dict(temperature=temperature, max_completion_tokens=max_tokens, **kwargs)
         else:
             self.kwargs = dict(temperature=temperature, max_tokens=max_tokens, **kwargs)
@@ -233,7 +232,6 @@ def dump_state(self):
 
 def _get_stream_completion_fn(
     request: Dict[str, Any],
-    retry_kwargs: Dict[str, Any],
     cache_kwargs: Dict[str, Any],
     sync=True,
 ):
@@ -247,11 +245,10 @@ def _get_stream_completion_fn(
     stream = cast(MemoryObjectSendStream, stream)
     caller_predict_id = id(caller_predict) if caller_predict else None
 
-    async def stream_completion(request: Dict[str, Any], retry_kwargs: Dict[str, Any], cache_kwargs: Dict[str, Any]):
+    async def stream_completion(request: Dict[str, Any], cache_kwargs: Dict[str, Any]):
         response = await litellm.acompletion(
             cache=cache_kwargs,
             stream=True,
-            **retry_kwargs,
             **request,
         )
         chunks = []
@@ -265,10 +262,10 @@ async def stream_completion(request: Dict[str, Any], retry_kwargs: Dict[str, Any
 
     def sync_stream_completion():
         syncified_stream_completion = syncify(stream_completion)
-        return syncified_stream_completion(request, retry_kwargs, cache_kwargs)
+        return syncified_stream_completion(request, cache_kwargs)
 
     async def async_stream_completion():
-        return await stream_completion(request, retry_kwargs, cache_kwargs)
+        return await stream_completion(request, cache_kwargs)
 
     if sync:
         return sync_stream_completion
@@ -277,20 +274,12 @@ async def async_stream_completion():
 
 
 def litellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
-    retry_kwargs = dict(
-        retry_policy=_get_litellm_retry_policy(num_retries),
-        retry_strategy="exponential_backoff_retry",
-        # In LiteLLM version 1.55.3 (the first version that supports retry_policy as an argument
-        # to completion()), the default value of max_retries is non-zero for certain providers, and
-        # max_retries is stacked on top of the retry_policy. To avoid this, we set max_retries=0
-        max_retries=0,
-    )
-
-    stream_completion = _get_stream_completion_fn(request, retry_kwargs, cache, sync=True)
+    stream_completion = _get_stream_completion_fn(request, cache, sync=True)
     if stream_completion is None:
         return litellm.completion(
             cache=cache,
-            **retry_kwargs,
+            num_retries=num_retries,
+            retry_strategy="exponential_backoff_retry",
             **request,
         )
 
@@ -316,30 +305,19 @@ def litellm_text_completion(request: Dict[str, Any], num_retries: int, cache={"n
         api_key=api_key,
         api_base=api_base,
         prompt=prompt,
-        retry_policy=_get_litellm_retry_policy(num_retries),
-        # In LiteLLM version 1.55.3 (the first version that supports retry_policy as an argument
-        # to completion()), the default value of max_retries is non-zero for certain providers, and
-        # max_retries is stacked on top of the retry_policy. To avoid this, we set max_retries=0
-        max_retries=0,
+        num_retries=num_retries,
+        retry_strategy="exponential_backoff_retry",
         **request,
     )
 
 
 async def alitellm_completion(request: Dict[str, Any], num_retries: int, cache={"no-cache": True, "no-store": True}):
-    retry_kwargs = dict(
-        retry_policy=_get_litellm_retry_policy(num_retries),
-        retry_strategy="exponential_backoff_retry",
-        # In LiteLLM version 1.55.3 (the first version that supports retry_policy as an argument
-        # to completion()), the default value of max_retries is non-zero for certain providers, and
-        # max_retries is stacked on top of the retry_policy. To avoid this, we set max_retries=0
-        max_retries=0,
-    )
-
-    stream_completion = _get_stream_completion_fn(request, retry_kwargs, cache, sync=False)
+    stream_completion = _get_stream_completion_fn(request, cache, sync=False)
     if stream_completion is None:
         return await litellm.acompletion(
             cache=cache,
-            **retry_kwargs,
+            num_retries=num_retries,
+            retry_strategy="exponential_backoff_retry",
             **request,
         )
 
@@ -365,32 +343,7 @@ async def alitellm_text_completion(
         api_key=api_key,
         api_base=api_base,
         prompt=prompt,
-        retry_policy=_get_litellm_retry_policy(num_retries),
-        # In LiteLLM version 1.55.3 (the first version that supports retry_policy as an argument
-        # to completion()), the default value of max_retries is non-zero for certain providers, and
-        # max_retries is stacked on top of the retry_policy. To avoid this, we set max_retries=0
-        max_retries=0,
+        num_retries=num_retries,
+        retry_strategy="exponential_backoff_retry",
         **request,
     )
-
-
-def _get_litellm_retry_policy(num_retries: int) -> RetryPolicy:
-    """
-    Get a LiteLLM retry policy for retrying requests when transient API errors occur.
-    Args:
-        num_retries: The number of times to retry a request if it fails transiently due to
-                     network error, rate limiting, etc. Requests are retried with exponential
-                     backoff.
-    Returns:
-        A LiteLLM RetryPolicy instance.
-    """
-    return RetryPolicy(
-        TimeoutErrorRetries=num_retries,
-        RateLimitErrorRetries=num_retries,
-        InternalServerErrorRetries=num_retries,
-        ContentPolicyViolationErrorRetries=num_retries,
-        # We don't retry on errors that are unlikely to be transient
-        # (e.g. bad request, invalid auth credentials)
-        BadRequestErrorRetries=0,
-        AuthenticationErrorRetries=0,
-    )
diff --git a/dspy/clients/lm_local.py b/dspy/clients/lm_local.py
index 8329db9786..79a0d8e023 100644
--- a/dspy/clients/lm_local.py
+++ b/dspy/clients/lm_local.py
@@ -49,12 +49,10 @@ def launch(lm: "LM", launch_kwargs: Optional[Dict[str, Any]] = None):
         if model.startswith("local:"):
             model = model[6:]
         if model.startswith("huggingface/"):
-            model = model[len("huggingface/"):]
+            model = model[len("huggingface/") :]
 
         logger.info(f"Grabbing a free port to launch an SGLang server for model {model}")
-        logger.info(
-            f"We see that CUDA_VISIBLE_DEVICES is {os.environ.get('CUDA_VISIBLE_DEVICES', 'unset')}"
-        )
+        logger.info(f"We see that CUDA_VISIBLE_DEVICES is {os.environ.get('CUDA_VISIBLE_DEVICES', 'unset')}")
         port = get_free_port()
         timeout = launch_kwargs.get("timeout", 1800)
         command = f"python -m sglang.launch_server --model-path {model} --port {port} --host 0.0.0.0"
@@ -111,9 +109,7 @@ def get_logs() -> str:
             return "".join(logs_buffer)
 
         # Let the user know server is up
-        logger.info(
-            f"Server ready on random port {port}! Logs are available via lm.get_logs() method on returned lm."
-        )
+        logger.info(f"Server ready on random port {port}! Logs are available via lm.get_logs() method on returned lm.")
 
         lm.kwargs["api_base"] = f"http://localhost:{port}/v1"
         lm.kwargs["api_key"] = "local"
@@ -121,10 +117,10 @@ def get_logs() -> str:
         lm.process = process
         lm.thread = thread
 
-
     @staticmethod
     def kill(lm: "LM", launch_kwargs: Optional[Dict[str, Any]] = None):
         from sglang.utils import terminate_process
+
         if not hasattr(lm, "process"):
             logger.info("No running server to kill.")
             return
@@ -168,7 +164,7 @@ def finetune(
             "bf16": True,
             "output_dir": output_dir,
         }
-        train_kwargs={**default_train_kwargs, **(train_kwargs or {})}
+        train_kwargs = {**default_train_kwargs, **(train_kwargs or {})}
         output_dir = train_kwargs["output_dir"]  # user might have changed the output_dir
 
         logger.info(f"Starting local training, will save to {output_dir}")
@@ -185,7 +181,7 @@ def finetune(
 def create_output_dir(model_name, data_path):
     model_str = model_name.replace("/", "-")
     time_str = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-    rnd_str = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
+    rnd_str = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
     model_identifier = f"{rnd_str}_{model_str}_{time_str}"
     output_dir = data_path.replace(".jsonl", "_" + model_identifier)
     return output_dir
@@ -204,16 +200,10 @@ def train_sft_locally(model_name, train_data, train_kwargs):
 
     device = train_kwargs.get("device", None)
     if device is None:
-        device = (
-            "cuda"
-            if torch.cuda.is_available()
-            else "mps" if torch.backends.mps.is_available() else "cpu"
-        )
+        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
     logger.info(f"Using device: {device}")
 
-    model = AutoModelForCausalLM.from_pretrained(
-        pretrained_model_name_or_path=model_name
-    ).to(device)
+    model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=model_name).to(device)
     tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=model_name)
 
     # Set up the chat format; generally only for non-chat model variants, hence the try-except.
@@ -229,21 +219,25 @@ def train_sft_locally(model_name, train_data, train_kwargs):
     logger.info("Creating dataset")
     if "max_seq_length" not in train_kwargs:
         train_kwargs["max_seq_length"] = 4096
-        logger.info(f"The 'train_kwargs' parameter didn't include a 'max_seq_length', defaulting to {train_kwargs['max_seq_length']}")
+        logger.info(
+            f"The 'train_kwargs' parameter didn't include a 'max_seq_length', defaulting to {train_kwargs['max_seq_length']}"
+        )
 
     from datasets import Dataset
 
     hf_dataset = Dataset.from_list(train_data)
+
     def tokenize_function(example):
-        return encode_sft_example(example, tokenizer, train_kwargs["max_seq_length"])
+        return encode_sft_example(example, tokenizer, train_kwargs["max_seq_length"])  # noqa: F821
+
     tokenized_dataset = hf_dataset.map(tokenize_function, batched=False)
     tokenized_dataset.set_format(type="torch")
     tokenized_dataset = tokenized_dataset.filter(lambda example: (example["labels"] != -100).any())
 
-    USE_PEFT = train_kwargs.get("use_peft", False)
+    use_peft = train_kwargs.get("use_peft", False)
     peft_config = None
 
-    if USE_PEFT:
+    if use_peft:
         from peft import LoraConfig
 
         rank_dimension = 32
@@ -293,8 +287,8 @@ def tokenize_function(example):
     # Save the model!
     trainer.save_model()
 
-    MERGE = True
-    if USE_PEFT and MERGE:
+    merge = True
+    if use_peft and merge:
         from peft import AutoPeftModelForCausalLM
 
         # Load PEFT model on CPU
@@ -305,9 +299,7 @@ def tokenize_function(example):
         )
 
         merged_model = model_.merge_and_unload()
-        merged_model.save_pretrained(
-            sft_config.output_dir, safe_serialization=True, max_shard_size="5GB"
-        )
+        merged_model.save_pretrained(sft_config.output_dir, safe_serialization=True, max_shard_size="5GB")
 
     # Clean up!
     import gc
@@ -330,7 +322,7 @@ def get_free_port() -> int:
         return s.getsockname()[1]
 
 
-def wait_for_server(base_url: str, timeout: int = None) -> None:
+def wait_for_server(base_url: str, timeout: Optional[int] = None) -> None:
     """
     Wait for the server to be ready by polling the /v1/models endpoint.
 
diff --git a/dspy/clients/openai.py b/dspy/clients/openai.py
index 6cfaf1a1cb..1f5952ceb3 100644
--- a/dspy/clients/openai.py
+++ b/dspy/clients/openai.py
@@ -5,50 +5,49 @@
 
 import openai
 
-from dspy.clients.provider import TrainingJob, Provider
+from dspy.clients.provider import Provider, TrainingJob
 from dspy.clients.utils_finetune import TrainDataFormat, TrainingStatus, save_data
 
-
 _OPENAI_MODELS = [
-  'gpt-4-turbo',
-  'gpt-4-turbo-2024-04-09',
-  'tts-1',
-  'tts-1-1106',
-  'chatgpt-4o-latest',
-  'dall-e-2',
-  'whisper-1',
-  'gpt-3.5-turbo-instruct',
-  'gpt-3.5-turbo',
-  'gpt-3.5-turbo-0125',
-  'babbage-002',
-  'davinci-002',
-  'gpt-4o-mini-2024-07-18',
-  'gpt-4o',
-  'dall-e-3',
-  'gpt-4o-mini',
-  'gpt-4o-2024-08-06',
-  'gpt-4o-2024-05-13',
-  'o1-preview',
-  'gpt-4o-audio-preview-2024-10-01',
-  'o1-mini-2024-09-12',
-  'gpt-4o-audio-preview',
-  'tts-1-hd',
-  'tts-1-hd-1106',
-  'o1-preview-2024-09-12',
-  'o1-mini',
-  'gpt-4-1106-preview',
-  'text-embedding-ada-002',
-  'gpt-3.5-turbo-16k',
-  'text-embedding-3-small',
-  'text-embedding-3-large',
-  'gpt-4o-realtime-preview-2024-10-01',
-  'gpt-4o-realtime-preview',
-  'gpt-3.5-turbo-1106',
-  'gpt-4-0613',
-  'gpt-4-turbo-preview',
-  'gpt-4-0125-preview',
-  'gpt-4',
-  'gpt-3.5-turbo-instruct-0914'
+    "gpt-4-turbo",
+    "gpt-4-turbo-2024-04-09",
+    "tts-1",
+    "tts-1-1106",
+    "chatgpt-4o-latest",
+    "dall-e-2",
+    "whisper-1",
+    "gpt-3.5-turbo-instruct",
+    "gpt-3.5-turbo",
+    "gpt-3.5-turbo-0125",
+    "babbage-002",
+    "davinci-002",
+    "gpt-4o-mini-2024-07-18",
+    "gpt-4o",
+    "dall-e-3",
+    "gpt-4o-mini",
+    "gpt-4o-2024-08-06",
+    "gpt-4o-2024-05-13",
+    "o1-preview",
+    "gpt-4o-audio-preview-2024-10-01",
+    "o1-mini-2024-09-12",
+    "gpt-4o-audio-preview",
+    "tts-1-hd",
+    "tts-1-hd-1106",
+    "o1-preview-2024-09-12",
+    "o1-mini",
+    "gpt-4-1106-preview",
+    "text-embedding-ada-002",
+    "gpt-3.5-turbo-16k",
+    "text-embedding-3-small",
+    "text-embedding-3-large",
+    "gpt-4o-realtime-preview-2024-10-01",
+    "gpt-4o-realtime-preview",
+    "gpt-3.5-turbo-1106",
+    "gpt-4-0613",
+    "gpt-4-turbo-preview",
+    "gpt-4-0125-preview",
+    "gpt-4",
+    "gpt-3.5-turbo-instruct-0914",
 ]
 
 
@@ -84,7 +83,6 @@ def status(self) -> TrainingStatus:
 
 
 class OpenAIProvider(Provider):
-    
     def __init__(self):
         super().__init__()
         self.finetunable = True
@@ -113,7 +111,7 @@ def is_provider_model(model: str) -> bool:
             return True
 
         return False
-    
+
     @staticmethod
     def _remove_provider_prefix(model: str) -> str:
         provider_prefix = "openai/"
@@ -179,7 +177,6 @@ def does_file_exist(file_id: str) -> bool:
         except Exception:
             return False
 
-
     @staticmethod
     def is_terminal_training_status(status: TrainingStatus) -> bool:
         return status in [
@@ -187,7 +184,7 @@ def is_terminal_training_status(status: TrainingStatus) -> bool:
             TrainingStatus.failed,
             TrainingStatus.cancelled,
         ]
-    
+
     @staticmethod
     def get_training_status(job_id: str) -> TrainingStatus:
         provider_status_to_training_status = {
@@ -234,11 +231,7 @@ def upload_data(data_path: str) -> str:
         return provider_file.id
 
     @staticmethod
-    def _start_remote_training(
-        train_file_id: str,
-        model: str,
-        train_kwargs: Optional[Dict[str, Any]] = None
-    ) -> str:
+    def _start_remote_training(train_file_id: str, model: str, train_kwargs: Optional[Dict[str, Any]] = None) -> str:
         train_kwargs = train_kwargs or {}
         provider_job = openai.fine_tuning.jobs.create(
             model=model,
diff --git a/dspy/clients/provider.py b/dspy/clients/provider.py
index a77ba8202a..b73f577c7b 100644
--- a/dspy/clients/provider.py
+++ b/dspy/clients/provider.py
@@ -1,14 +1,14 @@
 from abc import abstractmethod
 from concurrent.futures import Future
 from threading import Thread
-from typing import Any, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 from dspy.clients.utils_finetune import TrainDataFormat
-from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from dspy.clients.lm import LM
 
+
 class TrainingJob(Future):
     def __init__(
         self,
diff --git a/dspy/clients/utils_finetune.py b/dspy/clients/utils_finetune.py
index 6842b7ccf9..7d83c741fb 100644
--- a/dspy/clients/utils_finetune.py
+++ b/dspy/clients/utils_finetune.py
@@ -62,12 +62,12 @@ def save_data(
 
 
 def validate_data_format(
-        data: List[Dict[str, Any]],
-        data_format: TrainDataFormat,
-    ):
+    data: List[Dict[str, Any]],
+    data_format: TrainDataFormat,
+):
     find_err_funcs = {
         TrainDataFormat.CHAT: find_data_error_chat,
-        TrainDataFormat.COMPLETION: find_data_errors_completion
+        TrainDataFormat.COMPLETION: find_data_errors_completion,
     }
     err = f"Data format {data_format} is not supported."
     assert data_format in find_err_funcs, err
@@ -83,7 +83,7 @@ def validate_data_format(
         if isinstance(data_dict, dict):
             err = find_err_func(data_dict)
         if err:
-            err_dict = dict(index=ind, error=err)
+            err_dict = {"index": ind, "error": err}
             data_dict_errors.append(err_dict)
 
     if data_dict_errors:
@@ -96,9 +96,7 @@ def validate_data_format(
         raise ValueError(err)
 
 
-def find_data_errors_completion(
-        data_dict: Dict[str, str]
-    ) -> Optional[str]:
+def find_data_errors_completion(data_dict: Dict[str, str]) -> Optional[str]:
     keys = ["prompt", "completion"]
 
     assert isinstance(data_dict, dict)
@@ -106,7 +104,7 @@ def find_data_errors_completion(
     found_keys = sorted(data_dict.keys())
     if set(expected_keys) != set(found_keys):
         return f"Expected Keys: {expected_keys}; Found Keys: {found_keys}"
-    
+
     for key in keys:
         if not isinstance(data_dict[key], str):
             return f"Expected `{key}` to be of type `str`. Found: {type(data_dict[key])}"
@@ -114,9 +112,7 @@ def find_data_errors_completion(
 
 # Following functions are modified from the OpenAI cookbook:
 # https://cookbook.openai.com/examples/chat_finetuning_data_prep
-def find_data_error_chat(
-        messages: Dict[str, Any]
-    ) -> Optional[str]:
+def find_data_error_chat(messages: Dict[str, Any]) -> Optional[str]:
     assert isinstance(messages, dict)
 
     expected_keys = ["messages"]
@@ -133,9 +129,7 @@ def find_data_error_chat(
             return f"Error in message at index {ind}: {err}"
 
 
-def find_data_error_chat_message(
-        message: Dict[str, Any]
-    ) -> Optional[str]:
+def find_data_error_chat_message(message: Dict[str, Any]) -> Optional[str]:
     assert isinstance(message, dict)
 
     message_keys = sorted(["role", "content"])


Test patch:
diff --git a/tests/adapters/test_json_adapter.py b/tests/adapters/test_json_adapter.py
index b0242416b9..64f6285b50 100644
--- a/tests/adapters/test_json_adapter.py
+++ b/tests/adapters/test_json_adapter.py
@@ -105,3 +105,99 @@ async def test_json_adapter_async_call():
     lm = dspy.utils.DummyLM([{"answer": "Paris"}])
     result = await adapter.acall(lm, {}, signature, [], {"question": "What is the capital of France?"})
     assert result == [{"answer": "Paris"}]
+
+
+def test_json_adapter_on_pydantic_model():
+    from litellm.utils import ModelResponse, Message, Choices
+
+    class User(pydantic.BaseModel):
+        id: int
+        name: str
+        email: str
+
+    class Answer(pydantic.BaseModel):
+        analysis: str
+        result: str
+
+    class TestSignature(dspy.Signature):
+        user: User = dspy.InputField(desc="The user who asks the question")
+        question: str = dspy.InputField(desc="Question the user asks")
+        answer: Answer = dspy.OutputField(desc="Answer to this question")
+
+    program = dspy.Predict(TestSignature)
+
+    dspy.configure(lm=dspy.LM(model="openai/gpt4o", cache=False), adapter=dspy.JSONAdapter())
+
+    with mock.patch("litellm.completion") as mock_completion:
+        mock_completion.return_value = ModelResponse(
+            choices=[
+                Choices(
+                    message=Message(
+                        content="{'answer': {'analysis': 'Paris is the captial of France', 'result': 'Paris'}}"
+                    )
+                )
+            ],
+            model="openai/gpt4o",
+        )
+        result = program(
+            user={"id": 5, "name": "name_test", "email": "email_test"}, question="What is the capital of France?"
+        )
+
+        # Check that litellm.completion was called exactly once
+        mock_completion.assert_called_once()
+
+        _, call_kwargs = mock_completion.call_args
+        # Assert that there are exactly 2 messages (system + user)
+        assert len(call_kwargs["messages"]) == 2
+
+        assert call_kwargs["messages"][0]["role"] == "system"
+        content = call_kwargs["messages"][0]["content"]
+        assert content is not None
+
+        # Assert that system prompt includes correct input field descriptions
+        expected_input_fields = (
+            "1. `user` (User): The user who asks the question\n2. `question` (str): Question the user asks\n"
+        )
+        assert expected_input_fields in content
+
+        # Assert that system prompt includes correct output field description
+        expected_output_fields = "1. `answer` (Answer): Answer to this question\n"
+        assert expected_output_fields in content
+
+        # Assert that system prompt includes input formatting structure
+        expected_input_structure = "[[ ## user ## ]]\n{user}\n\n[[ ## question ## ]]\n{question}\n\n"
+        assert expected_input_structure in content
+
+        # Assert that system prompt includes output formatting structure
+        expected_output_structure = (  # noqa: Q000
+            "Outputs will be a JSON object with the following fields.\n\n{\n  "
+            '"answer": "{answer}        # note: the value you produce must adhere to the JSON schema: '
+            '{\\"type\\": \\"object\\", \\"properties\\": {\\"analysis\\": {\\"type\\": \\"string\\", \\"title\\": '
+            '\\"Analysis\\"}, \\"result\\": {\\"type\\": \\"string\\", \\"title\\": \\"Result\\"}}, \\"required\\": '
+            '[\\"analysis\\", \\"result\\"], \\"title\\": \\"Answer\\"}"\n}'
+        )
+        assert expected_output_structure in content
+
+        assert call_kwargs["messages"][1]["role"] == "user"
+        user_message_content = call_kwargs["messages"][1]["content"]
+        assert user_message_content is not None
+
+        # Assert that the user input data is formatted correctly
+        expected_input_data = (  # noqa: Q000
+            '[[ ## user ## ]]\n{"id": 5, "name": "name_test", "email": "email_test"}\n\n[[ ## question ## ]]\n'
+            "What is the capital of France?\n\n"
+        )
+        assert expected_input_data in user_message_content
+
+        # Assert that the adapter output has expected fields and values
+        assert result.answer.analysis == "Paris is the captial of France"
+        assert result.answer.result == "Paris"
+
+
+def test_json_adapter_parse_raise_error_on_mismatch_fields():
+    signature = dspy.make_signature("input1->output1")
+    adapter = dspy.JSONAdapter()
+    invalid_completion = '{"output": "Test output"}'
+    with pytest.raises(ValueError) as error:
+        adapter.parse(signature, invalid_completion)
+    assert str(error.value) == "Expected dict_keys(['output1']) but got dict_keys([])"
diff --git a/tests/clients/test_lm.py b/tests/clients/test_lm.py
index 7616b9f79c..3e214469c0 100644
--- a/tests/clients/test_lm.py
+++ b/tests/clients/test_lm.py
@@ -161,76 +161,31 @@ class ResponseFormat(pydantic.BaseModel):
     lm("Query")
 
 
-@pytest.mark.parametrize(
-    ("error_code", "expected_exception", "expected_num_retries"),
-    [
-        ("429", litellm.RateLimitError, 2),
-        ("504", litellm.Timeout, 3),
-        # Don't retry on user errors
-        ("400", litellm.BadRequestError, 0),
-        ("401", litellm.AuthenticationError, 0),
-        # TODO: LiteLLM retry logic isn't implemented properly for internal server errors
-        # and content policy violations, both of which may be transient and should be retried
-        # ("content-policy-violation, litellm.BadRequestError, 1),
-        # ("500", litellm.InternalServerError, 0, 1),
-    ],
-)
-def test_lm_chat_calls_are_retried_for_expected_failures(
-    litellm_test_server,
-    error_code,
-    expected_exception,
-    expected_num_retries,
-):
-    api_base, server_log_file_path = litellm_test_server
+def test_retry_number_set_correctly():
+    lm = dspy.LM("openai/gpt-4o-mini", num_retries=3)
+    with mock.patch("litellm.completion") as mock_completion:
+        lm("query")
 
-    openai_lm = dspy.LM(
-        model="openai/dspy-test-model",
-        api_base=api_base,
-        api_key="fakekey",
-        num_retries=expected_num_retries,
-        model_type="chat",
-    )
-    with pytest.raises(expected_exception):
-        openai_lm(error_code)
+    assert mock_completion.call_args.kwargs["num_retries"] == 3
 
-    request_logs = read_litellm_test_server_request_logs(server_log_file_path)
-    assert len(request_logs) == expected_num_retries + 1  # 1 initial request + 1 retries
 
+def test_retry_made_on_system_errors():
+    retry_tracking = [0]  # Using a list to track retries
 
-@pytest.mark.parametrize(
-    ("error_code", "expected_exception", "expected_num_retries"),
-    [
-        ("429", litellm.RateLimitError, 2),
-        ("504", litellm.Timeout, 3),
-        # Don't retry on user errors
-        ("400", litellm.BadRequestError, 0),
-        ("401", litellm.AuthenticationError, 0),
-        # TODO: LiteLLM retry logic isn't implemented properly for internal server errors
-        # and content policy violations, both of which may be transient and should be retried
-        # ("content-policy-violation, litellm.BadRequestError, 2),
-        # ("500", litellm.InternalServerError, 0, 2),
-    ],
-)
-def test_lm_text_calls_are_retried_for_expected_failures(
-    litellm_test_server,
-    error_code,
-    expected_exception,
-    expected_num_retries,
-):
-    api_base, server_log_file_path = litellm_test_server
+    def mock_create(*args, **kwargs):
+        retry_tracking[0] += 1
+        # These fields are called during the error handling
+        mock_response = mock.Mock()
+        mock_response.headers = {}
+        mock_response.status_code = 429
+        raise RateLimitError(response=mock_response, message="message", body="error")
 
-    openai_lm = dspy.LM(
-        model="openai/dspy-test-model",
-        api_base=api_base,
-        api_key="fakekey",
-        num_retries=expected_num_retries,
-        model_type="text",
-    )
-    with pytest.raises(expected_exception):
-        openai_lm(error_code)
+    lm = dspy.LM(model="openai/gpt-4o-mini", max_tokens=250, num_retries=3)
+    with mock.patch.object(litellm.OpenAIChatCompletion, "completion", side_effect=mock_create):
+        with pytest.raises(RateLimitError):
+            lm("question")
 
-    request_logs = read_litellm_test_server_request_logs(server_log_file_path)
-    assert len(request_logs) == expected_num_retries + 1  # 1 initial request + 1 retries
+    assert retry_tracking[0] == 4
 
 
 def test_reasoning_model_token_parameter():