Repository: oumi-ai/oumi
Pull request: Updated oumi infer to support CLI argument for system prompt
# Description
This change adds an optional argument for ```system_prompt``` in the CLI for the interactive mode. The two tests have been added that test this change with and without image as an input.
<!--
Thank you for contributing to Oumi! Before sending your PR out for review, please take a quick read through this template.

When your PR is merged, its title will appear in our release notes. Make sure your title gives a clear description of your change!

After you've updated your title, please replace this section with a detailed description of your change. Include as much context as possible so your reviewers can easily understand *what* you're changing and *why*.
The more information you provide, the faster we can review your change!
-->
<!--↓↓↓↓↓↓↓↓↓↓ Describe your change below ↓↓↓↓↓↓↓↓↓↓-->
The change specifically adds a new argument in ```infer``` and ```infer_interactive``` functions. It also adds new ```Message``` object in the ```Conversation``` object for inputs with and without an image input.

Further, since this added new code, tests have been added to ensure their functioning. Two tests have been added with and without image input specifically to decouple the testing of system prompt feature independent of the image input for future triages!
<!--↑↑↑↑↑↑↑↑↑↑ Describe your change above ↑↑↑↑↑↑↑↑↑↑-->

## Related issues

<!--
Make sure to list any relevant related issues to your change. More often than not this will be the single issue fixed by your PR.
-->
<!--↓↓↓↓↓↓↓↓↓↓ List your related issues below ↓↓↓↓↓↓↓↓↓↓-->

Fixes #1386 

<!--↑↑↑↑↑↑↑↑↑↑ List your related issues above ↑↑↑↑↑↑↑↑↑↑-->

## Before submitting

- [ ] This PR only changes documentation. (You can ignore the following checks in that case)
- [x] Did you read the [contributor guideline](https://github.com/oumi-ai/oumi/blob/main/CONTRIBUTING.md) Pull Request guidelines?
- [x] Did you link the issue(s) related to this PR in the section above?
- [x] Did you add / update tests where needed?

## Reviewers

At least one review from a member of `oumi-ai/oumi-staff` is required.

<!-- Add `oumi-ai/oumi-staff` as a reviewer when your PR is ready for review.

You are also welcome to add individual members of `oumi-ai/oumi-staff` as reviewers.

If no one has reviewed your PR after several days, feel free to add a comment tagging specific reviewers.

 -->


Problem statement:
[Feature] Update `oumi infer` to support CLI argument for system prompt
### Feature request

Add a new command line argument for `oumi infer` command



### Motivation / references

Users may want to adjust system prompt for task-specific  instructions

### Your contribution

I can help with code review, and can answer questions


Towards OPE-1019

Hints from issue discussion:
@xrdaukar hello! I would like to work on this. I do understand where and how I will add the code for accepting a system-prompt, but then where am I supposed to use this input? I suppose it shall be a part of the ```InferenceConfig```?

Thanks.
Thank you @Spaarsh !   I think, at minimum, we can add a new `system_prompt` parameter to `infer()` CLI function https://github.com/oumi-ai/oumi/blob/4ad8780ab2f795171e7f5e99b9c0048a384f62d0/src/oumi/cli/infer.py#L37. (similarly to `image` param, which provides visual context for multimodal models), then modify this function  https://github.com/oumi-ai/oumi/blob/4ad8780ab2f795171e7f5e99b9c0048a384f62d0/src/oumi/infer.py#L101 to insert `Role.SYSTEM` message to the conversation (as the first element).  It's OK to only support it for `interactive` mode initially. 

We could also consider adding system prompt param to `InferenceConfig` or `GenerationParams` but it'd be more involved and require more changes, so my suggestion is to do the minimal change as described above as the first step. @taenin fyi

+1 to @xrdaukar 's first suggestion. I think adding support just for interactive mode is a good first step. If we find folks need this for non-interactive inference we can circle back and address it.

Commit message: Merge branch 'main' into 1386/enhancement/add-system-prompt-argument-in-oumi-infer

Patch to review:
diff --git a/src/oumi/__init__.py b/src/oumi/__init__.py
index cf70b6580..e9eef515d 100644
--- a/src/oumi/__init__.py
+++ b/src/oumi/__init__.py
@@ -124,12 +124,17 @@ def evaluate(config: EvaluationConfig) -> list[dict[str, Any]]:
 
 
 def infer_interactive(
-    config: InferenceConfig, *, input_image_bytes: bytes | None = None
+    config: InferenceConfig,
+    *,
+    input_image_bytes: bytes | None = None,
+    system_prompt: str | None = None,
 ) -> None:
     """Interactively provide the model response for a user-provided input."""
     import oumi.infer
 
-    return oumi.infer.infer_interactive(config, input_image_bytes=input_image_bytes)
+    return oumi.infer.infer_interactive(
+        config, input_image_bytes=input_image_bytes, system_prompt=system_prompt
+    )
 
 
 def infer(
diff --git a/src/oumi/cli/infer.py b/src/oumi/cli/infer.py
index ac5327f21..b640ecead 100644
--- a/src/oumi/cli/infer.py
+++ b/src/oumi/cli/infer.py
@@ -44,6 +44,16 @@ def infer(
             ),
         ),
     ] = None,
