Rhoaieng 79828 implement dspy based prompt optimization

ai4rag/components/optimization/rag_templates_optimization.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 from typing import Any
 
+import dspy
 import pandas as pd
 from ogx_client import OgxClient
 
@@ -29,6 +30,7 @@
 from ai4rag.search_space.src.parameter import Parameter
 from ai4rag.search_space.src.search_space import AI4RAGSearchSpace
 from ai4rag.utils.event_handler.event_handler import KFPEventHandler
+from ai4rag.core.hpo.prompts import RAGPrompt
 
 _logger = logging.getLogger("rag-templates-optimization")
 _logger.addHandler(handler)
@@ -159,13 +161,20 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
     embedding_models: list[OGXEmbeddingModel] = []
     params: list[Parameter] = []
 
+    optimized_dspy_module = None
     for param_name, values in search_space_raw.items():
         if param_name == "foundation_model":
             values = [_deserialize_model(m, ogx_client) for m in values]
             foundation_models = values
         elif param_name == "embedding_model":
             values = [_deserialize_model(m, ogx_client) for m in values]
             embedding_models = values
+        elif param_name == "optimized_dspy_module":
+            if values:
+                optimized_dspy_module = dspy.ChainOfThought(RAGPrompt)
+                optimized_dspy_module.load_state(state=values)
+                print("Loaded learned state of DSPy CoT!")
+            continue
         params.append(Parameter(param_name, "C", values=values))
 
     search_space = AI4RAGSearchSpace(params=params)
@@ -193,6 +202,12 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
 
     event_handler = KFPEventHandler()
 
+    experiments_kwargs = {}
+
+    if optimized_dspy_module is not None:
+        experiments_kwargs["optimized_dspy_module"] = optimized_dspy_module
+        print("Passing the trained DSPy Module down to the AI4RAGExperiment class!")
+
     rag_exp = AI4RAGExperiment(
         client=ogx_client,
         event_handler=event_handler,
@@ -205,6 +220,7 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
         ogx_vector_io_provider_id=vector_io_provider_id,
         inference_max_threads=inference_max_threads,
         evaluators=evaluators,
+        **experiments_kwargs,
     )
 
     # --- Run the optimization loop ---

ai4rag/components/optimization/search_space_preparation.py

@@ -78,10 +78,13 @@ class SearchSpaceReport:
         model lists and non-model parameter ranges.
     selected_models : dict[str, list]
         Foundation and embedding model lists that survived pre-selection.
+    optimized_dspy_module : dict[str, Any]
+        A serialized representation of an optimized DSPy module as returned by `dump_state()` method.
     """
 
     search_space: dict[str, Any]
     selected_models: dict[str, list]
+    optimized_dspy_module: dict[str, Any]
 
     def save_json(self, path: str | Path) -> None:
         """Serialize the report to a JSON file.
@@ -95,8 +98,9 @@ def save_json(self, path: str | Path) -> None:
         """
         path = Path(path)
         path.parent.mkdir(parents=True, exist_ok=True)
+        report = {**self.search_space, "optimized_dspy_module": self.optimized_dspy_module}
         with open(path, "w", encoding="utf-8") as f:
-            json.dump(self.search_space, f, indent=2)
+            json.dump(report, f, indent=2)
 
 
 def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arguments,too-many-positional-arguments
@@ -113,6 +117,7 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
     chunk_sizes: list[int] | None = None,
     chunk_overlaps: list[int] | None = None,
     inference_max_threads: int = 10,
+    optimize_prompts: bool = True,
 ) -> SearchSpaceReport:
     """Run model pre-selection and prepare a search-space report.
 
@@ -161,6 +166,8 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
         RAG service during benchmark evaluation.  Lower values reduce
         per-request concurrency (useful when each request carries more
         retrieved context).  Defaults to ``10``.
+    optimize_prompts
+        Whether to utilise DSPy library for building and optimization of the prompts.
 
     Returns
     -------
@@ -222,14 +229,15 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
     fm_values = search_space["foundation_model"].values
     em_values = search_space["embedding_model"].values
 
