fix(artifacts): accept dict-shaped artifacts in InMemoryArtifactService (#3495)

src/google/adk/artifacts/artifact_util.py

@@ -109,8 +109,15 @@ def is_artifact_ref(artifact: types.Part) -> bool:
   Returns:
       True if the artifact part is an artifact reference, False otherwise.
   """
-  return bool(
-      artifact.file_data
-      and artifact.file_data.file_uri
-      and artifact.file_data.file_uri.startswith("artifact://")
-  )
+  # Support both object-like `types.Part` and plain `dict` shapes.
+  file_data = None
+  if isinstance(artifact, dict):
+    file_data = artifact.get("file_data")
+  else:
+    file_data = getattr(artifact, "file_data", None)
+
+  if not file_data:
+    return False
+
+  file_uri = file_data.get("file_uri") if isinstance(file_data, dict) else getattr(file_data, "file_uri", None)
+  return bool(file_uri and isinstance(file_uri, str) and file_uri.startswith("artifact://"))

src/google/adk/artifacts/in_memory_artifact_service.py

@@ -117,23 +117,44 @@ async def save_artifact(
     )
     if custom_metadata:
       artifact_version.custom_metadata = custom_metadata
-
-    if artifact.inline_data is not None:
-      artifact_version.mime_type = artifact.inline_data.mime_type
-    elif artifact.text is not None:
+    # Safely extract fields from `artifact`, supporting both object-like
+    # `types.Part` and plain `dict` payloads. This makes the in-memory
+    # service resilient to callers that send a dict (AgentSpace uploads).
+    def _get_field(obj, name):
+      if isinstance(obj, dict):
+        return obj.get(name)
+      return getattr(obj, name, None)
+
+    inline_data = _get_field(artifact, "inline_data")
+    text = _get_field(artifact, "text")
+    file_data = _get_field(artifact, "file_data")
+
+    if inline_data is not None:
+      # inline_data may be a dict or an object
+      artifact_version.mime_type = (
+          inline_data.get("mime_type")
+          if isinstance(inline_data, dict)
+          else inline_data.mime_type
+      )
+    elif text is not None:
       artifact_version.mime_type = "text/plain"
-    elif artifact.file_data is not None:
+    elif file_data is not None:
+      # If the artifact is an artifact-ref we validate the referenced URI.
       if artifact_util.is_artifact_ref(artifact):
-        if not artifact_util.parse_artifact_uri(artifact.file_data.file_uri):
-          raise ValueError(
-              f"Invalid artifact reference URI: {artifact.file_data.file_uri}"
-          )
-        # If it's a valid artifact URI, we store the artifact part as-is.
-        # And we don't know the mime type until we load it.
+        file_uri = (
+            file_data.get("file_uri") if isinstance(file_data, dict) else file_data.file_uri
+        )
+        if not artifact_util.parse_artifact_uri(file_uri):
+          raise ValueError(f"Invalid artifact reference URI: {file_uri}")
+        # Valid artifact URI: keep part as-is; mime type may be resolved later.
       else:
-        artifact_version.mime_type = artifact.file_data.mime_type
+        artifact_version.mime_type = (
+            file_data.get("mime_type") if isinstance(file_data, dict) else file_data.mime_type
+        )
     else:
-      raise ValueError("Not supported artifact type.")
+      # Fallback for unknown shapes: preserve behavior by storing the
+      # artifact but use a generic binary mime type instead of raising.
+      artifact_version.mime_type = "application/octet-stream"
 
     self.artifacts[path].append(
         _ArtifactEntry(data=artifact, artifact_version=artifact_version)
@@ -167,15 +188,18 @@ async def load_artifact(
 
     # Resolve artifact reference if needed.
     artifact_data = artifact_entry.data
+    # Resolve artifact reference if needed. Support dict-shaped stored
+    # artifacts as well as object-like `types.Part`.
     if artifact_util.is_artifact_ref(artifact_data):
-      parsed_uri = artifact_util.parse_artifact_uri(
-          artifact_data.file_data.file_uri
-      )
+      # Extract file_uri safely for dict or object shapes.
+      if isinstance(artifact_data, dict):
+        file_uri = artifact_data.get("file_data", {}).get("file_uri")
+      else:
+        file_uri = getattr(artifact_data.file_data, "file_uri", None)
+
+      parsed_uri = artifact_util.parse_artifact_uri(file_uri)
       if not parsed_uri:
-        raise ValueError(
-            "Invalid artifact reference URI:"
-            f" {artifact_data.file_data.file_uri}"
-        )
+        raise ValueError(f"Invalid artifact reference URI: {file_uri}")
       return await self.load_artifact(
           app_name=parsed_uri.app_name,
           user_id=parsed_uri.user_id,
@@ -184,12 +208,32 @@ async def load_artifact(
           version=parsed_uri.version,
       )
 
-    if (
-        artifact_data == types.Part()
-        or artifact_data == types.Part(text="")
-        or (artifact_data.inline_data and not artifact_data.inline_data.data)
-    ):
+    # Determine emptiness for both shapes.
+    def _is_empty(a):
+      if a is None:
+        return True
+      if isinstance(a, dict):
+        # common empty forms: empty text, empty inline_data, or inline_data with no bytes
+        if a.get("text") in (None, "") and not a.get("inline_data") and not a.get("file_data"):
+          return True
+        inline = a.get("inline_data")
+        if inline and isinstance(inline, dict) and not inline.get("data"):
+          return True
+        return False
+      # object-like types.Part
+      try:
+        if a == types.Part() or a == types.Part(text=""):
+          return True
+        inline = getattr(a, "inline_data", None)
+        if inline and not getattr(inline, "data", None):
+          return True
+      except Exception:
+        return False
+      return False
+
+    if _is_empty(artifact_data):
       return None
+
     return artifact_data
 
   @override

tests/unittests/artifacts/test_in_memory_artifact_dict_handling.py

@@ -0,0 +1,181 @@
+import sys
+import types as _types
+
+# Prefer the real `google.genai.types` if installed; fall back to a minimal
+# local stub only when the real package is not available. This avoids
+# overriding the real `google` package and ensures type hints used across
+# the ADK package are satisfied when the dependency is present.
+try:
+  from google.genai import types as types_mod  # type: ignore
+except Exception:
+  genai = _types.ModuleType("google.genai")
+  types_mod = _types.ModuleType("google.genai.types")
+
+  # Minimal Part and Content classes used in tests when google-genai
+  # is not installed. These are intentionally small and only provide
+  # the attributes the InMemoryArtifactService touches.
+  class Part:
+    def __init__(self, inline_data=None, text=None, file_data=None):
+      self.inline_data = inline_data
+      self.text = text
+      self.file_data = file_data
+
+    def __eq__(self, other):
+      if not isinstance(other, Part):
+        return False
+      return (
+        (self.inline_data == other.inline_data)
+        and (self.text == other.text)
+        and (self.file_data == other.file_data)
+      )
+
+  class Content:
+    pass
+
+  types_mod.Part = Part
+  types_mod.Content = Content
+  setattr(genai, "types", types_mod)
+  sys.modules["google.genai"] = genai
+  sys.modules["google.genai.types"] = types_mod
+
+import importlib.util
+import importlib.machinery
+import types as _types
+import os
+
+# Create lightweight stubs for package modules to avoid importing the
+# whole ADK package during this focused unit test. We populate
+# `sys.modules` entries so relative imports inside the target module
+# resolve against these stubs.
+if "google.adk" not in sys.modules:
+  sys.modules.setdefault("google", _types.ModuleType("google"))
+  sys.modules.setdefault("google.adk", _types.ModuleType("google.adk"))
+  sys.modules.setdefault("google.adk.artifacts", _types.ModuleType("google.adk.artifacts"))
+
+# Stub for base_artifact_service used by the in-memory service.
+base_mod = _types.ModuleType("google.adk.artifacts.base_artifact_service")
+from dataclasses import dataclass
+from typing import Optional, Any
+
+@dataclass
+class ArtifactVersion:
+  version: int
+  canonical_uri: str
+  mime_type: Optional[str] = None
+  custom_metadata: Optional[dict[str, Any]] = None
+
+class BaseArtifactService:
+  pass
+
+base_mod.ArtifactVersion = ArtifactVersion
+base_mod.BaseArtifactService = BaseArtifactService
+sys.modules["google.adk.artifacts.base_artifact_service"] = base_mod
+
+# Minimal artifact_util stub
+artifact_util_mod = _types.ModuleType("google.adk.artifacts.artifact_util")
+def parse_artifact_uri(uri: str):
+  return None
+def is_artifact_ref(artifact):
+  if isinstance(artifact, dict):
+    fd = artifact.get("file_data")
+    if not fd:
+      return False
+    uri = fd.get("file_uri")
+    return isinstance(uri, str) and uri.startswith("artifact://")
+  # object-like: attempt attribute access
+  fd = getattr(artifact, "file_data", None)
+  if not fd:
+    return False
+  uri = getattr(fd, "file_uri", None)
+  return isinstance(uri, str) and uri.startswith("artifact://")
+
+artifact_util_mod.parse_artifact_uri = parse_artifact_uri
+artifact_util_mod.is_artifact_ref = is_artifact_ref
+sys.modules["google.adk.artifacts.artifact_util"] = artifact_util_mod
+
+# Stub google.genai.types if not available
+try:
+  from google.genai import types as types_mod  # type: ignore
+except Exception:
+  genai = _types.ModuleType("google.genai")
+  types_mod = _types.ModuleType("google.genai.types")
+  class Part:
+    def __init__(self, inline_data=None, text=None, file_data=None):
+      self.inline_data = inline_data
+      self.text = text
+      self.file_data = file_data
+
+    def __eq__(self, other):
+      if not isinstance(other, Part):
+        return False
+      return (
+        (self.inline_data == other.inline_data)
+        and (self.text == other.text)
+        and (self.file_data == other.file_data)
+      )
+  class Content:
+    pass
+  types_mod.Part = Part
+  types_mod.Content = Content
+  sys.modules["google.genai"] = genai
+  sys.modules["google.genai.types"] = types_mod
+
+# Load the in-memory artifact service module directly from source
+module_path = os.path.join(os.getcwd(), "src", "google", "adk", "artifacts", "in_memory_artifact_service.py")
+spec = importlib.util.spec_from_file_location("google.adk.artifacts.in_memory_artifact_service", module_path)
+in_memory = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = in_memory
+spec.loader.exec_module(in_memory)
+InMemoryArtifactService = in_memory.InMemoryArtifactService
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_save_artifact_with_dict_and_part():
+  svc = InMemoryArtifactService()
+
+  # Case 1: artifact passed as plain dict (simulating AgentSpace upload)
+  dict_artifact = {
+    "file_data": {"file_uri": "memory://apps/a/u/s/f/versions/0", "mime_type": "text/plain"}
+  }
+  ver = await svc.save_artifact(
+    app_name="a",
+    user_id="u",
+    filename="s/f",
+    session_id="session",
+    artifact=dict_artifact,
+  )
+  assert ver == 0
+
+  # Case 2: artifact passed as object-like Part
+  part = types_mod.Part(text="hello")
+  ver2 = await svc.save_artifact(
+    app_name="a",
+    user_id="u",
+    filename="s/f",
+    session_id="session",
+    artifact=part,
+  )
+  assert ver2 == 1
+
+  # Ensure load returns the same object-like Part for the last saved
+  loaded = await svc.load_artifact(app_name="a", user_id="u", filename="s/f", session_id="session")
+  assert isinstance(loaded, types_mod.Part)
+
+
+@pytest.mark.asyncio
+async def test_save_artifact_with_inline_dict():
+  svc = InMemoryArtifactService()
+  inline = {"inline_data": {"mime_type": "image/png", "data": b"\x89PNG"}}
+  ver = await svc.save_artifact(
+    app_name="app",
+    user_id="user",
+    filename="user:avatar.png",
+    artifact=inline,
+  )
+  assert ver == 0
+
+  # list keys should include the user-scoped filename
+  keys = await svc.list_artifact_keys(app_name="app", user_id="user")
+  assert "user:avatar.png" in keys