+    system_prompt: Annotated[
+        Optional[str],
+        typer.Option(
+            "--system-prompt",
+            help=(
+                "System prompt for task-specific instructions. "
+                "Only used in interactive mode."
+            ),
+        ),
+    ] = None,
     level: cli_utils.LOG_LEVEL_TYPE = None,
 ):
     """Run inference on a model.
@@ -57,6 +67,7 @@ def infer(
         config: Path to the configuration file for inference.
         interactive: Whether to run in an interactive session.
         image: Path to the input image for `image+text` VLLMs.
+        system_prompt: System prompt for task-specific instructions.
         level: The logging level for the specified command.
     """
     extra_args = cli_utils.parse_extra_cli_args(ctx)
@@ -93,9 +104,10 @@ def infer(
                 "`input_path`."
             )
         return oumi_infer_interactive(
-            parsed_config, input_image_bytes=input_image_png_bytes
+            parsed_config,
+            input_image_bytes=input_image_png_bytes,
+            system_prompt=system_prompt,
         )
-
     if parsed_config.input_path is None:
         raise ValueError("One of `--interactive` or `input_path` must be provided.")
     generations = oumi_infer(parsed_config)
diff --git a/src/oumi/infer.py b/src/oumi/infer.py
index 9c1c8c3f4..ef25ec0ed 100644
--- a/src/oumi/infer.py
+++ b/src/oumi/infer.py
@@ -41,7 +41,10 @@ def _get_engine(config: InferenceConfig) -> BaseInferenceEngine:
 
 
 def infer_interactive(
-    config: InferenceConfig, *, input_image_bytes: Optional[bytes] = None
+    config: InferenceConfig,
+    *,
+    input_image_bytes: Optional[bytes] = None,
+    system_prompt: Optional[str] = None,
 ) -> None:
     """Interactively provide the model response for a user-provided input."""
     # Create engine up front to avoid reinitializing it for each input.
@@ -57,6 +60,7 @@ def infer_interactive(
             inputs=[
                 input_text,
             ],
+            system_prompt=system_prompt,
             input_image_bytes=input_image_bytes,
             inference_engine=inference_engine,
         )
@@ -71,6 +75,7 @@ def infer(
     config: InferenceConfig,
     inputs: Optional[list[str]] = None,
     inference_engine: Optional[BaseInferenceEngine] = None,
+    system_prompt: Optional[str] = None,
     *,
     input_image_bytes: Optional[bytes] = None,
 ) -> list[Conversation]:
@@ -81,6 +86,7 @@ def infer(
         inputs: A list of inputs for inference.
         inference_engine: The engine to use for inference. If unspecified, the engine
             will be inferred from `config`.
+        system_prompt: System prompt for task-specific instructions.
         input_image_bytes: An input PNG image bytes to be used with `image+text` VLLMs.
             Only used in interactive mode.
 
@@ -93,15 +99,21 @@ def infer(
     # Pass None if no conversations are provided.
     conversations = None
     if inputs is not None and len(inputs) > 0:
+        base_message = []
+        if system_prompt:
+            base_message.append(Message(role=Role.SYSTEM, content=system_prompt))
         if input_image_bytes is None:
             conversations = [
-                Conversation(messages=[Message(role=Role.USER, content=content)])
+                Conversation(
+                    messages=base_message + [Message(role=Role.USER, content=content)]
+                )
                 for content in inputs
             ]
         else:
             conversations = [
                 Conversation(
-                    messages=[
+                    messages=base_message
+                    + [
                         Message(
                             role=Role.USER,
                             content=[
diff --git a/tests/unit/cli/test_cli_infer.py b/tests/unit/cli/test_cli_infer.py
index f5be3b6a4..453cc7504 100644
--- a/tests/unit/cli/test_cli_infer.py
+++ b/tests/unit/cli/test_cli_infer.py
@@ -161,3 +161,64 @@ def test_infer_logging_levels(app, mock_infer, mock_infer_interactive):
         assert logger.level == logging.WARNING
         _ = runner.invoke(app, ["-i", "--config", yaml_path, "-log", "CRITICAL"])
         assert logger.level == logging.CRITICAL
+
+
+def test_infer_with_system_prompt(app, mock_infer_interactive):
+    with tempfile.TemporaryDirectory() as output_temp_dir:
+        yaml_path = str(Path(output_temp_dir) / "infer.yaml")
+
+        config: InferenceConfig = _create_inference_config()
+        config.to_yaml(yaml_path)
+
+        # Test with interactive mode and system prompt
+        result = runner.invoke(
+            app,
+            [
+                "-i",
+                "--config",
+                yaml_path,
+                "--system-prompt",
+                "You are a helpful assistant",
+            ],
+        )
+        assert result.exit_code == 0
+        mock_infer_interactive.assert_called_once_with(
+            config, system_prompt="You are a helpful assistant", input_image_bytes=None
+        )
+        mock_infer_interactive.reset_mock()
+
+
+def test_infer_with_system_prompt_and_image(app, mock_infer_interactive):
+    with tempfile.TemporaryDirectory() as output_temp_dir:
+        yaml_path = str(Path(output_temp_dir) / "infer.yaml")
+
+        config: InferenceConfig = _create_inference_config()
+        config.to_yaml(yaml_path)
+
+        test_image = PIL.Image.new(mode="RGB", size=(32, 16))
+        temp_io_output = io.BytesIO()
+        test_image.save(temp_io_output, format="PNG")
+        image_bytes = temp_io_output.getvalue()
+
+        image_path = Path(output_temp_dir) / "test_image.png"
+        with image_path.open(mode="wb") as f:
+            f.write(image_bytes)
+
+        result = runner.invoke(
+            app,
+            [
+                "-i",
+                "--config",
+                yaml_path,
+                "--system-prompt",
+                "You are a helpful assistant",
+                "--image",
+                str(image_path),
+            ],
+        )
+        assert result.exit_code == 0
+        mock_infer_interactive.assert_called_once_with(
+            config,
+            system_prompt="You are a helpful assistant",
+            input_image_bytes=image_bytes,
+        )


Code review comments:
diff_hunk: @@ -93,15 +99,21 @@ def infer(
     # Pass None if no conversations are provided.
     conversations = None
     if inputs is not None and len(inputs) > 0:
+        base_message = []
path: src/oumi/infer.py
text: @user1:
nit: rename base_message to system_messages