Smarter download_artifacts tool

ymir/tools/privileged/copr.py

@@ -2,8 +2,10 @@
 import gzip
 import logging
 import random
+import tempfile
 import time
 from pathlib import Path
+from shutil import rmtree
 from urllib.parse import urljoin, urlparse
 
 import aiohttp
@@ -12,7 +14,6 @@
 from beeai_framework.emitter import Emitter
 from beeai_framework.tools import (
     JSONToolOutput,
-    StringToolOutput,
     ToolError,
     ToolRunOptions,
 )
@@ -324,13 +325,22 @@ async def branch_to_chroot(dist_git_branch: str, upcoming_z_streams: dict[str, s
 
 class DownloadArtifactsToolInput(BaseModel):
     artifacts_urls: list[str] = Field(description="URLs to build artifacts (logs and RPM files)")
-    target_path: AbsolutePath = Field(description="Absolute path where to download the artifacts")
 
 
-class DownloadArtifactsTool(Tool[DownloadArtifactsToolInput, ToolRunOptions, StringToolOutput]):
+class DownloadArtifactsResult(BaseModel):
+    target_path: str = Field(description="Location of downloaded files")
+
+
+class DownloadArtifactsToolOutput(JSONToolOutput[DownloadArtifactsResult]):
+    def get_text_content(self) -> str:
+        """Return content with minimum tokens possible"""
+        return f"target_path: {self.result.target_path}"
+
+
+class DownloadArtifactsTool(Tool[DownloadArtifactsToolInput, ToolRunOptions, DownloadArtifactsToolOutput]):
     name = "download_artifacts"
     description = """
-    Downloads build artifacts to the specified location.
+    Downloads build artifacts to a temporary location.
     Any gzipped log files will be automatically decompressed,
     for example `http://example.com/builder-live.log.gz`
     will be downloaded as `builder-live.log`.
@@ -348,9 +358,9 @@ async def _run(
         tool_input: DownloadArtifactsToolInput,
         options: ToolRunOptions | None,
         context: RunContext,
-    ) -> StringToolOutput:
+    ) -> DownloadArtifactsToolOutput:
         artifacts_urls = tool_input.artifacts_urls
-        target_path = tool_input.target_path
+        target_path = tempfile.mkdtemp()
         async with aiohttp.ClientSession(
             timeout=AIOHTTP_TIMEOUT, headers={"User-Agent": YMIR_USER_AGENT}
         ) as session:
@@ -369,7 +379,11 @@ async def _run(
                                     target = target.with_suffix("")
                             (target_path / target).write_bytes(content)
                         else:
+                            # Cleanup temporary dir
+                            rmtree(target_path)
                             raise ToolError(f"Failed to download {url}: {response.status} {response.reason}")
                 except (aiohttp.ClientError, TimeoutError) as e:
+                    # Cleanup temporary dir
+                    rmtree(target_path)
                     raise ToolError(f"Failed to download {url}: {e}") from e
-        return StringToolOutput(result="Successfully downloaded the specified build artifacts")
+        return DownloadArtifactsToolOutput(result=DownloadArtifactsResult(target_path=target_path))

ymir/tools/privileged/tests/unit/test_copr.py

@@ -170,9 +170,8 @@ async def sleep(*_):
     ],
 )
 @pytest.mark.asyncio
-async def test_download_artifacts(url, tmp_path):
+async def test_download_artifacts(url):
     artifacts_urls = [url]
-    target_path = tmp_path
     content = b"12345"
     content_gz = b"\x1f\x8b\x00\x00"
 
@@ -189,19 +188,12 @@ async def read():
     )
     if "broken" in url:
         with pytest.raises(ToolError):
-            await DownloadArtifactsTool().run(
-                input={"artifacts_urls": artifacts_urls, "target_path": target_path}
-            )
+            await DownloadArtifactsTool().run(input={"artifacts_urls": artifacts_urls})
     else:
-        out = await DownloadArtifactsTool().run(
-            input={"artifacts_urls": artifacts_urls, "target_path": target_path}
-        )
+        out = await DownloadArtifactsTool().run(input={"artifacts_urls": artifacts_urls})
         result = out.result
-        assert result.startswith("Successfully")
-    path = target_path / url.rsplit("/", 1)[-1].removesuffix(".gz")
-    if "broken" in url:
-        assert not path.is_file()
-    else:
+        target_path = Path(result.target_path)
+        path = target_path / url.rsplit("/", 1)[-1].removesuffix(".gz")
         assert path.read_bytes() == content
 
 

ymir/tools/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ymir-tools"
-version = "0.6.1"
+version = "0.7.0"
 description = "Ymir MCP tools for AI workflows"
 requires-python = ">=3.13"
 dynamic = ["dependencies"]