feat: pass image file to llm (#39848)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-09-24 23:22:26 +08:00 · 2026-08-03 06:04:32 +00:00
parent 982c0667e1
commit dc8dd09450
3 changed files with 164 additions and 3 deletions
@@ -16,6 +16,7 @@ from graphon.model_runtime.entities.llm_entities import LLMUsage
 from graphon.model_runtime.entities.message_entities import (
    DocumentPromptMessageContent,
    ImagePromptMessageContent,
+    PromptMessageContentType,
    TextPromptMessageContent,
    UserPromptMessage,
 )
@@ -133,6 +134,7 @@ def runner(mocker: MockerFixture, sqlite_engine: Engine) -> Iterator[FunctionCal
    runner.history_prompt_messages = []
    runner._current_thoughts = []
    runner.files = []
+    runner.vision_enabled = False
    runner.agent_callback = MagicMock()
    runner.session = Session(sqlite_engine)

@@ -290,6 +292,82 @@ class TestClearUserPromptImageMessages:

        assert result[0].content == "hello\n[image]\n[file]"

+    def test_keeps_knowledge_retrieval_image_message(self, runner: FunctionCallAgentRunner):
+        text = TextPromptMessageContent(data="query")
+        image = ImagePromptMessageContent(format="url", mime_type="image/png")
+        user_msg = UserPromptMessage(name="knowledge_retrieval", content=[image, text])
+
+        result = runner._clear_user_prompt_image_messages([user_msg])
+
+        assert result[0].content == [image, text]
+
+
+# ==============================
+# Dataset Tool Image Content
+# ==============================
+
+
+class TestBuildDatasetToolImageContents:
+    def test_returns_empty_when_vision_disabled(self, runner: FunctionCallAgentRunner):
+        tool = MagicMock()
+        tool.__class__.__name__ = "DatasetRetrieverTool"
+        response = "![image](http://localhost:5001/files/890985e9-c2f1-484e-bc7b-62010a337e6d/file-preview)"
+
+        assert runner._build_dataset_tool_image_contents(runner.session, response, tool) == []
+
+    def test_builds_image_contents_from_dataset_tool_preview_links(
+        self, runner: FunctionCallAgentRunner, mocker: MockerFixture
+    ):
+        from core.tools.utils.dataset_retriever_tool import DatasetRetrieverTool
+
+        runner.vision_enabled = True
+        image_content = ImagePromptMessageContent(format="url", mime_type="image/png")
+        to_prompt_content = mocker.patch(
+            "core.agent.fc_agent_runner.file_manager.to_prompt_message_content",
+            return_value=image_content,
+        )
+        grant_access = mocker.patch("core.agent.fc_agent_runner.grant_upload_file_access")
+        sign_preview = mocker.patch(
+            "core.agent.fc_agent_runner.sign_upload_file_preview_url",
+            return_value="http://localhost:5001/files/file-id/file-preview?sign=1",
+        )
+        build_reference = mocker.patch("core.agent.fc_agent_runner.build_file_reference", return_value="file-ref")
+
+        upload_file = MagicMock()
+        upload_file.id = "890985e9-c2f1-484e-bc7b-62010a337e6d"
+        upload_file.name = "chart.png"
+        upload_file.extension = "png"
+        upload_file.mime_type = "image/png"
+        upload_file.source_url = ""
+        upload_file.size = 123
+        upload_file.key = "image_files/chart.png"
+
+        non_image_file = MagicMock()
+        non_image_file.id = "11111111-1111-1111-1111-111111111111"
+        non_image_file.mime_type = "application/pdf"
+
+        scalars_result = MagicMock()
+        scalars_result.all.return_value = [upload_file, non_image_file]
+        session = MagicMock()
+        session.scalars.return_value = scalars_result
+
+        response = (
+            "![image](http://localhost:5001/files/890985e9-c2f1-484e-bc7b-62010a337e6d/file-preview?sign=1)\n"
+            "duplicate ![image](http://localhost:5001/files/890985e9-c2f1-484e-bc7b-62010a337e6d/file-preview)\n"
+            "file ![file](http://localhost:5001/files/11111111-1111-1111-1111-111111111111/file-preview)"
+        )
+
+        tool = MagicMock(spec=DatasetRetrieverTool)
+        contents = runner._build_dataset_tool_image_contents(session, response, tool)
+
+        assert contents == [image_content]
+        assert contents[0].type == PromptMessageContentType.IMAGE
+        grant_access.assert_called_once()
+        assert list(grant_access.call_args.args[0]) == ["890985e9-c2f1-484e-bc7b-62010a337e6d"]
+        sign_preview.assert_called_once_with(upload_file.id, upload_file.extension)
+        build_reference.assert_called_once_with(record_id=str(upload_file.id))
+        to_prompt_content.assert_called_once()
+

 # ==============================
 # Run Method Tests