+    mps = ModelsPreSelector(
+        benchmark_data=benchmark_data.get_random_sample(n_records=sample_size, random_seed=random_seed),
+        documents=documents,
+        foundation_models=search_space._search_space["foundation_model"].values,  # pylint: disable=protected-access
+        embedding_models=search_space._search_space["embedding_model"].values,  # pylint: disable=protected-access
+        max_threads=inference_max_threads,
+    )
+
     if len(fm_values) > top_n_generation or len(em_values) > top_k_embedding:
-        mps = ModelsPreSelector(
-            benchmark_data=benchmark_data.get_random_sample(n_records=sample_size, random_seed=random_seed),
-            documents=documents,
-            foundation_models=search_space._search_space["foundation_model"].values,  # pylint: disable=protected-access
-            embedding_models=search_space._search_space["embedding_model"].values,  # pylint: disable=protected-access
-            max_threads=inference_max_threads,
-        )
         mps.evaluate_patterns()
         selected = mps.select_models(
             n_embedding_models=top_k_embedding,
@@ -245,6 +253,10 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
             "embedding_model": list(em_values),
         }
 
+    optimized_dspy_module = None
+    if optimize_prompts:
+        optimized_dspy_module = mps.optimize_prompts(f"{ogx_client.base_url}/v1", ogx_client.api_key)
+
     # Build verbose representation from valid (rule-filtered) combinations only
     valid_combinations = search_space.combinations
     if not valid_combinations:
@@ -259,6 +271,7 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
     return SearchSpaceReport(
         search_space=verbose_repr,
         selected_models=selected_models,
+        optimized_dspy_module=optimized_dspy_module.dump_state() if optimized_dspy_module else {},
     )
 
 

ai4rag/core/experiment/experiment.py

@@ -6,9 +6,11 @@
 from dataclasses import asdict, is_dataclass
 from typing import Any, Sequence
 
+import dspy
 import pandas as pd
 from docling_core.types.doc import DoclingDocument
 from ogx_client import OgxClient
+from ai4rag.core.hpo.prompts import RAGPrompt
 
 from ai4rag import logger
 from ai4rag.core.experiment.benchmark_data import BenchmarkData
@@ -161,7 +163,7 @@ def __init__(
         self.n_mps_embedding_models = kwargs.pop("n_mps_embedding_models", ModelsPreSelector.DEFAULT_N_EMBEDDING_MODELS)
         self.known_observations: list[dict] | None = kwargs.pop("known_observations", None)
         self.inference_max_threads: int = kwargs.pop("inference_max_threads", 10)
-
+        self.optimized_dspy_module = kwargs.pop("optimized_dspy_module", None)
         self.results: ExperimentResults = ExperimentResults()
         self._exception_handler = ExperimentExceptionHandler(self.event_handler)
 
@@ -427,9 +429,24 @@ def run_single_evaluation(self, rag_params: RAGParamsType) -> float:
                 "Only 'vector' mode is supported for chroma."
             )
 
-        context_template_text = foundation_model.context_template_text
-        system_message_text = foundation_model.system_message_text
-        user_message_text = foundation_model.user_message_text
+        if self.optimized_dspy_module is not None:
+            chat_adapter = dspy.ChatAdapter(use_json_adapter_fallback=True)
+            optimized_dspy_module_signature = RAGPrompt.load_state(
+                self.optimized_dspy_module.dump_state()["predict"]["signature"]
+            )
+            context_template_text = "[[ ## contexts ## ]]"
+            system_message_text = chat_adapter.format_system_message(signature=optimized_dspy_module_signature)
+            user_message_text = chat_adapter.format_user_message_content(
+                signature=optimized_dspy_module_signature,
+                inputs={"question": "<user_question>", "contexts": "<contexts>"},
+            )
+            logger.info(
+                "Using the DSPy-optimized `context_template_text`, `system_message_text` and `user_message_text`."
+            )
+        else:
+            context_template_text = foundation_model.context_template_text
+            system_message_text = foundation_model.system_message_text
+            user_message_text = foundation_model.user_message_text
 
         rag_params = {
             "retrieval": retrieval_params,
@@ -528,10 +545,10 @@ def run_single_evaluation(self, rag_params: RAGParamsType) -> float:
             ranker_alpha=retrieval_params.get(AI4RAGParamNames.RANKER_ALPHA),
         )
 
