This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b13d796fbf1 Fixing mypy errors in google gen ai provider and tests
(#67276)
b13d796fbf1 is described below
commit b13d796fbf19868635191f936073477b188cff67
Author: Amogh Desai <[email protected]>
AuthorDate: Thu May 21 16:25:45 2026 +0530
Fixing mypy errors in google gen ai provider and tests (#67276)
---
.../providers/google/cloud/operators/gen_ai.py | 4 +--
.../gen_ai/example_gen_ai_generative_model.py | 29 +++++++++++-----------
.../tests/unit/google/cloud/hooks/test_gen_ai.py | 29 +++++++++++-----------
.../unit/google/cloud/operators/test_gen_ai.py | 29 +++++++++++-----------
4 files changed, 44 insertions(+), 47 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
index da081840a7c..71a51c32419 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
@@ -164,7 +164,7 @@ class GenAIGenerateContentOperator(GoogleCloudBaseOperator):
*,
project_id: str,
location: str,
- contents: ContentListUnionDict,
+ contents: ContentListUnionDict | list[str],
model: str,
generation_config: GenerateContentConfig | dict[str, Any] | None =
None,
gcp_conn_id: str = "google_cloud_default",
@@ -311,7 +311,7 @@ class GenAICountTokensOperator(GoogleCloudBaseOperator):
*,
project_id: str,
location: str,
- contents: ContentListUnion | ContentListUnionDict,
+ contents: ContentListUnion | ContentListUnionDict | list[str],
model: str,
config: CountTokensConfigOrDict | None = None,
gcp_conn_id: str = "google_cloud_default",
diff --git
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
index a1a0123a761..c9a6f9ffcaf 100644
---
a/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
+++
b/providers/google/tests/system/google/cloud/gen_ai/example_gen_ai_generative_model.py
@@ -199,22 +199,21 @@ CACHED_SYSTEM_INSTRUCTION = """
You are an expert researcher. You always stick to the facts in the sources
provided, and never make up new facts.
Now look at these research papers, and answer the following questions.
"""
-CACHED_CONTENT_CONFIG = CreateCachedContentConfig(
- contents=[
- Content(
- role="user",
- parts=[
- Part.from_uri(
-
file_uri="gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf",
- mime_type="application/pdf",
- ),
- Part.from_uri(
-
file_uri="gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf",
- mime_type="application/pdf",
- ),
- ],
- )
+_CACHED_CONTENT = Content(
+ role="user",
+ parts=[
+ Part.from_uri(
+
file_uri="gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf",
+ mime_type="application/pdf",
+ ),
+ Part.from_uri(
+
file_uri="gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf",
+ mime_type="application/pdf",
+ ),
],
+)
+CACHED_CONTENT_CONFIG = CreateCachedContentConfig(
+ contents=_CACHED_CONTENT,
system_instruction=CACHED_SYSTEM_INSTRUCTION,
display_name="test-cache",
ttl="3600s",
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_gen_ai.py
b/providers/google/tests/unit/google/cloud/hooks/test_gen_ai.py
index 4bc691ad52b..4a5e8c669b6 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_gen_ai.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_gen_ai.py
@@ -75,22 +75,21 @@ You are an expert researcher. You always stick to the facts
in the sources provi
Now look at these research papers, and answer the following questions.
"""
-CACHED_CONTENT_CONFIG = CreateCachedContentConfig(
- contents=[
- Content(
- role="user",
- parts=[
- Part.from_uri(
-
file_uri="gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf",
- mime_type="application/pdf",
- ),
- Part.from_uri(
-
file_uri="gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf",
- mime_type="application/pdf",
- ),
- ],
- )
+_CACHED_CONTENT = Content(
+ role="user",
+ parts=[
+ Part.from_uri(
+
file_uri="gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf",
+ mime_type="application/pdf",
+ ),
+ Part.from_uri(
+
file_uri="gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf",
+ mime_type="application/pdf",
+ ),
],
+)
+CACHED_CONTENT_CONFIG = CreateCachedContentConfig(
+ contents=_CACHED_CONTENT,
system_instruction=TEST_CACHED_SYSTEM_INSTRUCTION,
display_name="test-cache",
ttl="3600s",
diff --git a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
index 9446f5b858a..cbd2105ef82 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
@@ -60,22 +60,21 @@ CACHED_SYSTEM_INSTRUCTION = """
You are an expert researcher. You always stick to the facts in the sources
provided, and never make up new facts.
Now look at these research papers, and answer the following questions.
"""
-CACHED_CONTENT_CONFIG = CreateCachedContentConfig(
- contents=[
- Content(
- role="user",
- parts=[
- Part.from_uri(
-
file_uri="gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf",
- mime_type="application/pdf",
- ),
- Part.from_uri(
-
file_uri="gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf",
- mime_type="application/pdf",
- ),
- ],
- )
+_CACHED_CONTENT = Content(
+ role="user",
+ parts=[
+ Part.from_uri(
+
file_uri="gs://cloud-samples-data/generative-ai/pdf/2312.11805v3.pdf",
+ mime_type="application/pdf",
+ ),
+ Part.from_uri(
+
file_uri="gs://cloud-samples-data/generative-ai/pdf/2403.05530.pdf",
+ mime_type="application/pdf",
+ ),
],
+)
+CACHED_CONTENT_CONFIG = CreateCachedContentConfig(
+ contents=_CACHED_CONTENT,
system_instruction=CACHED_SYSTEM_INSTRUCTION,
display_name="test-cache",
ttl="3600s",
