Make conversation simulator public and easily subclassable by smoorjani · Pull Request #20243 · mlflow/mlflow

smoorjani · 2026-01-23T05:23:34Z

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

This PR does two main things:

Makes the conversation simulator simulate method public as we have validated quality
Makes the simulator easily customizable by allowing easy subclassing of the user agent

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

from mlflow.genai.simulators import (
    BaseSimulatedUserAgent,
    ConversationSimulator,
    SimulatorContext,
)


class ImpatientUserAgent(BaseSimulatedUserAgent):
    def generate_message(self, context: SimulatorContext) -> str:
        if context.is_first_turn:
            return f"I need help RIGHT NOW with: {context.goal}"
        return self.invoke_llm(
            f"You are impatient. Respond briefly and push for faster answers. "
            f"Goal: {context.goal}. Last response: {context.last_assistant_response}"
        )


def predict_fn(input: list[dict], **kwargs) -> dict:
    return {"role": "assistant", "content": "I'll help you with that."}


simulator = ConversationSimulator(
    test_cases=[{"goal": "Reset my password"}],
    user_agent_class=ImpatientUserAgent,
    max_turns=3,
)

trace_ids = simulator.simulate(predict_fn)

result:

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

Make simulate method of ConversationSimulator public and make the user agent easily subclassable to customize conversation simulation..

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

github-actions · 2026-01-23T05:23:47Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20243/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/20243/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/20243/merge

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

github-actions · 2026-01-24T04:44:43Z

Documentation preview for 205a258 is available at:

https://pr-20243--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

xsh310 · 2026-01-26T16:39:06Z

mlflow/genai/simulators/simulator.py

    return "\n".join(formatted)


+def _fetch_traces(all_trace_ids: list[list[str]]) -> list[list["Trace"]]:


Would it be better to return traces instead of traceids in _run_conversation directly so that we can remove this transformation?

Good callout! The reason I do not is because we need to flush the traces (which are logged asynchronously) before retrieving the traces. Instead of doing this in multiple threads, it's much easier to flush all at once and then retrieve the traces.

xsh310 · 2026-01-26T16:39:10Z

mlflow/genai/simulators/simulator.py

+            The generated user message string.
+        """
+
+    def invoke_llm(self, prompt: str, system_prompt: str | None = None) -> str:


qq: is this provided in the abstract class for user's convenience or do we expect the users to use this for inference?

yes, just for convenience

xsh310

Overall LGTM

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

smoorjani · 2026-01-28T23:55:08Z

mlflow/genai/simulators/simulator.py

+
+    mlflow.flush_trace_async_logging()
+
+    client = TracingClient()


not to reviewer: the alternative to this code is using search traces with a specific bit of metadata, but the logic to reconstruct the order would be more complex than what is presently here. There is no performance difference between using search traces and parallelizing get trace

AveshCSingh

I took a quick pass and left a few small comments, but overall this LGTM. Please feel free to merge after addressing my comments.

mlflow/genai/simulators/simulator.py

AveshCSingh · 2026-01-29T00:34:19Z

mlflow/genai/simulators/simulator.py

                    progress_bar.close()

-        return all_trace_ids
+        return _fetch_traces(all_trace_ids)


tests/genai/simulators/test_simulator.py

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

Make simulator public and subclassable

67f36e7

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

github-actions bot added area/evaluation MLflow Evaluation rn/feature Mention under Features in Changelogs. labels Jan 23, 2026

smoorjani added 4 commits January 23, 2026 09:50

.

80dd690

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

merge with master

c62e5c5

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

.

5b389cc

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

.

d543190

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

smoorjani requested review from AveshCSingh and xsh310 January 23, 2026 21:31

smoorjani added 3 commits January 23, 2026 15:42

fix tests

fc44a82

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

self-review

cdfb5cb

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

.

18ef5be

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

xsh310 reviewed Jan 26, 2026

View reviewed changes

xsh310 approved these changes Jan 26, 2026

View reviewed changes

github-actions bot assigned xsh310 Jan 26, 2026

smoorjani added 2 commits January 27, 2026 12:31

use batch_get_traces with parallel fallback

5700523

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

simplify trace fetching to use parallel get_trace

46c7740

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

smoorjani commented Jan 28, 2026

View reviewed changes

AveshCSingh approved these changes Jan 29, 2026

View reviewed changes

github-actions bot assigned AveshCSingh Jan 29, 2026

address PR review comments

205a258

Signed-off-by: Samraj Moorjani <samraj.moorjani@databricks.com>

smoorjani enabled auto-merge January 29, 2026 01:35

smoorjani disabled auto-merge January 29, 2026 01:36

smoorjani added this pull request to the merge queue Jan 29, 2026

Merged via the queue into mlflow:master with commit c4c6c19 Jan 29, 2026
47 checks passed

smoorjani deleted the gwt-refactor-simulator branch January 29, 2026 03:02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Make conversation simulator public and easily subclassable#20243

Make conversation simulator public and easily subclassable#20243
smoorjani merged 11 commits intomlflow:masterfrom
smoorjani:gwt-refactor-simulator

smoorjani commented Jan 23, 2026

Uh oh!

github-actions bot commented Jan 23, 2026

Install mlflow from this PR

Uh oh!

github-actions bot commented Jan 24, 2026 •

edited

Loading

Uh oh!

xsh310 Jan 26, 2026

Uh oh!

smoorjani Jan 28, 2026

Uh oh!

xsh310 Jan 26, 2026

Uh oh!

smoorjani Jan 26, 2026

Uh oh!

xsh310 left a comment

Uh oh!

smoorjani Jan 28, 2026

Uh oh!

AveshCSingh left a comment

Uh oh!

Uh oh!

AveshCSingh Jan 29, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return "\n".join(formatted)


		def _fetch_traces(all_trace_ids: list[list[str]]) -> list[list["Trace"]]:

Conversation

smoorjani commented Jan 23, 2026

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

github-actions bot commented Jan 23, 2026

Install mlflow from this PR

Uh oh!

github-actions bot commented Jan 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

xsh310 Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

smoorjani Jan 28, 2026

Choose a reason for hiding this comment

Uh oh!

xsh310 Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

smoorjani Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

xsh310 left a comment

Choose a reason for hiding this comment

Uh oh!

smoorjani Jan 28, 2026

Choose a reason for hiding this comment

Uh oh!

AveshCSingh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

AveshCSingh Jan 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

github-actions bot commented Jan 24, 2026 •

edited

Loading