-        rag_pattern = SimpleRAG(
-            foundation_model=foundation_model,
-            retriever=retriever,
-        )
+        rag_pattern_kwargs = {}
+        if self.optimized_dspy_module is not None:
+            rag_pattern_kwargs["optimized_dspy_module"] = self.optimized_dspy_module
+        rag_pattern = SimpleRAG(foundation_model=foundation_model, retriever=retriever, **rag_pattern_kwargs)
 
         _rag_log = (
             f"Retrieval and generation using collection: '{collection_name}' and "

ai4rag/core/experiment/mps.py

@@ -2,9 +2,11 @@
 # Copyright IBM Corp. 2025-2026
 # SPDX-License-Identifier: Apache-2.0
 # -----------------------------------------------------------------------------
+from random import choice, seed
 from typing import Any, TypedDict
 
 from docling_core.types.doc import DoclingDocument
+import dspy
 
 from ai4rag import logger
 from ai4rag.core.experiment.benchmark_data import BenchmarkData
@@ -15,6 +17,7 @@
     IndexingError,
 )
 from ai4rag.core.experiment.utils import build_evaluation_data, query_rag
+from ai4rag.core.hpo.prompts import RAGPrompt, overall_score_per_question
 from ai4rag.evaluator import UnitxtEvaluator
 from ai4rag.evaluator.base_evaluator import BaseEvaluator, EvaluationMetricsResult
 from ai4rag.evaluator.custom_metrics import apply_custom_metrics
@@ -180,6 +183,58 @@ def evaluate_patterns(self):
                 f"None of the given models has been successfully evaluated. {msg}"
             )
 
