Function generate_media_artifact

contributing/samples/core/artifacts/agent.py:159–204

Generates a valid media artifact of specified type. Args: media_type: One of 'image', 'audio', 'video'. ctx: The tool context for saving artifacts.

(media_type: str, ctx: Context)

async def generate_media_artifact(media_type: str, ctx: Context) -> dict:
  """Generates a valid media artifact of specified type.

  Args:
      media_type: One of 'image', 'audio', 'video'.
      ctx: The tool context for saving artifacts.
  """

  with tempfile.TemporaryDirectory() as tmpdir:
    if media_type == "image":
      mime_type = "image/bmp"
      file_path = os.path.join(tmpdir, "sample.bmp")
      generate_bmp(file_path)
      filename = "sample_image.bmp"
    elif media_type == "audio":
      mime_type = "audio/wav"
      file_path = os.path.join(tmpdir, "sample.wav")
      generate_wav(file_path)
      filename = "sample_audio.wav"
    elif media_type == "video":
      mime_type = "video/mp4"
      file_path = os.path.join(tmpdir, "sample.mp4")
      try:
        generate_video(file_path)
      except (ImportError, RuntimeError) as e:
        return {"error": str(e)}
      filename = "sample_video.mp4"
    else:
      return {"error": f"Unsupported media type: {media_type}"}

    with open(file_path, "rb") as f:
      data = f.read()

    version = await ctx.save_artifact(
        filename,
        types.Part.from_bytes(data=data, mime_type=mime_type),
    )

  return {
      "message": (
          f"Media artifact '{filename}' generated and saved (version"
          f" {version})."
      ),
      "filename": filename,
      "version": version,
  }


root_agent = Agent(