From 94e6f6198241c464949dbdebb6500165e274e9a0 Mon Sep 17 00:00:00 2001 From: Mayuresh Agashe Date: Mon, 10 Jun 2024 18:32:29 +0000 Subject: [PATCH] Accept partial file names Adds prefix files/ if not present Change-Id: Iac5c4d0934620f2462cf15ae519474b0ce7908da --- google/generativeai/files.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/google/generativeai/files.py b/google/generativeai/files.py index 4028d37f7..0e8a8ed27 100644 --- a/google/generativeai/files.py +++ b/google/generativeai/files.py @@ -81,16 +81,20 @@ def list_files(page_size=100) -> Iterable[file_types.File]: yield file_types.File(proto) -def get_file(name) -> file_types.File: +def get_file(name: str) -> file_types.File: """Calls the API to retrieve a specified file using a supported file service.""" + if "/" not in name: + name = f"files/{name}" client = get_default_file_client() return file_types.File(client.get_file(name=name)) -def delete_file(name): +def delete_file(name: str | file_types.File | protos.File): """Calls the API to permanently delete a specified file using a supported file service.""" if isinstance(name, (file_types.File, protos.File)): name = name.name + elif "/" not in name: + name = f"files/{name}" request = protos.DeleteFileRequest(name=name) client = get_default_file_client() client.delete_file(request=request)