+    def optimize_prompts(self, api_base, api_key):
+        """Performs prompts optimization using the MIPROv2 algorithm and CoT module."""
+        # same model as MPS uses!
+        lm = dspy.LM(
+            f"openai/{self.foundation_models[0]}",
+            api_base=api_base,  ## must end with `/v1`
+            api_key=api_key,
+        )
+        dspy.configure(lm=lm)
+
+        cot = dspy.ChainOfThought(RAGPrompt)
+
+        document_ids = []
+        for element in self.benchmark_data.document_ids:
+            document_ids.extend(element)
+        documents = [document for document in self.documents if document.name in document_ids]
+        chunked_documents = self._chunk_documents(documents)
+        try:
+            vector_store = self._create_vector_store(
+                self.embedding_models[0], chunked_documents, collection_name="prompt_optimization_collection"
+            )
+        except Exception as exc:
+            raise IndexingError(exc, "prompt_optimization_collection", self.embedding_models[0].model_id) from exc
+
+        retriever = Retriever(vector_store, **self.retrieval_params)
+
+        dspy_train_set = []
+        dspy_test_set = []
+
+        # prepare dspy_train_set so that it contains the same questions as mps_train_set
+        for _, question_data in self.benchmark_data._benchmark_data.iterrows():
+            contexts = retriever.retrieve(question_data["question"])
+            dspy_example = dspy.Example(
+                question=question_data["question"],
+                contexts=contexts,
+                correct_answer=question_data["correct_answers"],
+                correct_answer_document_ids=question_data["correct_answer_document_ids"],
+            ).with_inputs("question", "contexts")
+            dspy_train_set.append(dspy_example)
+
+        seed(17)
+        test_question = choice(dspy_train_set)
+        dspy_test_set.append(test_question)
+        dspy_train_set.remove(test_question)
+
+        mipro = dspy.MIPROv2(metric=overall_score_per_question, auto="light", seed=17)
+        optimized_cot = mipro.compile(
+            cot, trainset=dspy_train_set, valset=dspy_test_set, seed=17, max_bootstrapped_demos=0, max_labeled_demos=0
+        )
+        # TODO save the chain's state
+        return optimized_cot
+
     def _evaluate_foundation_models(self, retriever: Retriever, embedding_model: BaseEmbeddingModel):
         """
         Evaluates each embedding model with given retriever.

ai4rag/core/hpo/prompts.py

@@ -0,0 +1,49 @@
+from unitxt.eval_utils import evaluate
+from statistics import fmean
+import dspy
+from ai4rag.rag.chunking.chunk import AI4RAGChunk
+
+
+def unitxt_metrics(example, prediction):
+    """Given an example and model's prediction calculates:
+    - answer correctness
+    - faifhtfulness
+    using the `unitxt.eval_utils.evaluate` function
+    """
+    eval_data = {
+        "question": example.question,
+        "answer": prediction.answer_grounded_in_contexts,
+        "contexts": [chunk.text for chunk in example.contexts],
+        "ground_truths": example.correct_answer,
+    }
+    scores, ci = evaluate(
+        [eval_data],
+        metric_names=[
+            "metrics.rag.external_rag.answer_correctness",
+            "metrics.rag.external_rag.faithfulness",
+        ],
+        compute_conf_intervals=True,
+    )
+    return scores, ci
+
+
+def overall_score_per_question(example, prediction, trace=None, pred_name=None, pred_trace=None):
+    """For a given question calculates arithmetic mean from the following unitxt metrics:
+    - answer_correctness
+    - faithfulness
+    - context_correctness
+    """
+    scores, ci = unitxt_metrics(example, prediction)
+
+    for external_rag_metrics in (
+        dict(filter(lambda items: "metrics.rag.external_rag" in items[0], per_question_score.items()))
+        for per_question_score in scores
+    ):
+        return round(fmean(external_rag_metrics.values()), 4)
+
+
+class RAGPrompt(dspy.Signature):
+
+    question: str = dspy.InputField()
+    contexts: list[AI4RAGChunk] = dspy.InputField()
+    answer_grounded_in_contexts: str = dspy.OutputField()

ai4rag/rag/template/simple_rag_template.py

@@ -14,6 +14,7 @@
 from ..embedding.base_model import BaseEmbeddingModel
 from ..foundation_models.base_model import BaseFoundationModel
 from .base_template import BaseRAGTemplate, RAGTemplateError
+import dspy
 
 
 class SimpleRAG(BaseRAGTemplate):
@@ -46,13 +47,15 @@ def __init__(
         chunker: BaseChunker | None = None,
         embedding_model: BaseEmbeddingModel | None = None,
         vector_store: BaseVectorStore | None = None,
+        **kwargs,
     ):
         super().__init__(
             foundation_model=foundation_model,
             retriever=retriever,
             embedding_model=embedding_model,
             vector_store=vector_store,
         )
+        self.optimized_dspy_module = kwargs.pop("optimized_dspy_module", None)
 
         self.chunker = chunker
 
@@ -110,10 +113,23 @@ def generate(self, question: str, **kwargs) -> dict[str, Any]:
             {"role": "user", "content": user_message},
         ]
 
-        chat_response = self.foundation_model.chat(messages=messages)
+        if self.optimized_dspy_module is not None:
+            with dspy.context(
+                lm=dspy.LM(
+                    f"openai/{self.foundation_model}",
+                    api_base=f"{self.foundation_model.client.base_url}/v1",
+                    api_key=self.foundation_model.client.api_key,
+                )
+            ):
+                model_prediction = self.optimized_dspy_module(question=question, contexts=reference_documents)
+            answer = model_prediction.answer_grounded_in_contexts
+            # TODO log model reasoning for a given question
+        else:
+            chat_response = self.foundation_model.chat(messages=messages)
+            answer = chat_response[0].message.content
 
         return {
-            "answer": chat_response[0].message.content,
+            "answer": answer,
             "reference_documents": reference_documents,
             "question": question,
         }

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
   "pygam~=0.12.0",
   "scikit-learn==1.8.*",
   "unitxt~=1.26.1",
+  "dspy[optuna]~=3.0"
 ]