From 4cca8584f8c13e2ebefe69617309d229b3531dcb Mon Sep 17 00:00:00 2001 From: Filip Komarzyniec Date: Fri, 21 Aug 2026 00:53:30 +0200 Subject: [PATCH] mlflow-based prompts optimisation Signed-off-by: Filip Komarzyniec rh-pre-commit.version: 2.3.2 rh-pre-commit.check-secrets: ENABLED --- .../optimization/search_space_preparation.py | 242 ++++++++++++++++++ ai4rag/core/experiment/mps.py | 2 + pyproject.toml | 2 + 3 files changed, 246 insertions(+) diff --git a/ai4rag/components/optimization/search_space_preparation.py b/ai4rag/components/optimization/search_space_preparation.py index e4738287..eee54002 100644 --- a/ai4rag/components/optimization/search_space_preparation.py +++ b/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: diff --git a/ai4rag/core/experiment/mps.py b/ai4rag/core/experiment/mps.py index 50bd6848..e5c97aec 100644 --- a/ai4rag/core/experiment/mps.py +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index 6892d7ca..866e2f2f 100644 --- a/pyproject.toml +++ b/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 ]