From 58acba1e9a74cae8093aa17fead82f85c8052512 Mon Sep 17 00:00:00 2001 From: drona-gyawali Date: Sat, 8 Mar 2025 00:26:14 +0545 Subject: [PATCH] Feature/Add support for per-request metadata/headers #698 --- google/generativeai/generative_models.py | 33 ++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/google/generativeai/generative_models.py b/google/generativeai/generative_models.py index 8d331a9f6..7ee38be73 100644 --- a/google/generativeai/generative_models.py +++ b/google/generativeai/generative_models.py @@ -244,6 +244,7 @@ def generate_content( tools: content_types.FunctionLibraryType | None = None, tool_config: content_types.ToolConfigType | None = None, request_options: helper_types.RequestOptionsType | None = None, + extra_headers: dict[str, str] | None = None, ) -> generation_types.GenerateContentResponse: """A multipurpose function to generate responses from the model. @@ -298,6 +299,7 @@ def generate_content( stream: If True, yield response chunks as they are generated. tools: `protos.Tools` more info coming soon. request_options: Options for the request. + extra_headers: Optional headers to include with this specific request. """ if not contents: raise TypeError("contents must not be empty") @@ -318,6 +320,21 @@ def generate_content( if request_options is None: request_options = {} + else: + request_options = ( + request_options.copy() + if isinstance(request_options, dict) + else vars(request_options).copy() + ) + + metadata = list(request_options.get("metadata", [])) + if extra_headers: + metadata += list(extra_headers.items()) + + if metadata: + request_options["metadata"] = metadata + elif "metadata" in request_options: + del request_options["metadata"] try: if stream: @@ -351,6 +368,7 @@ async def generate_content_async( tools: content_types.FunctionLibraryType | None = None, tool_config: content_types.ToolConfigType | None = None, request_options: helper_types.RequestOptionsType | None = None, + extra_headers: dict[str, str] | None = None, ) -> generation_types.AsyncGenerateContentResponse: """The async version of `GenerativeModel.generate_content`.""" if not contents: @@ -372,6 +390,21 @@ async def generate_content_async( if request_options is None: request_options = {} + else: + request_options = ( + request_options.copy() + if isinstance(request_options, dict) + else vars(request_options).copy() + ) + + metadata = list(request_options.get("metadata", [])) + if extra_headers: + metadata += list(extra_headers.items()) + + if metadata: + request_options["metadata"] = metadata + elif "metadata" in request_options: + del request_options["metadata"] try: if stream: