From 0235ec0daef0ae5f22225a4c1535bd451035ff77 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Mon, 6 May 2024 15:35:35 -0700 Subject: [PATCH 1/2] Fix typing Change-Id: I09fb7df098da08bb24337b08cbbc997c4c62af1e --- google/generativeai/types/content_types.py | 35 ++-------------------- google/generativeai/types/file_types.py | 31 +++++++++++++++++++ 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index e19163331..f87740122 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -113,7 +113,7 @@ def _convert_dict(d: Mapping) -> glm.Content | glm.Part | glm.Blob: if "inline_data" in part: part["inline_data"] = to_blob(part["inline_data"]) if "file_data" in part: - part["file_data"] = to_file_data(part["file_data"]) + part["file_data"] = file_types.to_file_data(part["file_data"]) return glm.Part(part) elif is_blob_dict(d): blob = d @@ -161,35 +161,6 @@ def to_blob(blob: BlobType) -> glm.Blob: ) -class FileDataDict(TypedDict): - mime_type: str - file_uri: str - - -FileDataType = Union[FileDataDict, glm.FileData, file_types.File] - - -def to_file_data(file_data: FileDataType): - if isinstance(file_data, dict): - if "file_uri" in file_data: - file_data = glm.FileData(file_data) - else: - file_data = glm.File(file_data) - - if isinstance(file_data, file_types.File): - file_data = file_data.to_proto() - - if isinstance(file_data, (glm.File, file_types.File)): - file_data = glm.FileData( - mime_type=file_data.mime_type, - file_uri=file_data.uri, - ) - - if isinstance(file_data, glm.FileData): - return file_data - else: - raise TypeError(f"Could not convert a {type(file_data)} to `FileData`") - class PartDict(TypedDict): text: str @@ -197,7 +168,7 @@ class PartDict(TypedDict): # When you need a `Part` accept a part object, part-dict, blob or string -PartType = Union[glm.Part, PartDict, BlobType, str, glm.FunctionCall, glm.FunctionResponse] +PartType = Union[glm.Part, PartDict, BlobType, str, glm.FunctionCall, glm.FunctionResponse, file_types.FileDataType] def is_part_dict(d): @@ -221,7 +192,7 @@ def to_part(part: PartType): elif isinstance(part, glm.FileData): return glm.Part(file_data=part) elif isinstance(part, (glm.File, file_types.File)): - return glm.Part(file_data=to_file_data(part)) + return glm.Part(file_data=file_types.to_file_data(part)) elif isinstance(part, glm.FunctionCall): return glm.Part(function_call=part) elif isinstance(part, glm.FunctionResponse): diff --git a/google/generativeai/types/file_types.py b/google/generativeai/types/file_types.py index a2612d47c..0c66d70d0 100644 --- a/google/generativeai/types/file_types.py +++ b/google/generativeai/types/file_types.py @@ -15,6 +15,8 @@ from __future__ import annotations import datetime +from typing import Union +from typing_extensions import TypedDict from google.generativeai.client import get_default_file_client @@ -69,3 +71,32 @@ def uri(self) -> str: def delete(self): client = get_default_file_client() client.delete_file(name=self.name) + + +class FileDataDict(TypedDict): + mime_type: str + file_uri: str + + +FileDataType = Union[FileDataDict, glm.FileData, glm.File, File] + +def to_file_data(file_data: FileDataType): + if isinstance(file_data, dict): + if "file_uri" in file_data: + file_data = glm.FileData(file_data) + else: + file_data = glm.File(file_data) + + if isinstance(file_data, File): + file_data = file_data.to_proto() + + if isinstance(file_data, glm.File): + file_data = glm.FileData( + mime_type=file_data.mime_type, + file_uri=file_data.uri, + ) + + if isinstance(file_data, glm.FileData): + return file_data + else: + raise TypeError(f"Could not convert a {type(file_data)} to `FileData`") From 0f9485d395d968ff06d3898f966f207951b7f524 Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Thu, 9 May 2024 14:07:22 -0700 Subject: [PATCH 2/2] format Change-Id: I7923e2257a6d935bb1b60f4f81dd7910387292ae --- google/generativeai/types/content_types.py | 11 +++++++++-- google/generativeai/types/file_types.py | 1 + 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/google/generativeai/types/content_types.py b/google/generativeai/types/content_types.py index f87740122..ec06001fd 100644 --- a/google/generativeai/types/content_types.py +++ b/google/generativeai/types/content_types.py @@ -161,14 +161,21 @@ def to_blob(blob: BlobType) -> glm.Blob: ) - class PartDict(TypedDict): text: str inline_data: BlobType # When you need a `Part` accept a part object, part-dict, blob or string -PartType = Union[glm.Part, PartDict, BlobType, str, glm.FunctionCall, glm.FunctionResponse, file_types.FileDataType] +PartType = Union[ + glm.Part, + PartDict, + BlobType, + str, + glm.FunctionCall, + glm.FunctionResponse, + file_types.FileDataType, +] def is_part_dict(d): diff --git a/google/generativeai/types/file_types.py b/google/generativeai/types/file_types.py index 0c66d70d0..ce71be534 100644 --- a/google/generativeai/types/file_types.py +++ b/google/generativeai/types/file_types.py @@ -80,6 +80,7 @@ class FileDataDict(TypedDict): FileDataType = Union[FileDataDict, glm.FileData, glm.File, File] + def to_file_data(file_data: FileDataType): if isinstance(file_data, dict): if "file_uri" in file_data: