feat: add ability to substitute media types for attachments #106

CONFIGURATION.md

@@ -781,6 +781,7 @@ Do not show parameter with value in stage while tool call:
 |---------------------------|----------|---------------|----------------------------------------------------------------------------------------------------------------------|----------------------------------------------|---------------|
 | supported_types           | No       | Array[String] | List of supported attachment MIME types                                                                              | `*/*`(all), `image/png`, `image/jpeg`, etc.. | `[*/*]`       |
 | propagate_types_to_choice | No       | Array[String] | List of supported attachment MIME types that will be shown in main chat (propagated from tool call result to choice) | `*/*`(all), `image/png`, `image/jpeg`, etc.. | `[]`          |
+| media_type_substitution   | No       | dict[str, str] | Maps original MIME type to substitute. Key is a original mime_type, value is desired mime type. | `*/*`(all), `image/png`, `image/jpeg`, etc.. | `{}` |
 
 <details>
 <summary><b>Parameter info configuration JSON sample</b></summary>
@@ -794,7 +795,10 @@ Do not show parameter with value in stage while tool call:
     "image/png",
     "image/jpeg",
     "application/vnd.plotly.v1+json"
-  ]
+  ],
+  "media_type_substitution": {
+    "application/json": "application/vnd.plotly.v1+json"
+  }
 }
 ```
 

docs/generated-app-schema.json

@@ -61,6 +61,14 @@
           },
           "title": "Propagate Types To Choice",
           "type": "array"
+        },
+        "media_type_substitution": {
+          "additionalProperties": {
+            "type": "string"
+          },
+          "description": "Maps original attachment types to substitution types for custom visualizers.",
+          "title": "Media Type Substitution",
+          "type": "object"
         }
       },
       "title": "AttachmentConfig",

src/quickapp/common/staged_base_tool.py

@@ -14,7 +14,7 @@
 from .exceptions import InvalidToolCallParameterException
 from .perf_timer.perf_timer import PerformanceTimer
 from .tool_fallback.processor import FallbackProcessor
-from .utils import matches_type
+from .utils import matches_type, substitute_media_type
 
 logger = logging.getLogger(__name__)
 
@@ -114,6 +114,9 @@ async def _run_in_stage_report_success(
                 filtered: list = []
                 for a in result.attachments:
                     if matches_type(a.type, attachment_cfg.supported_types):
+                        a.type = substitute_media_type(
+                            a.type, attachment_cfg.media_type_substitution
+                        )
                         filtered.append(a)
                         if matches_type(a.type, attachment_cfg.propagate_types_to_choice):
                             result.propagate_to_choice.append(a)

src/quickapp/common/utils.py

@@ -32,6 +32,16 @@ def matches_type(mime_type: str | None, allowed_mime_types: list[str] | None) ->
     return False
 
 
+def substitute_media_type(
+    mime_type: str | None,
+    mapping: dict[str, str],
+) -> str | None:
+    """Apply media type substitution"""
+    if mime_type is None:
+        return None
+    return mapping.get(mime_type, mime_type)
+
+
 def sanitize_toolname(input_str: str) -> str:
     """
     Sanitizes a string to match the pattern ^[a-zA-Z0-9_-]{1,64}$

src/quickapp/config/tools/base.py

@@ -173,6 +173,10 @@ class AttachmentConfig(BaseModel):
         default_factory=lambda: DEFAULT_PROPAGATE_TO_CHOICE,
         description="List of attachment types to propagate from stage to choice.",
     )
+    media_type_substitution: dict[str, str] = Field(
+        default_factory=dict,
+        description="Maps original attachment types to substitution types for custom visualizers.",
+    )
 
     @field_validator("supported_types", mode="before")
     @classmethod
@@ -188,6 +192,13 @@ def coerce_none_propagate_types(cls, v: list[str] | None) -> list[str]:
             return DEFAULT_PROPAGATE_TO_CHOICE
         return v
 
+    @field_validator("media_type_substitution", mode="before")
+    @classmethod
+    def coerce_none_type_substitution(cls, v: dict[str, str] | None) -> dict[str, str]:
+        if v is None:
+            return {}
+        return v
+
 
 class OpenAiToolConfig(
     BaseModel,

src/tests/unit_tests/common/test_staged_base_tool.py

@@ -126,3 +126,108 @@ async def test_propagation_only_for_surviving_attachments(mock_stage_wrapper_fac
     # so it should NOT appear in propagate_to_choice
     assert len(result.propagate_to_choice) == 1
     assert result.propagate_to_choice[0].type == "image/png"
+
+
+@pytest.mark.asyncio
+async def test_media_type_substitution_applied(mock_stage_wrapper_factory):
+    """Attachments that pass supported_types should have their type substituted
+    according to media_type_substitution mapping."""
+    mock_config = Mock()
+    mock_config.display = None
+    mock_config.fallback_configuration = ToolFallbackConfig(display_error_in_stage=True)
+    mock_config.attachment = AttachmentConfig(
+        supported_types=["image/*"],
+        propagate_types_to_choice=["image/*"],
+        media_type_substitution={"image/png": "image/webp"},
+    )
+
+    attachment = Attachment(type="image/png", title="photo.png", data="img_data")
+
+    result_to_return = CompletionResult(
+        content="result",
+        content_type="text/plain",
+        attachments=[attachment],
+    )
+
+    tool = CustomTestStagedBaseTool(
+        stage_wrapper_builder=mock_stage_wrapper_factory,
+        tool_config=mock_config,
+        perf_timer=Mock(),
+        result_to_return=result_to_return,
+    )
+
+    result = await tool.arun("call-1")
+
+    assert len(result.attachments) == 1
+    assert result.attachments[0].type == "image/webp"
+
+
+@pytest.mark.asyncio
+async def test_media_type_substitution_not_applied_when_no_match(mock_stage_wrapper_factory):
+    """Attachments whose type is not in the substitution mapping keep their original type."""
+    mock_config = Mock()
+    mock_config.display = None
+    mock_config.fallback_configuration = ToolFallbackConfig(display_error_in_stage=True)
+    mock_config.attachment = AttachmentConfig(
+        supported_types=["image/*"],
+        propagate_types_to_choice=["image/*"],
+        media_type_substitution={"image/png": "image/webp"},
+    )
+
+    attachment = Attachment(type="image/jpeg", title="photo.jpg", data="img_data")
+
+    result_to_return = CompletionResult(
+        content="result",
+        content_type="text/plain",
+        attachments=[attachment],
+    )
+
+    tool = CustomTestStagedBaseTool(
+        stage_wrapper_builder=mock_stage_wrapper_factory,
+        tool_config=mock_config,
+        perf_timer=Mock(),
+        result_to_return=result_to_return,
+    )
+
+    result = await tool.arun("call-1")
+
+    assert len(result.attachments) == 1
+    assert result.attachments[0].type == "image/jpeg"
+
+
+@pytest.mark.asyncio
+async def test_propagation_uses_substituted_type(mock_stage_wrapper_factory):
+    """After substitution, propagate_types_to_choice should match against the NEW type."""
+    mock_config = Mock()
+    mock_config.display = None
+    mock_config.fallback_configuration = ToolFallbackConfig(display_error_in_stage=True)
+    mock_config.attachment = AttachmentConfig(
+        supported_types=["image/*"],
+        propagate_types_to_choice=["application/custom"],
+        media_type_substitution={"image/png": "application/custom"},
+    )
+
+    attachment = Attachment(type="image/png", title="chart.png", data="img_data")
+
+    result_to_return = CompletionResult(
+        content="result",
+        content_type="text/plain",
+        attachments=[attachment],
+    )
+
+    tool = CustomTestStagedBaseTool(
+        stage_wrapper_builder=mock_stage_wrapper_factory,
+        tool_config=mock_config,
+        perf_timer=Mock(),
+        result_to_return=result_to_return,
+    )
+
+    result = await tool.arun("call-1")
+
+    # The attachment survives (image/png passes supported_types=["image/*"])
+    assert len(result.attachments) == 1
+    # Its type was substituted
+    assert result.attachments[0].type == "application/custom"
+    # Propagation check uses the substituted type, which matches propagate_types_to_choice
+    assert len(result.propagate_to_choice) == 1
+    assert result.propagate_to_choice[0].type == "application/custom"

src/tests/unit_tests/common/test_substitute_media_type.py

@@ -0,0 +1,24 @@
+from quickapp.common.utils import substitute_media_type
+
+
+class TestSubstituteMediaType:
+    def test_none_mime_type_returns_none(self):
+        assert substitute_media_type(None, {"image/png": "image/webp"}) is None
+
+    def test_empty_mapping_returns_original(self):
+        assert substitute_media_type("image/png", {}) == "image/png"
+
+    def test_no_match_returns_original(self):
+        assert substitute_media_type("image/png", {"text/plain": "text/html"}) == "image/png"
+
+    def test_exact_match_substitutes(self):
+        assert substitute_media_type("image/png", {"image/png": "image/webp"}) == "image/webp"
+
+    def test_multiple_mappings(self):
+        mapping = {
+            "image/png": "image/webp",
+            "text/plain": "text/html",
+        }
+        assert substitute_media_type("image/png", mapping) == "image/webp"
+        assert substitute_media_type("text/plain", mapping) == "text/html"
+        assert substitute_media_type("application/pdf", mapping) == "application/pdf"