Skip to content

Commit 0bc6849

Browse files
authored
Merge pull request #113 from owndev/add-user-info-headers-forwarding-option
Add user info headers forwarding option
2 parents 5aa7567 + 1a51fb4 commit 0bc6849

File tree

1 file changed

+48
-4
lines changed

1 file changed

+48
-4
lines changed

pipelines/google/google_gemini.py

Lines changed: 48 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
author_url: https://github.com/owndev/
55
project_url: https://github.com/owndev/Open-WebUI-Functions
66
funding_url: https://github.com/sponsors/owndev
7-
version: 1.8.2
7+
version: 1.8.3
88
required_open_webui_version: 0.6.26
99
license: Apache License 2.0
1010
description: Highly optimized Google Gemini pipeline with advanced image generation capabilities, intelligent compression, and streamlined processing workflows.
@@ -191,11 +191,12 @@ class Valves(BaseModel):
191191
description="The Google Cloud region to use with Vertex AI.",
192192
)
193193
VERTEX_AI_RAG_STORE: str | None = Field(
194-
default=os.getenv("VERTEX_AI_RAG_STORE"),
194+
default=os.getenv("GOOGLE_VERTEX_AI_RAG_STORE"),
195195
description="Vertex AI RAG Store path for grounding (e.g., projects/PROJECT/locations/LOCATION/ragCorpora/DATA_STORE_ID). Only used when USE_VERTEX_AI is true.",
196196
)
197197
USE_PERMISSIVE_SAFETY: bool = Field(
198-
default=os.getenv("USE_PERMISSIVE_SAFETY", "false").lower() == "true",
198+
default=os.getenv("GOOGLE_USE_PERMISSIVE_SAFETY", "false").lower()
199+
== "true",
199200
description="Use permissive safety settings for content generation.",
200201
)
201202
MODEL_CACHE_TTL: int = Field(
@@ -211,6 +212,13 @@ class Valves(BaseModel):
211212
description="Default system prompt applied to all chats. If a user-defined system prompt exists, "
212213
"this is prepended to it. Leave empty to disable.",
213214
)
215+
ENABLE_FORWARD_USER_INFO_HEADERS: bool = Field(
216+
default=os.getenv(
217+
"GOOGLE_ENABLE_FORWARD_USER_INFO_HEADERS", "false"
218+
).lower()
219+
== "true",
220+
description="Whether to forward user information headers.",
221+
)
214222

215223
# Image Processing Configuration
216224
IMAGE_MAX_SIZE_MB: float = Field(
@@ -565,8 +573,44 @@ def _get_client(self) -> genai.Client:
565573
)
566574
else:
567575
self.log.debug("Initializing Google Generative AI client with API Key")
576+
headers = {}
577+
if (
578+
self.valves.ENABLE_FORWARD_USER_INFO_HEADERS
579+
and hasattr(self, "user")
580+
and self.user
581+
):
582+
583+
def sanitize_header_value(value: Any, max_length: int = 255) -> str:
584+
if value is None:
585+
return ""
586+
# Convert to string and remove all control characters
587+
sanitized = re.sub(r"[\x00-\x1F\x7F]", "", str(value))
588+
sanitized = sanitized.strip()
589+
return (
590+
sanitized[:max_length]
591+
if len(sanitized) > max_length
592+
else sanitized
593+
)
594+
595+
user_attrs = {
596+
"X-OpenWebUI-User-Name": sanitize_header_value(
597+
getattr(self.user, "name", None)
598+
),
599+
"X-OpenWebUI-User-Id": sanitize_header_value(
600+
getattr(self.user, "id", None)
601+
),
602+
"X-OpenWebUI-User-Email": sanitize_header_value(
603+
getattr(self.user, "email", None)
604+
),
605+
"X-OpenWebUI-User-Role": sanitize_header_value(
606+
getattr(self.user, "role", None)
607+
),
608+
}
609+
headers = {k: v for k, v in user_attrs.items() if v not in (None, "")}
568610
options = types.HttpOptions(
569-
api_version=self.valves.API_VERSION, base_url=self.valves.BASE_URL
611+
api_version=self.valves.API_VERSION,
612+
base_url=self.valves.BASE_URL,
613+
headers=headers,
570614
)
571615
return genai.Client(
572616
api_key=EncryptedStr.decrypt(self.valves.GOOGLE_API_KEY),

0 commit comments

Comments
 (0)