mlflow-based prompts optimisation

ai4rag/components/optimization/search_space_preparation.py

@@ -8,16 +8,26 @@
 from pathlib import Path
 from typing import Any
 
+import mlflow
 import pandas as pd
 from ogx_client import OgxClient
+from mlflow.genai.optimize.optimizers import GepaPromptOptimizer
+from mlflow.genai.optimize import MetaPromptOptimizer
 
 from ai4rag import handler
 from ai4rag.components.utils.docling_io import load_docling_documents
 from ai4rag.core.experiment.benchmark_data import BenchmarkData
 from ai4rag.core.experiment.mps import ModelsPreSelector
+from ai4rag.rag.chunking.langchain_chunker import LangChainChunker
 from ai4rag.rag.embedding.base_model import BaseEmbeddingModel
 from ai4rag.rag.foundation_models.base_model import BaseFoundationModel
 from ai4rag.search_space.prepare.prepare_search_space import prepare_search_space_with_ogx
+from ai4rag.rag.retrieval.retriever import Retriever
+from ai4rag.search_space.src.model_props import (
+    QUESTION_PLACEHOLDER,
+    REFERENCE_DOCUMENTS_PLACEHOLDER,
+)
+import os
 
 _logger = logging.getLogger("search-space-preparation")
 _logger.addHandler(handler)
