diff --git a/langfun/core/data/conversion/gemini.py b/langfun/core/data/conversion/gemini.py index 5620e6da..cba207c3 100644 --- a/langfun/core/data/conversion/gemini.py +++ b/langfun/core/data/conversion/gemini.py @@ -83,12 +83,14 @@ def _convert_chunk( if video_metadata: result['videoMetadata'] = video_metadata return result + if chunk.uri and chunk.uri.startswith('gs://'): + always_send_content = False if chunk.is_text: return {'text': chunk.to_text()} if ( not always_send_content and chunk.uri - and chunk.uri.lower().startswith(('http:', 'https:', 'ftp:')) + and chunk.uri.lower().startswith(('http:', 'https:', 'ftp:', 'gs:')) ): return { 'fileData': { diff --git a/langfun/core/data/conversion/gemini_test.py b/langfun/core/data/conversion/gemini_test.py index 33489cbd..c37364f6 100644 --- a/langfun/core/data/conversion/gemini_test.py +++ b/langfun/core/data/conversion/gemini_test.py @@ -92,6 +92,30 @@ def test_as_format_with_image(self): }, ) + def test_as_format_with_gcs_uri(self): + self.assertEqual( + lf.Template( + 'Summarize {{doc}}', + doc=lf_modalities.Custom.from_uri( + 'gs://my-bucket/temp/doc.pdf', mime='application/pdf' + ), + ) + .render() + .as_gemini_format(), + { + 'role': 'user', + 'parts': [ + {'text': 'Summarize'}, + { + 'fileData': { + 'fileUri': 'gs://my-bucket/temp/doc.pdf', + 'mimeType': 'application/pdf', + } + }, + ], + }, + ) + def test_as_format_youtube_with_video_metadata(self): self.assertEqual( lf.Template(