Repository: stanfordnlp/dspy
Pull request #6580 (2025-01-27 00:09:46)

Handle lists of models in parameters
Simple fix for https://github.com/stanfordnlp/dspy/issues/3920

When dumping state check if the field is a list, and if so convert each item to a serializable dict if applicable. Lists of images are not handled.

I also changed it so that models are converted to dicts rather than a JSON string. This makes more sense to me as then the entire resulting object will be converted to a JSON string later.

I added a simple test case that passes.


Patch:
diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml
new file mode 100644
index 00000000..700e8be7
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug_report.yml
@@ -0,0 +1,38 @@
+
+name: Bug Report
+description: Report a bug in the project
+title: "[Bug] "
+labels: bug
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## 🐛 Bug Report
+        Please fill out all required fields to help us diagnose and fix the issue.
+
+  - type: textarea
+    id: description
+    attributes:
+      label: "What happened?"
+      description: "Clearly describe the unexpected behavior."
+      placeholder: "Example: When I try to save a file, I get an error message..."
+    validations:
+      required: true
+
+  - type: textarea
+    id: steps-to-reproduce
+    attributes:
+      label: "Steps to reproduce"
+      description: "Tell us how to reproduce the issue."
+      placeholder: "Please provide a code snippet or a github gist for reproducing purpose."
+    validations:
+      required: true
+
+  - type: input
+    id: environment
+    attributes:
+      label: "DSPy version"
+      description: "Tell us your DSPy version."
+    validations:
+      required: true
+
diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml
new file mode 100644
index 00000000..3ba13e0c
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/config.yml
@@ -0,0 +1,1 @@
+blank_issues_enabled: false
diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml
new file mode 100644
index 00000000..f8fdcc20
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature_request.yml
@@ -0,0 +1,33 @@
+name: Feature Request
+description: Suggest a new feature or improvement
+title: "[Feature] "
+labels: enhancement
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## 🚀 Feature Request
+        Please fill out the following details.
+
+  - type: textarea
+    id: description
+    attributes:
+      label: "What feature would you like to see?"
+      description: "Describe the feature clearly."
+    validations:
+      required: true
+
+  - type: checkboxes
+    id: contribute
+    attributes:
+      label: "Would you like to contribute?"
+      options:
+        - label: Yes, I'd like to help implement this.
+        - label: No, I just want to request it.
+
+  - type: textarea
+    id: additional-info
+    attributes:
+      label: "Additional Context"
+      description: "Any links, references, or extra details?"
+      placeholder: "Example: This feature exists in XYZ tool."
diff --git a/README.md b/README.md
index 8b3c6d63..b95fa849 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,7 @@
 
 **Documentation:** [DSPy Docs](https://dspy.ai/)
 
-[![Downloads](https://static.pepy.tech/badge/dspy-ai)](https://pepy.tech/project/dspy-ai)  [![Downloads](https://static.pepy.tech/badge/dspy/month)](https://pepy.tech/project/dspy)
+[![PyPI Downloads](https://static.pepy.tech/badge/dspy)](https://pepy.tech/projects/dspy)  [![PyPI Downloads](https://static.pepy.tech/badge/dspy/month)](https://pepy.tech/projects/dspy)
 
 
 ----
diff --git a/docs/docs/community/community-resources.md b/docs/docs/community/community-resources.md
index 1eb79ee1..b2011cbd 100644
--- a/docs/docs/community/community-resources.md
+++ b/docs/docs/community/community-resources.md
@@ -21,6 +21,7 @@ This is the list of tutorials and blog posts on DSPy. If you would like to add y
 | **Intro to DSPy: Goodbye Prompting, Hello Programming** | [Link](https://medium.com/towards-data-science/intro-to-dspy-goodbye-prompting-hello-programming-4ca1c6ce3eb9) |
 | **DSPyGen: Revolutionizing AI** | [Link](https://www.linkedin.com/pulse/launch-alert-dspygen-20242252-revolutionizing-ai-sean-chatman--g9f1c/) |
 | **Building an AI Assistant with DSPy** | [Link](https://www.linkedin.com/pulse/building-ai-assistant-dspy-valliappa-lakshmanan-vgnsc/) |
+| **Building Self-improving Agents in Production with DSPy** | [Link](https://relevanceai.com/blog/building-self-improving-agentic-systems-in-production-with-dspy) |
 
 
 ### Videos
@@ -42,7 +43,7 @@ This is the list of tutorials and blog posts on DSPy. If you would like to add y
 | **DSPy End-to-End: SF Meetup** | [Link](https://www.youtube.com/watch?v=Y81DoFmt-2U) |
 | **Monitoring & Tracing DSPy with Langtrace** | [Link](https://langtrace.ai/blog/announcing-dspy-support-in-langtrace) |
 | **Teaching chat models to solve chess puzzles using DSPy + Finetuning** | [Link](https://raw.sh/posts/chess_puzzles) |
-
+| **Build Self-Improving AI Agents with DSPy (No Code)** | [Link](https://www.youtube.com/watch?v=UY8OsMlV21Y) |
 
 ### Podcasts
 
diff --git a/docs/docs/community/use-cases.md b/docs/docs/community/use-cases.md
index c89510fb..a948ec77 100644
--- a/docs/docs/community/use-cases.md
+++ b/docs/docs/community/use-cases.md
@@ -104,5 +104,6 @@ TODO: This list in particular is highly incomplete. There are a couple dozen oth
 | **Langtrace** | [Link](https://docs.langtrace.ai/supported-integrations/llm-frameworks/dspy) |
 | **Langfuse** | [Link](https://langfuse.com/docs/integrations/dspy) |
 | **OpenLIT** | [Link](https://docs.openlit.io/latest/integrations/dspy) |
+| **Relevance AI** | [Link](https://relevanceai.com/blog/dspy-programming---not-prompting---language-models) |
 
 Credit: Some of these resources were originally compiled in the [Awesome DSPy](https://github.com/ganarajpr/awesome-dspy/tree/master) repo.
diff --git a/dspy/.internal_dspyai/setup.py b/dspy/.internal_dspyai/setup.py
index c84e0f95..c992c394 100644
--- a/dspy/.internal_dspyai/setup.py
+++ b/dspy/.internal_dspyai/setup.py
@@ -8,7 +8,7 @@ setup(
     #replace_package_name_marker
     name="dspy-ai",
     #replace_package_version_marker
-    version="2.6.0rc8", 	
+    version="2.6.0", 	
     description="DSPy",	
     long_description=long_description,	
     long_description_content_type="text/markdown",	
@@ -19,5 +19,5 @@ setup(
     packages=find_packages(include=["dsp.*", "dspy.*", "dsp", "dspy"]),	
     python_requires=">=3.9",
     #replace_dspy_version_marker
-    install_requires=["dspy>=2.6.0rc8"]
+    install_requires=["dspy>=2.6.0"]
 )	
diff --git a/dspy/__metadata__.py b/dspy/__metadata__.py
index 99476341..06429933 100644
--- a/dspy/__metadata__.py
+++ b/dspy/__metadata__.py
@@ -1,7 +1,7 @@
 #replace_package_name_marker
 __name__="dspy"
 #replace_package_version_marker
-__version__="2.6.0rc8"
+__version__="2.6.0"
 __description__="DSPy"
 __url__="https://github.com/stanfordnlp/dspy"
 __author__="Omar Khattab"
diff --git a/dspy/predict/predict.py b/dspy/predict/predict.py
index 0e15137b..0eb5717e 100644
--- a/dspy/predict/predict.py
+++ b/dspy/predict/predict.py
@@ -38,7 +38,16 @@ class Predict(Module, Parameter):
                 if isinstance(demo[field], Image):
                     demo[field] = demo[field].model_dump()
                 elif isinstance(demo[field], BaseModel):
-                    demo[field] = demo[field].model_dump_json()
+                    demo[field] = demo[field].model_dump()
+                # Handle lists of BaseModels
+                elif isinstance(demo[field], list):
+                    processed_list = []
+                    for item in demo[field]:
+                        if isinstance(item, BaseModel):
+                            processed_list.append(item.model_dump())
+                        else:
+                            processed_list.append(item)
+                    demo[field] = processed_list
 
             state["demos"].append(demo)
 
diff --git a/dspy/propose/grounded_proposer.py b/dspy/propose/grounded_proposer.py
index 15ef1313..f618552b 100644
--- a/dspy/propose/grounded_proposer.py
+++ b/dspy/propose/grounded_proposer.py
@@ -5,7 +5,7 @@ from dspy.propose.dataset_summary_generator import create_dataset_summary
 from dspy.propose.utils import create_example_string, create_predictor_level_history_string, strip_prefix, get_dspy_source_code
 from dspy.teleprompt.utils import get_signature, get_prompt_model
 
-from dspy.propose.propose_base import Proposer
+from .propose_base import Proposer
 
 # Hardcoded variables (TODO: update)
 MAX_INSTRUCT_IN_HISTORY = 5  # 10
@@ -97,6 +97,9 @@ def generate_instruction_class(
             module = dspy.InputField(
                 desc="The module to create an instruction for.", prefix="MODULE:",
             )
+            module_description = dspy.InputField(
+                desc="Description of the module to create an instruction for.", prefix="MODULE DESCRIPTION:",
+            )
         task_demos = dspy.InputField(
             format=str,
             desc="Example inputs/outputs of our module.",
@@ -164,24 +167,39 @@ class GenerateModuleInstruction(dspy.Module):
         program,
         previous_instructions,
         data_summary,
-        max_demos=3,
+        num_demos_in_context=3,
         tip=None,
     ):
-        # Construct full program demo or single module demo depending on whether or not we're using the full program
-        task_demos = ""
+        def gather_examples_from_sets(candidate_sets, max_examples):
+            """Helper function to gather up to augmented examples from given sets."""
+            count = 0
+            for candidate_set in candidate_sets:
+                for example in candidate_set:
+                    if "augmented" in example.keys():
+                        fields_to_use = get_signature(program.predictors()[pred_i]).fields
+                        yield create_example_string(fields_to_use, example)
+                        count += 1
+                        if count >= max_examples:
+                            return
+
+        # Construct full program demo or single module demo depending on settings
         basic_instruction = get_signature(program.predictors()[pred_i]).instructions
-        curr_demos_num = 0
+        task_demos = ""
         
         if self.use_task_demos:
-            for example in demo_candidates[pred_i][demo_set_i]:
-                if "augmented" in example.keys():
-                    fields_to_use = get_signature(program.predictors()[pred_i]).fields
-                    example_string = create_example_string(fields_to_use, example)
-                    task_demos += f"{example_string}\n"
-                    curr_demos_num += 1
-                    if curr_demos_num >= max_demos:
-                        break
-        else:
+            # Combine current and adjacent sets
+            adjacent_sets = (
+                [demo_candidates[pred_i][demo_set_i]] +
+                demo_candidates[pred_i][demo_set_i + 1:] +
+                demo_candidates[pred_i][:demo_set_i]
+            )
+            
+            # Gather examples up to the required count
+            example_strings = gather_examples_from_sets(adjacent_sets, num_demos_in_context)
+            task_demos = "\n\n".join(example_strings) + "\n\n"
+
+        # Default to no demos provided if no examples were gathered, or if we're using the first demo set
+        if not task_demos.strip() or demo_set_i == 0:
             task_demos = "No task demos provided."
 
         # Summarize the program
@@ -224,20 +242,19 @@ class GenerateModuleInstruction(dspy.Module):
 
         # Generate an instruction for our chosen module
         if self.verbose: print(f"task_demos {task_demos}")
+
         instruct = self.generate_module_instruction(
             dataset_description=data_summary,
             program_code=self.program_code_string,
-            program_description=program_description,
             module=module_code,
+            program_description=program_description,
+            module_description=module_description,
             task_demos=task_demos,
             tip=tip,
             basic_instruction=basic_instruction,
             previous_instructions=previous_instructions,
-            module_description=module_description,
         )
-        if hasattr(instruct, "module_description"):
-            module_description = strip_prefix(instruct.module_description)
-            if self.verbose: print(f"MODULE DESCRIPTION: {module_description}")
+
         proposed_instruction = strip_prefix(instruct.proposed_instruction)
 
         return dspy.Prediction(proposed_instruction=proposed_instruction)
@@ -254,6 +271,7 @@ class GroundedProposer(Proposer):
         use_dataset_summary=True,
         program_aware=True,
         use_task_demos=True,
+        num_demos_in_context = 3,
         use_instruct_history=True,
         use_tip=True,
         set_tip_randomly=True,
@@ -265,6 +283,7 @@ class GroundedProposer(Proposer):
         self.program_aware = program_aware
         self.use_dataset_summary = use_dataset_summary
         self.use_task_demos = use_task_demos
+        self.num_demos_in_context = num_demos_in_context
         self.use_instruct_history = use_instruct_history
         self.use_tip = use_tip
         self.set_tip_randomly=set_tip_randomly
@@ -315,17 +334,16 @@ class GroundedProposer(Proposer):
             self.use_instruct_history = use_history
             if self.verbose: print(f"Use history T/F: {self.use_instruct_history}")
 
-        num_demos = max(len(demo_candidates[0]) if demo_candidates else N, 1)
-
         if not demo_candidates:
             if self.verbose: print("No demo candidates provided. Running without task demos.")
             self.use_task_demos = False
 
         # Create an instruction for each predictor 
         for pred_i, predictor in enumerate(program.predictors()):
-            for demo_set_i in range(num_demos):
+            for demo_set_i in range(len(demo_candidates[0])):
                 if pred_i not in proposed_instructions:
                     proposed_instructions[pred_i] = []
+                selected_tip = None
                 if self.set_tip_randomly:
                     if self.verbose: print("Using a randomly generated configuration for our grounded proposer.")
                     # Randomly select the tip
@@ -341,7 +359,6 @@ class GroundedProposer(Proposer):
                         program=program,
                         predictor=predictor,
                         pred_i=pred_i,
-                        # prompt_model=prompt_model,
                         T=T,
                         demo_candidates=demo_candidates,
                         demo_set_i=demo_set_i,
@@ -349,7 +366,7 @@ class GroundedProposer(Proposer):
                         tip=selected_tip,
                     ),
                 )
-
+        
         return proposed_instructions
 
     def propose_instruction_for_predictor(
@@ -385,8 +402,11 @@ class GroundedProposer(Proposer):
         # Generate a new instruction for our predictor, using the temperature specified for this round
         original_temp = self.prompt_model.kwargs["temperature"]
 
+        epsilon = self.rng.uniform(0.01, 0.05)
+        modified_temp = T + epsilon
+
         with dspy.settings.context(lm=self.prompt_model):
-            self.prompt_model.kwargs["temperature"] = T
+            self.prompt_model.kwargs["temperature"] = modified_temp
             proposed_instruction = instruction_generator.forward(
                 demo_candidates=demo_candidates,
                 pred_i=pred_i,
@@ -394,6 +414,7 @@ class GroundedProposer(Proposer):
                 program=program,
                 data_summary=self.data_summary,
                 previous_instructions=instruction_history,
+                num_demos_in_context = self.num_demos_in_context,
                 tip=tip,
             ).proposed_instruction
         self.prompt_model.kwargs["temperature"] = original_temp
@@ -402,4 +423,4 @@ class GroundedProposer(Proposer):
         if self.verbose: self.prompt_model.inspect_history(n=1)
         if self.verbose: print(f"PROPOSED INSTRUCTION: {proposed_instruction}")
 
-        return strip_prefix(proposed_instruction)
+        return strip_prefix(proposed_instruction)
\ No newline at end of file
diff --git a/dspy/propose/utils.py b/dspy/propose/utils.py
index 99e146f7..7d93e3c8 100644
--- a/dspy/propose/utils.py
+++ b/dspy/propose/utils.py
@@ -146,6 +146,7 @@ def get_dspy_source_code(module):
     base_code = ""
 
     # Don't get source code for Predict or ChainOfThought modules (NOTE we will need to extend this list as more DSPy.modules are added)
+    # TODO: if type(module).__name__ not in ["Predict", "ChainOfThought", "ReAct"]:
     if not type(module).__name__ == "Predict" and not type(module).__name__ == "ChainOfThought":
         try:
             base_code = inspect.getsource(type(module))
diff --git a/dspy/retrievers/embeddings.py b/dspy/retrievers/embeddings.py
index 75e1ff1f..ded5b153 100644
--- a/dspy/retrievers/embeddings.py
+++ b/dspy/retrievers/embeddings.py
@@ -34,7 +34,9 @@ class Embeddings:
 
     def forward(self, query: str):
         import dspy
-        return dspy.Prediction(passages=self.search_fn(query))
+
+        passages, indices = self.search_fn(query)
+        return dspy.Prediction(passages=passages, indices=indices)
 
     def _batch_forward(self, queries: List[str]):
         q_embeds = self.embedder(queries)
@@ -76,7 +78,7 @@ class Embeddings:
         top_k_indices = np.argsort(-scores, axis=1)[:, :self.k]
         top_indices = candidate_indices[np.arange(len(q_embeds))[:, None], top_k_indices]
 
-        return [[self.corpus[idx] for idx in indices] for indices in top_indices]
+        return [([self.corpus[idx] for idx in indices], [idx for idx in indices]) for indices in top_indices]
 
     def _normalize(self, embeddings: np.ndarray):
         norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
diff --git a/dspy/teleprompt/mipro_optimizer_v2.py b/dspy/teleprompt/mipro_optimizer_v2.py
index 5ac7f1df..0bba2322 100644
--- a/dspy/teleprompt/mipro_optimizer_v2.py
+++ b/dspy/teleprompt/mipro_optimizer_v2.py
@@ -6,6 +6,7 @@ from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
 
 import numpy as np
 import optuna
+from optuna.distributions import CategoricalDistribution
 
 import dspy
 from dspy.evaluate.evaluate import Evaluate
@@ -437,6 +438,7 @@ class MIPROv2(Teleprompter):
             program_aware=program_aware_proposer,
             use_dataset_summary=data_aware_proposer,
             use_task_demos=fewshot_aware_proposer,
+            num_demos_in_context = BOOTSTRAPPED_FEWSHOT_EXAMPLES_IN_CONTEXT,
             use_tip=tip_aware_proposer,
             set_tip_randomly=tip_aware_proposer,
             use_instruct_history=False,
@@ -477,8 +479,19 @@ class MIPROv2(Teleprompter):
         minibatch_full_eval_steps: int,
         seed: int,
     ) -> Optional[Any]:
-        logger.info("Evaluating the default program...\n")
-        default_score = eval_candidate_program(len(valset), valset, program, evaluate, self.rng)
+
+        # Run optimization
+        optuna.logging.set_verbosity(optuna.logging.WARNING)
+        logger.info("==> STEP 3: FINDING OPTIMAL PROMPT PARAMETERS <==")
+        logger.info(
+            "We will evaluate the program over a series of trials with different combinations of instructions and few-shot examples to find the optimal combination using Bayesian Optimization.\n"
+        )
+
+        # Compute the adjusted total trials that we will run (including full evals)
+        adjusted_num_trials = (num_trials + num_trials // minibatch_full_eval_steps + 1) if minibatch else num_trials
+        logger.info(f"== Trial {1} / {adjusted_num_trials} - Full Evaluation of Default Program ==")
+
+        default_score, baseline_results = eval_candidate_program(len(valset), valset, program, evaluate, self.rng, return_all_scores=True)
         logger.info(f"Default program score: {default_score}\n")
 
         trial_logs = {}
@@ -494,7 +507,7 @@ class MIPROv2(Teleprompter):
         best_score = default_score
         best_program = program.deepcopy()
         total_eval_calls = len(valset)
-        score_data= [(best_score, program.deepcopy(), True)]
+        score_data = [{"score": best_score, "program": program.deepcopy(), "full_eval": True}]
         param_score_dict = defaultdict(list)
         fully_evaled_param_combos = {}
 
@@ -504,7 +517,7 @@ class MIPROv2(Teleprompter):
 
             trial_num = trial.number + 1
             if minibatch:
-                logger.info(f"== Minibatch Trial {trial_num} / {num_trials} ==")
+                logger.info(f"== Trial {trial_num} / {adjusted_num_trials} - Minibatch ==")
             else:
                 logger.info(f"===== Trial {trial_num} / {num_trials} =====")
 
@@ -514,7 +527,7 @@ class MIPROv2(Teleprompter):
             candidate_program = program.deepcopy()
 
             # Choose instructions and demos, insert them into the program
-            chosen_params = self._select_and_insert_instructions_and_demos(
+            chosen_params, raw_chosen_params = self._select_and_insert_instructions_and_demos(
                 candidate_program,
                 instruction_candidates,
                 demo_candidates,
@@ -531,7 +544,7 @@ class MIPROv2(Teleprompter):
             # Evaluate the candidate program (on minibatch if minibatch=True)
             batch_size = minibatch_size if minibatch else len(valset)
             score = eval_candidate_program(
-                batch_size, valset, candidate_program, evaluate, self.rng
+               batch_size, valset, candidate_program, evaluate, self.rng
             )
             total_eval_calls += batch_size
 
@@ -542,7 +555,7 @@ class MIPROv2(Teleprompter):
                 logger.info(f"{GREEN}Best full score so far!{ENDC} Score: {score}")
 
             # Log evaluation results
-            score_data.append((score, candidate_program, batch_size >= len(valset))) # score, prog, full_eval
+            score_data.append({"score": score, "program": candidate_program, "full_eval": batch_size >= len(valset)}) # score, prog, full_eval
             if minibatch:
                 self._log_minibatch_eval(
                     score,
@@ -551,7 +564,7 @@ class MIPROv2(Teleprompter):
                     chosen_params,
                     score_data,
                     trial,
-                    num_trials,
+                    adjusted_num_trials,
                     trial_logs,
                     trial_num,
                     candidate_program,
@@ -563,16 +576,17 @@ class MIPROv2(Teleprompter):
                 )
             categorical_key = ",".join(map(str, chosen_params))
             param_score_dict[categorical_key].append(
-                (score, candidate_program),
+                (score, candidate_program, raw_chosen_params),
             )
 
-            # If minibatch, perform full evaluation at intervals
+            # If minibatch, perform full evaluation at intervals (and at the very end)
             if minibatch and (
                 (trial_num % minibatch_full_eval_steps == 0)
-                or (trial_num == num_trials)
+                or (trial_num == (adjusted_num_trials -1))
             ):
                 best_score, best_program, total_eval_calls = self._perform_full_evaluation(
                     trial_num,
+                    adjusted_num_trials,
                     param_score_dict,
                     fully_evaled_param_combos,
                     evaluate,
@@ -582,20 +596,28 @@ class MIPROv2(Teleprompter):
                     score_data,
                     best_score,
                     best_program,
+                    study,
+                    instruction_candidates,
+                    demo_candidates
                 )
-
+            
             return score
 
-        # Run optimization
-        optuna.logging.set_verbosity(optuna.logging.WARNING)
-        logger.info("==> STEP 3: FINDING OPTIMAL PROMPT PARAMETERS <==")
-        logger.info(
-            "We will evaluate the program over a series of trials with different combinations of instructions and few-shot examples to find the optimal combination using Bayesian Optimization.\n"
-        )
-
         sampler = optuna.samplers.TPESampler(seed=seed, multivariate=True)
         study = optuna.create_study(direction="maximize", sampler=sampler)
-        study.optimize(objective, n_trials=num_trials)
+
+        default_params = {f"{i}_predictor_instruction": 0 for i in range(len(program.predictors()))}
+        if demo_candidates:
+            default_params.update({f"{i}_predictor_demos": 0 for i in range(len(program.predictors()))})
+
+        # Add default run as a baseline in optuna (TODO: figure out how to weight this by # of samples evaluated on)
+        trial = optuna.trial.create_trial(
+            params=default_params,
+            distributions= self._get_param_distributions(program, instruction_candidates, demo_candidates),
+            value=default_score,
+        )
+        study.add_trial(trial)
+        study.optimize(objective, n_trials=num_trials-1)
 
         # Attach logs to best program
         if best_program is not None and self.track_stats:
@@ -603,11 +625,11 @@ class MIPROv2(Teleprompter):
             best_program.score = best_score
             best_program.prompt_model_total_calls = self.prompt_model_total_calls
             best_program.total_calls = self.total_calls
-            sorted_candidate_programs = sorted(score_data, key=lambda x: x[0], reverse=True)
+            sorted_candidate_programs = sorted(score_data, key=lambda x: x["score"], reverse=True)
             # Attach all minibatch programs
-            best_program.mb_candidate_programs = [score_data for score_data in sorted_candidate_programs if not score_data[2]]
+            best_program.mb_candidate_programs = [score_data for score_data in sorted_candidate_programs if not score_data["full_eval"]]
             # Attach all programs that were evaluated on the full trainset, in descending order of score
-            best_program.candidate_programs = [score_data for score_data in sorted_candidate_programs if score_data[2]]
+            best_program.candidate_programs = [score_data for score_data in sorted_candidate_programs if score_data["full_eval"]]
 
         logger.info(f"Returning best identified program with score {best_score}!")
 
@@ -621,7 +643,7 @@ class MIPROv2(Teleprompter):
         chosen_params,
         score_data,
         trial,
-        num_trials,
+        adjusted_num_trials,
         trial_logs,
         trial_num,
         candidate_program,
@@ -633,16 +655,18 @@ class MIPROv2(Teleprompter):
         trial_logs[trial_num]["mb_score"] = score
         trial_logs[trial_num]["total_eval_calls_so_far"] = total_eval_calls
         trial_logs[trial_num]["mb_program"] = candidate_program.deepcopy()
-
+        
         logger.info(
             f"Score: {score} on minibatch of size {batch_size} with parameters {chosen_params}."
         )
-        logger.info(f"Minibatch scores so far: {'['+', '.join([f'{s[0]}' for s in score_data if not s[2]]) +']'}")
-        trajectory = "[" + ", ".join([f"{s[0]}" for s in score_data if s[2]]) + "]"
+        minibatch_scores = ', '.join([f'{s["score"]}' for s in score_data if not s["full_eval"]])
+        logger.info(f"Minibatch scores so far: {'['+ minibatch_scores +']'}")
+        full_eval_scores = ', '.join([f'{s["score"]}' for s in score_data if s["full_eval"]])
+        trajectory = "[" + full_eval_scores + "]"
         logger.info(f"Full eval scores so far: {trajectory}")
         logger.info(f"Best full score so far: {best_score}")
         logger.info(
-            f'{"="*len(f"== Minibatch Trial {trial.number+1} / {num_trials} ==")}\n\n'
+            f'{"="*len(f"== Trial {trial.number+1} / {adjusted_num_trials} - Minibatch Evaluation ==")}\n\n'
         )
 
     def _log_normal_eval(
@@ -656,7 +680,8 @@ class MIPROv2(Teleprompter):
         trial_logs[trial_num]["full_eval_program"] = candidate_program.deepcopy()
 
         logger.info(f"Score: {score} with parameters {chosen_params}.")
-        logger.info(f"Scores so far: {'['+', '.join([f'{s[0]}' for s in score_data if s[2]])+']'}")
+        full_eval_scores = ', '.join([f'{s["score"]}' for s in score_data if s["full_eval"]])
+        logger.info(f"Scores so far: {'['+full_eval_scores+']'}")
         logger.info(f"Best score so far: {best_score}")
         logger.info(f'{"="*len(f"===== Trial {trial.number+1} / {num_trials} =====")}\n\n')
 
@@ -670,6 +695,7 @@ class MIPROv2(Teleprompter):
         trial_num: int,
     ) -> List[str]:
         chosen_params = []
+        raw_chosen_params = {}
 
         for i, predictor in enumerate(candidate_program.predictors()):
             # Select instruction
@@ -683,7 +709,7 @@ class MIPROv2(Teleprompter):
             set_signature(predictor, updated_signature)
             trial_logs[trial_num][f"{i}_predictor_instruction"] = instruction_idx
             chosen_params.append(f"Predictor {i}: Instruction {instruction_idx}")
-
+            raw_chosen_params[f"{i}_predictor_instruction"] = instruction_idx
             # Select demos if available
             if demo_candidates:
                 demos_idx = trial.suggest_categorical(
@@ -692,12 +718,24 @@ class MIPROv2(Teleprompter):
                 predictor.demos = demo_candidates[i][demos_idx]
                 trial_logs[trial_num][f"{i}_predictor_demos"] = demos_idx
                 chosen_params.append(f"Predictor {i}: Few-Shot Set {demos_idx}")
+                raw_chosen_params[f"{i}_predictor_demos"] = instruction_idx
+
+        return chosen_params, raw_chosen_params
 
-        return chosen_params
+    def _get_param_distributions(self, program, instruction_candidates, demo_candidates):
+        param_distributions = {}
+
+        for i in range(len(instruction_candidates)):
+            param_distributions[f"{i}_predictor_instruction"] = CategoricalDistribution(range(len(instruction_candidates[i])))
+            if demo_candidates:
+                param_distributions[f"{i}_predictor_demos"] = CategoricalDistribution(range(len(demo_candidates[i])))
+
+        return param_distributions
 
     def _perform_full_evaluation(
         self,
         trial_num: int,
+        adjusted_num_trials: int,
         param_score_dict: Dict,
         fully_evaled_param_combos: Dict,
         evaluate: Evaluate,
@@ -707,11 +745,14 @@ class MIPROv2(Teleprompter):
         score_data,
         best_score: float,
         best_program: Any,
+        study: optuna.Study,
+        instruction_candidates: List,
+        demo_candidates: List,
     ):
-        logger.info(f"===== Full Eval {len(fully_evaled_param_combos)+1} =====")
+        logger.info(f"===== Trial {trial_num+1} /{adjusted_num_trials} - Full Evaluation =====")
 
         # Identify best program to evaluate fully
-        highest_mean_program, mean_score, combo_key = (
+        highest_mean_program, mean_score, combo_key, params = (
             get_program_with_highest_avg_score(
                 param_score_dict, fully_evaled_param_combos
             )
@@ -722,7 +763,15 @@ class MIPROv2(Teleprompter):
         full_eval_score = eval_candidate_program(
             len(valset), valset, highest_mean_program, evaluate, self.rng
         )
-        score_data.append((full_eval_score, highest_mean_program, True))
+        score_data.append({"score": full_eval_score, "program": highest_mean_program, "full_eval": True})
+
+        # Log full eval as a trial so that optuna can learn from the new results
+        trial = optuna.trial.create_trial(
+            params=params,
+            distributions= self._get_param_distributions(best_program, instruction_candidates, demo_candidates),
+            value=full_eval_score,
+        )
+        study.add_trial(trial)
 
         # Log full evaluation results
         fully_evaled_param_combos[combo_key] = {
@@ -745,7 +794,8 @@ class MIPROv2(Teleprompter):
             logger.info(f"{GREEN}New best full eval score!{ENDC} Score: {full_eval_score}")
             best_score = full_eval_score
             best_program = highest_mean_program.deepcopy()
-        trajectory = "[" + ", ".join([f"{s[0]}" for s in score_data if s[2]]) + "]"
+        full_eval_scores = ', '.join([f'{s["score"]}' for s in score_data if s["full_eval"]])
+        trajectory = "[" + full_eval_scores + "]"
         logger.info(f"Full eval scores so far: {trajectory}")
         logger.info(f"Best full score so far: {best_score}")
         logger.info(len(f"===== Full Eval {len(fully_evaled_param_combos)+1} =====") * "=")
diff --git a/dspy/teleprompt/random_search.py b/dspy/teleprompt/random_search.py
index 4fa8ba23..c054b105 100644
--- a/dspy/teleprompt/random_search.py
+++ b/dspy/teleprompt/random_search.py
@@ -135,7 +135,7 @@ class BootstrapFewShotWithRandomSearch(Teleprompter):
             print(f"Scores so far: {scores}")
             print(f"Best score so far: {max(scores)}")
 
-            score_data.append((score, subscores, seed, program))
+            score_data.append({"score": score, "subscores": subscores, "seed": seed, "program": program})
 
             if self.stop_at_score is not None and score >= self.stop_at_score:
                 print(f"Stopping early because score {score} is >= stop_at_score {self.stop_at_score}")
@@ -143,7 +143,7 @@ class BootstrapFewShotWithRandomSearch(Teleprompter):
 
         # To best program, attach all program candidates in decreasing average score
         best_program.candidate_programs = score_data
-        best_program.candidate_programs = sorted(best_program.candidate_programs, key=lambda x: x[0], reverse=True)
+        best_program.candidate_programs = sorted(best_program.candidate_programs, key=lambda x: x["score"], reverse=True)
 
         print(f"{len(best_program.candidate_programs)} candidate programs found.")
 
diff --git a/dspy/teleprompt/utils.py b/dspy/teleprompt/utils.py
index bdc6a529..96cc2d44 100644
--- a/dspy/teleprompt/utils.py
+++ b/dspy/teleprompt/utils.py
@@ -42,25 +42,23 @@ def create_minibatch(trainset, batch_size=50, rng=None):
     return minibatch
 
 
-def eval_candidate_program(batch_size, trainset, candidate_program, evaluate, rng=None):
+def eval_candidate_program(batch_size, trainset, candidate_program, evaluate, rng=None, return_all_scores=False):
     """Evaluate a candidate program on the trainset, using the specified batch size."""
 
     try:
         # Evaluate on the full trainset
         if batch_size >= len(trainset):
-            score = evaluate(candidate_program, devset=trainset)
+            return evaluate(candidate_program, devset=trainset, return_all_scores=return_all_scores)
         # Or evaluate on a minibatch
         else:
-            score = evaluate(
+            return evaluate(
                 candidate_program,
                 devset=create_minibatch(trainset, batch_size, rng),
+                return_all_scores=return_all_scores
             )
     except Exception as e:
         print(f"Exception occurred: {e}")
-        score = 0.0  # TODO: Handle this better, as -ve scores are possible
-
-    return score
-
+        return 0.0  # TODO: Handle this better, as -ve scores are possible
 
 def eval_candidate_program_with_pruning(
     trial, trial_logs, trainset, candidate_program, evaluate, trial_num, batch_size=100,
@@ -114,22 +112,23 @@ def get_program_with_highest_avg_score(param_score_dict, fully_evaled_param_comb
         scores = np.array([v[0] for v in values])
         mean = np.average(scores)
         program = values[0][1]
-        results.append((key, mean, program))
+        params = values[0][2]
+        results.append((key, mean, program, params))
 
     # Sort results by the mean
     sorted_results = sorted(results, key=lambda x: x[1], reverse=True)
 
     # Find the combination with the highest mean, skip fully evaluated ones
     for combination in sorted_results:
-        key, mean, program = combination
+        key, mean, program, params = combination
 
         if key in fully_evaled_param_combos:
             continue
 
-        return program, mean, key
+        return program, mean, key, params
 
     # If no valid program is found, we return the last valid one that we found
-    return program, mean, key
+    return program, mean, key, params
 
 
 def calculate_last_n_proposed_quality(
diff --git a/poetry.lock b/poetry.lock
index 0b51e3b2..23e1abc3 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 2.0.0 and should not be changed by hand.
+# This file is automatically @generated by Poetry 2.0.1 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -2876,15 +2876,15 @@ tests = ["aiohttp", "boto3", "duckdb", "pandas (>=1.4)", "polars (>=0.19)", "pyt
 
 [[package]]
 name = "litellm"
-version = "1.57.4"
+version = "1.59.8"
 description = "Library to easily interface with LLM API providers"
 optional = false
 python-versions = "!=2.7.*,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,!=3.7.*,>=3.8"
 groups = ["main", "dev"]
 markers = "python_version <= \"3.11\" or python_version >= \"3.12\""
 files = [
-    {file = "litellm-1.57.4-py3-none-any.whl", hash = "sha256:afe48924d8a36db801018970a101622fce33d117fe9c54441c0095c491511abb"},
-    {file = "litellm-1.57.4.tar.gz", hash = "sha256:747a870ddee9c71f9560fc68ad02485bc1008fcad7d7a43e87867a59b8ed0669"},
+    {file = "litellm-1.59.8-py3-none-any.whl", hash = "sha256:2473914bd2343485a185dfe7eedb12ee5fda32da3c9d9a8b73f6966b9b20cf39"},
+    {file = "litellm-1.59.8.tar.gz", hash = "sha256:9d645cc4460f6a9813061f07086648c4c3d22febc8e1f21c663f2b7750d90512"},
 ]
 
 [package.dependencies]
@@ -2911,11 +2911,12 @@ pyyaml = {version = ">=6.0.1,<7.0.0", optional = true, markers = "extra == \"pro
 rq = {version = "*", optional = true, markers = "extra == \"proxy\""}
 tiktoken = ">=0.7.0"
 tokenizers = "*"
-uvicorn = {version = ">=0.22.0,<0.23.0", optional = true, markers = "extra == \"proxy\""}
+uvicorn = {version = ">=0.29.0,<0.30.0", optional = true, markers = "extra == \"proxy\""}
+uvloop = {version = ">=0.21.0,<0.22.0", optional = true, markers = "extra == \"proxy\""}
 
 [package.extras]
 extra-proxy = ["azure-identity (>=1.15.0,<2.0.0)", "azure-keyvault-secrets (>=4.8.0,<5.0.0)", "google-cloud-kms (>=2.21.3,<3.0.0)", "prisma (==0.11.0)", "resend (>=0.8.0,<0.9.0)"]
-proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "cryptography (>=43.0.1,<44.0.0)", "fastapi (>=0.115.5,<0.116.0)", "fastapi-sso (>=0.16.0,<0.17.0)", "gunicorn (>=22.0.0,<23.0.0)", "orjson (>=3.9.7,<4.0.0)", "pynacl (>=1.5.0,<2.0.0)", "python-multipart (>=0.0.18,<0.0.19)", "pyyaml (>=6.0.1,<7.0.0)", "rq", "uvicorn (>=0.22.0,<0.23.0)"]
+proxy = ["PyJWT (>=2.8.0,<3.0.0)", "apscheduler (>=3.10.4,<4.0.0)", "backoff", "cryptography (>=43.0.1,<44.0.0)", "fastapi (>=0.115.5,<0.116.0)", "fastapi-sso (>=0.16.0,<0.17.0)", "gunicorn (>=22.0.0,<23.0.0)", "orjson (>=3.9.7,<4.0.0)", "pynacl (>=1.5.0,<2.0.0)", "python-multipart (>=0.0.18,<0.0.19)", "pyyaml (>=6.0.1,<7.0.0)", "rq", "uvicorn (>=0.29.0,<0.30.0)", "uvloop (>=0.21.0,<0.22.0)"]
 
 [[package]]
 name = "livereload"
@@ -7840,15 +7841,15 @@ zstd = ["zstandard (>=0.18.0)"]
 
 [[package]]
 name = "uvicorn"
-version = "0.22.0"
+version = "0.29.0"
 description = "The lightning-fast ASGI server."
 optional = false
-python-versions = ">=3.7"
+python-versions = ">=3.8"
 groups = ["main", "dev"]
 markers = "python_version <= \"3.11\" or python_version >= \"3.12\""
 files = [
-    {file = "uvicorn-0.22.0-py3-none-any.whl", hash = "sha256:e9434d3bbf05f310e762147f769c9f21235ee118ba2d2bf1155a7196448bd996"},
-    {file = "uvicorn-0.22.0.tar.gz", hash = "sha256:79277ae03db57ce7d9aa0567830bbb51d7a612f54d6e1e3e92da3ef24c2c8ed8"},
+    {file = "uvicorn-0.29.0-py3-none-any.whl", hash = "sha256:2c2aac7ff4f4365c206fd773a39bf4ebd1047c238f8b8268ad996829323473de"},
+    {file = "uvicorn-0.29.0.tar.gz", hash = "sha256:6a69214c0b6a087462412670b3ef21224fa48cae0e452b5883e8e8bdfdd11dd0"},
 ]
 
 [package.dependencies]
@@ -7858,6 +7859,7 @@ h11 = ">=0.8"
 httptools = {version = ">=0.5.0", optional = true, markers = "extra == \"standard\""}
 python-dotenv = {version = ">=0.13", optional = true, markers = "extra == \"standard\""}
 pyyaml = {version = ">=5.1", optional = true, markers = "extra == \"standard\""}
+typing-extensions = {version = ">=4.0", markers = "python_version < \"3.11\""}
 uvloop = {version = ">=0.14.0,<0.15.0 || >0.15.0,<0.15.1 || >0.15.1", optional = true, markers = "(sys_platform != \"win32\" and sys_platform != \"cygwin\") and platform_python_implementation != \"PyPy\" and extra == \"standard\""}
 watchfiles = {version = ">=0.13", optional = true, markers = "extra == \"standard\""}
 websockets = {version = ">=10.4", optional = true, markers = "extra == \"standard\""}
@@ -7869,10 +7871,9 @@ standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)",
 name = "uvloop"
 version = "0.21.0"
 description = "Fast implementation of asyncio event loop on top of libuv"
-optional = true
+optional = false
 python-versions = ">=3.8.0"
-groups = ["main"]
-markers = "platform_python_implementation != \"PyPy\" and (sys_platform != \"win32\" and sys_platform != \"cygwin\") and (python_version <= \"3.11\" or python_version >= \"3.12\")"
+groups = ["main", "dev"]
 files = [
     {file = "uvloop-0.21.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:ec7e6b09a6fdded42403182ab6b832b71f4edaf7f37a9a0e371a01db5f0cb45f"},
     {file = "uvloop-0.21.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:196274f2adb9689a289ad7d65700d37df0c0930fd8e4e743fa4834e850d7719d"},
@@ -7912,6 +7913,7 @@ files = [
     {file = "uvloop-0.21.0-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:2d1f581393673ce119355d56da84fe1dd9d2bb8b3d13ce792524e1607139feff"},
     {file = "uvloop-0.21.0.tar.gz", hash = "sha256:3bf12b0fda68447806a7ad847bfa591613177275d35b6724b1ee573faa3704e3"},
 ]
+markers = {main = "platform_python_implementation != \"PyPy\" and (sys_platform != \"win32\" and sys_platform != \"cygwin\") and (python_version <= \"3.11\" or python_version >= \"3.12\")", dev = "python_version <= \"3.11\" or python_version >= \"3.12\""}
 
 [package.extras]
 dev = ["Cython (>=3.0,<4.0)", "setuptools (>=60)"]
@@ -8608,4 +8610,4 @@ weaviate = ["weaviate-client"]
 [metadata]
 lock-version = "2.1"
 python-versions = ">=3.9,<3.13"
-content-hash = "a4c5bfe1833c997d05f962b27084053c22588989d7da70a8a1fa03b0b3a32692"
+content-hash = "d287a71722902352227cc39bdea54d2ff3071a0ae40b2af5231501d7c2476aa7"
diff --git a/pyproject.toml b/pyproject.toml
index 8d0b60d1..2bf368ff 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,9 +6,9 @@ build-backend = "setuptools.build_meta"
 # Do not add spaces around the '=' sign for any of the fields
 # preceeded by a marker comment as it affects the publish workflow.
 #replace_package_name_marker
-name = "dspy"
+name="dspy"
 #replace_package_version_marker
-version = "2.6.0rc8"
+version="2.6.0"
 description = "DSPy"
 readme = "README.md"
 authors = [{ name = "Omar Khattab", email = "okhattab@stanford.edu" }]
@@ -40,7 +40,7 @@ dependencies = [
     "pydantic~=2.0",
     "jinja2",
     "magicattr~=0.1.6",
-    "litellm==1.57.4",
+    "litellm>=1.59.8,<2.0.0",
     "diskcache",
     "json-repair",
     "tenacity>=8.2.3",
@@ -134,7 +134,7 @@ pgvector = { version = "^0.2.5", optional = true }
 llama-index = { version = "^0.10.30", optional = true }
 jinja2 = "^3.1.3"
 magicattr = "^0.1.6"
-litellm = { version = "==1.57.4", extras = ["proxy"] }
+litellm = { version = "^1.59.8", extras = ["proxy"] }
 diskcache = "^5.6.0"
 json-repair = "^0.30.0"
 tenacity = ">=8.2.3"


Test patch:
diff --git a/tests/predict/test_predict.py b/tests/predict/test_predict.py
index 828507db..0ae2c000 100644
--- a/tests/predict/test_predict.py
+++ b/tests/predict/test_predict.py
@@ -87,6 +87,65 @@ def test_demos_after_dump_and_load_state():
     # Demos don't need to keep the same types after saving and loading the state.
     assert new_instance.demos[0]["content"] == original_instance.demos[0].content
 
+def test_typed_demos_after_dump_and_load_state():
+    class Item(pydantic.BaseModel):
+        name: str
+        quantity: int
+
+    class InventorySignature(dspy.Signature):
+        """Handle inventory items and their translations."""
+        items: list[Item] = dspy.InputField()
+        language: str = dspy.InputField()
+        translated_items: list[Item] = dspy.OutputField()
+        total_quantity: int = dspy.OutputField()
+
+    original_instance = Predict(InventorySignature)
+    original_instance.demos = [
+        dspy.Example(
+            items=[
+                Item(name="apple", quantity=5),
+                Item(name="banana", quantity=3)
+            ],
+            language="SPANISH",
+            translated_items=[
+                Item(name="manzana", quantity=5),
+                Item(name="plátano", quantity=3)
+            ],
+            total_quantity=8
+        ).with_inputs("items", "language"),
+    ]
+
+    # Test dump_state
+    dumped_state = original_instance.dump_state()
+    assert len(dumped_state["demos"]) == len(original_instance.demos)
+    # Verify the input items were properly serialized
+    assert isinstance(dumped_state["demos"][0]["items"], list)
+    assert len(dumped_state["demos"][0]["items"]) == 2
+    assert dumped_state["demos"][0]["items"][0] == {"name": "apple", "quantity": 5}
+
+    # Test serialization/deserialization
+    saved_state = ujson.dumps(dumped_state)
+    loaded_state = ujson.loads(saved_state)
+
+    # Test load_state
+    new_instance = Predict(InventorySignature)
+    new_instance.load_state(loaded_state)
+    assert len(new_instance.demos) == len(original_instance.demos)
+
+    # Verify the structure is maintained after loading
+    loaded_demo = new_instance.demos[0]
+    assert isinstance(loaded_demo["items"], list)
+    assert len(loaded_demo["items"]) == 2
+    assert loaded_demo["items"][0]["name"] == "apple"
+    assert loaded_demo["items"][0]["quantity"] == 5
+    assert loaded_demo["items"][1]["name"] == "banana"
+    assert loaded_demo["items"][1]["quantity"] == 3
+
+    # Verify output items were also properly maintained
+    assert isinstance(loaded_demo["translated_items"], list)
+    assert len(loaded_demo["translated_items"]) == 2
+    assert loaded_demo["translated_items"][0]["name"] == "manzana"
+    assert loaded_demo["translated_items"][1]["name"] == "plátano"
 
 # def test_typed_demos_after_dump_and_load_state():
 #     class TypedTranslateToEnglish(dspy.Signature):