@@ -67,6 +77,166 @@ def _serialize_model(model: BaseFoundationModel | BaseEmbeddingModel) -> dict[st
     return result
 
 
+def precompute_retrieval_context(
+    benchmark_sample: BenchmarkData,
+    retriever: Retriever,
+    context_template_text: str,
+) -> list[dict[str, Any]]:
+    """Build mlflow-compatible training data with pre-computed retrieval context.
+
+    For each benchmark question, retrieves relevant chunks via *retriever*,
+    formats them into a ``reference_documents`` string using
+    *context_template_text*, and returns the result as a list of dicts
+    suitable for ``mlflow.genai.optimize_prompts(train_data=...)``.
+
+    Parameters
+    ----------
+    benchmark_sample
+        A random sample of the benchmark data (questions + expected answers).
+    retriever
+        A fully initialised retriever backed by an indexed vector store.
+    context_template_text
+        Per-document formatting template (e.g. ``"Document {doc_number}:\\n{document}"``).
+
+    Returns
+    -------
+    list[dict[str, Any]]
+        Each element has ``{"inputs": {"question": ..., "reference_documents": ...}, "outputs": ...}``.
+    """
+    train_data = []
+    for question, correct_answers in zip(benchmark_sample.questions, benchmark_sample.correct_answers):
+        chunks = retriever.retrieve(question)
+        reference_documents = "\n\n".join(
+            context_template_text.format(document=chunk.text, doc_number=i) for i, chunk in enumerate(chunks, start=1)
+        )
+        train_data.append(
+            {
+                "inputs": {
+                    QUESTION_PLACEHOLDER: question,
+                    REFERENCE_DOCUMENTS_PLACEHOLDER: reference_documents,
+                },
+                "outputs": correct_answers[0] if isinstance(correct_answers, list) else correct_answers,
+            }
+        )
+    return train_data
+
+
+def optimize_prompts_for_models(
+    foundation_models: list[BaseFoundationModel],
+    train_data: list[dict[str, Any]],
+    reflection_model: str,
+    maas_client: OgxClient,
+) -> None:
+    """Optimize system/user prompts for each foundation model via mlflow.
+
+    Registers each model's prompts as an mlflow chat prompt, runs
+    ``mlflow.genai.optimize_prompts`` with ``GepaPromptOptimizer``, and
+    mutates the model objects in-place with the optimised prompt texts.
+
+    Parameters
+    ----------
+    foundation_models
+        Foundation models whose prompts should be optimised.
+    train_data
+        Pre-computed training data as returned by
+        :func:`precompute_retrieval_context`.
+    reflection_model
+        Model URI for the GEPA reflection LLM
+        (e.g. ``"openai/ibm/granite-3-8b-instruct"``).
+    maas_client
+        An authenticated MaaS :class:`~openai.OpenAI` client.
+    """
+    original_tracking_uri = mlflow.get_tracking_uri()
+    mlflow.set_tracking_uri("sqlite://")
+
+    optimizer = GepaPromptOptimizer(reflection_model=reflection_model)
+
+    try:
+        for model in foundation_models:
+            prompt_name = model.model_id.replace("/", "-")
+
+            mlflow_user_message = model.user_message_text
+            for placeholder in (QUESTION_PLACEHOLDER, REFERENCE_DOCUMENTS_PLACEHOLDER):
+                mlflow_user_message = mlflow_user_message.replace(f"{{{placeholder}}}", f"{{{{{placeholder}}}}}")
+
+            prompt = mlflow.genai.register_prompt(
+                name=prompt_name,
+                template=[
+                    {"role": "system", "content": model.system_message_text},
+                    {"role": "user", "content": mlflow_user_message},
+                ],
+            )
+
+            def make_predict_fn(fm: BaseFoundationModel, uri: str):
+                def predict_fn(**kwargs) -> str:
+                    loaded = mlflow.genai.load_prompt(uri)
+                    messages = loaded.format(**kwargs)
+                    response = fm.chat(messages=messages)
+                    return response[0]["content"]
+
+                return predict_fn
+
+            result = mlflow.genai.optimize_prompts(
+                predict_fn=make_predict_fn(model, prompt.uri),
+                train_data=train_data,
+                prompt_uris=[prompt.uri],
+                optimizer=optimizer,
+                enable_tracking=False,
+            )
+
+            optimized_template = result.optimized_prompts[0].to_single_brace_format()
+            for msg in optimized_template:
+                if msg["role"] == "system":
+                    model.system_message_text = msg["content"]
+                elif msg["role"] == "user":
+                    model.user_message_text = msg["content"]
+
+            _logger.info(
+                "Optimised prompts for model '%s' (score: %s -> %s).",
+                model.model_id,
+                result.initial_eval_score,
+                result.final_eval_score,
+            )
+    finally:
+        mlflow.set_tracking_uri(original_tracking_uri)
+
+
+def outer_predict_function(fm, retriever):
+
+    def mlflow_predict_function(question: str):
+        system_prompt = mlflow.load_prompt("default_system")
+        user_prompt = mlflow.load_prompt("default_user")
+        context_prompt = mlflow.load_prompt("default_context")
+
+        # extract ref docs and create `context`
+        reference_documents = retriever.retrieve(question)
+
+        contexts = []
+        for numb, chunk in enumerate(reference_documents, start=1):
+            contexts.append(context_prompt.format(doc_number=numb, document=chunk))
+
+        # contexts = "\n\n".join(
+        #     [
+        #         context_prompt.format(doc_number=numb, document=chunk)
+        #         for numb, chunk in enumerate(reference_documents, start=1)
+        #     ]
+        # )
+
+        msgs = [
+            {"role": "system", "content": system_prompt.format()},
+            {
+                "role": "user",
+                "content": user_prompt.format(question=question, reference_documents="\n\n".join(contexts)),
+            },
+        ]
+
+        res = fm.chat(messages=msgs)
+
+        return res[0].message.content
+
+    return mlflow_predict_function
+
+
 @dataclass
 class SearchSpaceReport:
     """Result of the search-space preparation step.
@@ -113,6 +283,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 +332,10 @@ 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``.
+    reflection_model
+        Model URI for the GEPA reflection LLM used during prompt
+        optimisation (e.g. ``"openai/ibm/granite-3-8b-instruct"``).
+        When ``None`` (default), prompt optimisation is skipped.
 
     Returns
     -------
@@ -222,6 +397,7 @@ 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 | None = None
     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),
@@ -245,6 +421,72 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
             "embedding_model": list(em_values),
         }
 
+    if optimize_prompts:
+        os.environ["OPENAI_API_BASE"] = f"{ogx_client.base_url}/v1"
+        os.environ["OPENAI_API_KEY"] = ogx_client.api_key
+
+        selected_fm = selected_models["foundation_model"]
+        selected_em = selected_models["embedding_model"]
+
+        if mps is not None:
+            retriever = mps.retrievers[selected_em[0].model_id]
+        else:
+            document_ids = []
+            sample_for_docs = benchmark_data.get_random_sample(n_records=sample_size, random_seed=random_seed)
+            for element in sample_for_docs.document_ids:
+                document_ids.extend(element)
+            sample_docs = [d for d in documents if d.name in document_ids]
+            chunker = LangChainChunker(chunk_size=512, chunk_overlap=128, method="recursive")
+            chunks = chunker.split_documents(sample_docs)
+            vector_store = ModelsPreSelector._create_vector_store(
+                selected_em[0], chunks, collection_name="ai4rag_prompt_opt"
+            )
+            retriever = Retriever(vector_store, number_of_chunks=3, method="simple", search_mode="vector")
+
+        sample = benchmark_data.get_random_sample(n_records=sample_size, random_seed=random_seed)
+        # prompt_train_data = precompute_retrieval_context(
+        #     benchmark_sample=sample,
+        #     retriever=retriever,
+        #     context_template_text=selected_fm[0].context_template_text,
+        # )
+
+        mlflow.set_tracking_uri("sqlite:///mlflow.db")
+
+        user_prompt = mlflow.genai.register_prompt(
+            name="default_user",
+            template=selected_fm[0].user_message_text,
+        )
+
+        system_prompt = mlflow.genai.register_prompt(
+            name="default_system",
+            template=selected_fm[0].system_message_text,
+        )
+
+        context_template = mlflow.genai.register_prompt(
+            name="default_context", template=selected_fm[0].context_template_text
+        )
+
+        meta_opt = MetaPromptOptimizer(reflection_model=f"openai:/{selected_fm[0].model_id}")
+        result = mlflow.genai.optimize_prompts(
+            predict_fn=outer_predict_function(fm=selected_fm[0], retriever=None),
+            prompt_uris=[user_prompt.uri, system_prompt.uri, context_template.uri],
+            optimizer=meta_opt,
+            train_data=[],
+            enable_tracking=False,
+            scorers=[],
+        )
+
+        selected_fm[0].user_message_text = result.optimized_prompts[0].template
+        selected_fm[0].system_message_text = result.optimized_prompts[1].template
+        selected_fm[0].context_template_text = result.optimized_prompts[2].template
+
+        # optimize_prompts_for_models(
+        #     foundation_models=selected_fm,
+        #     train_data=prompt_train_data,
+        #     reflection_model=reflection_model,
+        #     maas_client=maas_client,
+        # )
+
     # Build verbose representation from valid (rule-filtered) combinations only
     valid_combinations = search_space.combinations
     if not valid_combinations:

ai4rag/core/experiment/mps.py

@@ -127,6 +127,7 @@ def __init__(
             "chunk_overlap": kwargs.get(AI4RAGParamNames.CHUNK_OVERLAP, 128),
         }
         self.evaluation_results: list[MPSEvaluationResultsTyped] = []
+        self.retrievers: dict[str, Retriever] = {}
         self._exception_handler = ExperimentExceptionHandler()
         self.max_threads = kwargs.pop("max_threads", 10)
         self._unitxt_metrics = tuple(m for m in Metrics if m.evaluator == "unitxt")
@@ -166,6 +167,7 @@ def evaluate_patterns(self):
                     raise IndexingError(exc, collection_name, embedding_model.model_id) from exc
 
                 retriever = Retriever(vector_store, **self.retrieval_params)
+                self.retrievers[embedding_model.model_id] = retriever
                 self._evaluate_foundation_models(retriever=retriever, embedding_model=embedding_model)
 
             except IndexingError as exc:

pyproject.toml

@@ -39,6 +39,8 @@ dependencies = [
   "pygam~=0.12.0",
   "scikit-learn==1.8.*",
   "unitxt~=1.26.1",
+  "gepa>=0.0.26",
+  "mlflow>=3.5.0",  # for prompt optimize API support
 ]