diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md index e15084f..1930999 100644 --- a/.github/copilot-instructions.md +++ b/.github/copilot-instructions.md @@ -19,7 +19,7 @@ This repo extends Open WebUI with pluggable Python pipelines and filters. If you - Azure and N8N stream with SSE (`StreamingResponse`); always close `aiohttp` `ClientSession`/response in `finally` (see `cleanup_response`). - Gemini disables streaming for image-generation models; thinking is wrapped in `
` and emitted incrementally. - Cross-component integration: - - `filters/google_search_tool.py` converts `features.web_search` → `metadata.features.google_search_tool`; Gemini reads this to enable grounding tools. + - Gemini reads `search_web` and `fetch_url` tools from `__metadata__.get("tools", {})` to enable grounding tools natively when function calling is enabled. - N8N non-streaming responses can append a tool-calls section built by `_format_tool_calls_section` with verbosity and truncation valves. ## Dev workflow (local) diff --git a/README.md b/README.md index c01fd65..d16336f 100644 --- a/README.md +++ b/README.md @@ -223,7 +223,7 @@ The functions include a built-in encryption mechanism for sensitive information: - **Video generation (Veo)**: Generate videos with Google Veo models (3.1, 3, 2). Configurable aspect ratio, resolution, duration, negative prompt, and person generation controls. Supports text-to-video and image-to-video for all supported Veo models. Videos are automatically uploaded and embedded with playback controls. - **Token usage tracking**: Returns prompt, completion, and total token counts to Open WebUI for automatic persistence in the database. - **Model whitelist & additional models**: Restrict the visible model list via `GOOGLE_MODEL_WHITELIST` and add SDK-unsupported models via `GOOGLE_MODEL_ADDITIONAL`. -- Grounding with Google Search via the [google_search_tool.py filter](./filters/google_search_tool.py) +- Grounding with Google Search and URL Context (natively supported via native tool calling) - Grounding with Vertex AI Search via the [vertex_ai_search_tool.py filter](./filters/vertex_ai_search_tool.py) - Native tool calling support - Configurable API version support diff --git a/docs/google-gemini-integration.md b/docs/google-gemini-integration.md index c378b50..632b8cd 100644 --- a/docs/google-gemini-integration.md +++ b/docs/google-gemini-integration.md @@ -55,14 +55,21 @@ This integration enables **Open WebUI** to interact with **Google Gemini** model - **Customizable Generation Settings** Use environment variables to configure token limits, temperature, etc. -- **Grounding with Google search** - Improve the accuracy and recency of Gemini responses with Google search grounding. +- **Grounding with Google search and URL Context** + Improve the accuracy and recency of Gemini responses with Google search grounding and the URL Context tool, automatically enabled when native tool calling is active and appropriate tools (`search_web`, `fetch_url`) are provided. + +- **Native tool calling support** + Leverage Google genai native function calling to orchestrate the use of tools. + +- **Native MCP Tool Support** + Directly integrate Model Context Protocol (MCP) tool sessions with Gemini's native tool calling capabilities. + +- **Extensible Native Gemini Tools** + The pipeline automatically detects and executes Open WebUI tools that return a `google.genai.types.Tool` object. This allows for seamless integration of Gemini-specific features like Vertex AI Search grounding as standalone tools. - **Ability to forward User Headers and change gemini base url** Forward user information headers (like Name, Id, Email and Role) to Google API or LiteLLM for better context and analytics. Also, change the base URL for the Google Generative AI API if needed. -- **Native tool calling support** - Leverage Google genai native function calling to orchestrate the use of tools ## Environment Variables @@ -558,11 +565,15 @@ GOOGLE_MODEL_WHITELIST="gemini-exp-1206,gemini-2.0-flash-exp,gemini-1.5-pro" ## Web search and access -[Grounding with Google search](https://ai.google.dev/gemini-api/docs/google-search) together with the [URL context tool](https://ai.google.dev/gemini-api/docs/url-context) are enabled/disabled together via the `google_search_tool` feature, which can be switched on/off in a Filter. +[Grounding with Google search](https://ai.google.dev/gemini-api/docs/google-search) and the [URL context tool](https://ai.google.dev/gemini-api/docs/url-context) are natively integrated into the Gemini pipeline. They are automatically enabled when **Native tool calling** is toggled on in Open WebUI, and the corresponding tools are available to the pipeline: + +- **`search_web`**: Maps to Google Search grounding (or Enterprise Search if configured). +- **`fetch_url`**: Maps to the URL Context tool for accessing webpage content. -For instance, the following [Filter (google_search_tool.py)](../filters/google_search_tool.py) will replace Open Web UI default web search function with Google search grounding + the URL context tool. +When these tools are used, sources and Google search queries used by Gemini will be displayed with the response. -When enabled, sources and google queries from the search used by Gemini will be displayed with the response. +> [!NOTE] +> The separate `google_search_tool` filter is no longer required for Gemini grounding, as the pipeline now handles these tools natively. ### Enterprise Search @@ -573,51 +584,42 @@ To enable Enterprise Search: 1. Set `GOOGLE_USE_ENTERPRISE_SEARCH=true` (or toggle the Valve in the UI). 2. Ensure `GOOGLE_GENAI_USE_VERTEXAI=true` (Enterprise Search is a Vertex AI feature). -When enabled, the pipeline will use the `enterprise_web_search` tool instead of the standard `google_search` tool whenever grounding is requested. +When enabled, the pipeline will use the `enterprise_web_search` tool instead of the standard `google_search` tool whenever the `search_web` tool is called. ## Grounding with Vertex AI Search -Improve the accuracy and recency of Gemini responses by grounding them with your own data in Vertex AI Search. +Improve the accuracy and recency of Gemini responses by grounding them with your own data in Vertex AI Search. This feature is implemented as a [native Gemini tool](../tools/vertex_ai_search.py). ### Configuration To enable Vertex AI Search grounding, you need to: 1. **Set up a Vertex AI Search Data Store**: Follow the [Google Cloud documentation](https://cloud.google.com/vertex-ai/docs/search/overview) to create a Data Store in Discovery Engine and ingest your documents. -2. **Provide the RAG Store Path**: The path should be in the format `projects/PROJECT/locations/LOCATION/ragCorpora/DATA_STORE_ID` or `projects/PROJECT/locations/global/collections/default_collection/dataStores/DATA_STORE_ID`. - - Set the `VERTEX_AI_RAG_STORE` environment variable, or - - Use the [Filter (vertex_ai_search_tool.py)](../filters/vertex_ai_search_tool.py) to enable the feature and optionally pass the store ID via chat metadata. -3. **Enable Vertex AI**: Set `GOOGLE_GENAI_USE_VERTEXAI=true` to use Vertex AI (required for Vertex AI Search grounding). +2. **Configure the Tool**: Enable the **Vertex AI Search** tool in Open WebUI. +3. **Set the RAG Store Path**: The path should be in the format `projects/PROJECT/locations/LOCATION/ragCorpora/DATA_STORE_ID` or `projects/PROJECT/locations/global/collections/default_collection/dataStores/DATA_STORE_ID`. + - This is configured via the `VERTEX_AI_RAG_STORE` valve on the tool itself. +4. **Enable Vertex AI**: Set `GOOGLE_GENAI_USE_VERTEXAI="true"` in the Gemini pipeline settings (required for Vertex AI Search grounding). + +### Usage -When `USE_VERTEX_AI` is `true` and `VERTEX_AI_RAG_STORE` is configured, Vertex AI Search grounding will be automatically enabled. You can also explicitly enable it via the `vertex_ai_search` feature flag. +Once the tool is configured, you can enable it for any chat by selecting the **Vertex AI Search** tool. When enabled, Gemini will use the specified Vertex AI Search Data Store to retrieve relevant information and ground its responses, providing citations to the source documents. -### Example Filter Usage +> [!NOTE] +> The separate `vertex_ai_search_tool` filter is no longer required for Vertex AI grounding, as the pipeline now handles these tools natively. -The [vertex_ai_search_tool.py](../filters/vertex_ai_search_tool.py) filter enables Vertex AI Search grounding when the `vertex_ai_search` feature is requested: +## Native tool calling support -```python -# filters/vertex_ai_search_tool.py -# ... (filter code) ... -``` +Native tool calling is enabled/disabled via the standard 'Function calling' Open Web UI toggle. -To use this filter, ensure it's enabled in your Open WebUI configuration. Then, in your chat settings or via metadata, you can enable the `vertex_ai_search` feature: - -```json -{ - "features": { - "vertex_ai_search": true - }, - "params": { - "vertex_rag_store": "projects/your-project/locations/global/collections/default_collection/dataStores/your-data-store-id" - } -} -``` +## Native MCP Tool Support -## Native tool calling support +The Google Gemini pipeline supports **Native MCP Tool Support**, allowing Gemini models to directly use tools from any connected MCP (Model Context Protocol) servers when **Native tool calling** is enabled. -Native tool calling is enabled/disabled via the standard 'Function calling' Open Web UI toggle. +When using this feature: +- The pipeline automatically detects connected MCP clients and includes their entire sessions in the Gemini tool list. +- **Important**: The standard Open Web UI **MCP function whitelist does not apply** when using native tool calling with Gemini. All tools provided by the connected MCP servers will be available to the model. ## Default System Prompt diff --git a/filters/google_search_tool.py b/filters/google_search_tool.py deleted file mode 100644 index 9d15d05..0000000 --- a/filters/google_search_tool.py +++ /dev/null @@ -1,33 +0,0 @@ -""" -title: Google Search Tool Filter for https://github.com/owndev/Open-WebUI-Functions/blob/main/pipelines/google/google_gemini.py -author: owndev, olivier-lacroix -author_url: https://github.com/owndev/ -project_url: https://github.com/owndev/Open-WebUI-Functions -funding_url: https://github.com/sponsors/owndev -version: 1.0.0 -license: Apache License 2.0 -requirements: - - https://github.com/owndev/Open-WebUI-Functions/blob/main/pipelines/google/google_gemini.py -description: Replacing web_search tool with google search grounding -""" - -import logging -from open_webui.env import SRC_LOG_LEVELS - - -class Filter: - def __init__(self): - self.log = logging.getLogger("google_ai.pipe") - self.log.setLevel(SRC_LOG_LEVELS.get("OPENAI", logging.INFO)) - - def inlet(self, body: dict) -> dict: - features = body.get("features", {}) - - # Ensure metadata structure exists and add new feature - metadata = body.setdefault("metadata", {}) - metadata_features = metadata.setdefault("features", {}) - - if features.pop("web_search"): - self.log.debug("Replacing web_search tool with google search grounding") - metadata_features["google_search_tool"] = True - return body diff --git a/filters/vertex_ai_search_tool.py b/filters/vertex_ai_search_tool.py deleted file mode 100644 index 8be1534..0000000 --- a/filters/vertex_ai_search_tool.py +++ /dev/null @@ -1,43 +0,0 @@ -""" -title: Vertex AI Search Tool Filter for https://github.com/owndev/Open-WebUI-Functions/blob/main/pipelines/google/google_gemini.py -author: owndev, eun2ce -author_url: https://github.com/owndev/ -project_url: https://github.com/owndev/Open-WebUI-Functions -funding_url: https://github.com/sponsors/owndev -version: 1.0.0 -license: Apache License 2.0 -requirements: - - https://github.com/owndev/Open-WebUI-Functions/blob/main/pipelines/google/google_gemini.py -description: Enable Vertex AI Search grounding for RAG -""" - -import logging -import os -from open_webui.env import SRC_LOG_LEVELS - - -class Filter: - def __init__(self): - self.log = logging.getLogger("google_ai.pipe") - self.log.setLevel(SRC_LOG_LEVELS.get("OPENAI", logging.INFO)) - - def inlet(self, body: dict) -> dict: - features = body.get("features", {}) - - metadata = body.setdefault("metadata", {}) - metadata_features = metadata.setdefault("features", {}) - metadata_params = metadata.setdefault("params", {}) - - if features.pop("vertex_ai_search", False): - self.log.debug("Enabling Vertex AI Search grounding") - metadata_features["vertex_ai_search"] = True - - if "vertex_rag_store" not in metadata_params: - vertex_rag_store = os.getenv("VERTEX_AI_RAG_STORE") - if vertex_rag_store: - metadata_params["vertex_rag_store"] = vertex_rag_store - else: - self.log.warning( - "vertex_ai_search enabled but vertex_rag_store not provided in params or VERTEX_AI_RAG_STORE env var" - ) - return body diff --git a/pipelines/google/google_gemini.py b/pipelines/google/google_gemini.py index 6767c43..fdaedc9 100644 --- a/pipelines/google/google_gemini.py +++ b/pipelines/google/google_gemini.py @@ -4,7 +4,7 @@ author_url: https://github.com/owndev/ project_url: https://github.com/owndev/Open-WebUI-Functions funding_url: https://github.com/sponsors/owndev -version: 1.15.1 +version: 2.0.0 required_open_webui_version: 0.9.0 license: Apache License 2.0 description: Highly optimized Google Gemini pipeline with advanced image and video generation capabilities, intelligent compression, and streamlined processing workflows. @@ -53,6 +53,7 @@ import base64 import hashlib import logging +import inspect import io import uuid import aiofiles @@ -60,7 +61,17 @@ from google import genai from google.genai import types from google.genai.errors import ClientError, ServerError, APIError -from typing import List, Union, Optional, Dict, Any, Tuple, AsyncIterator, Callable +from typing import ( + List, + Union, + Optional, + Dict, + Any, + Tuple, + AsyncIterator, + Callable, + Awaitable, +) from pydantic_core import core_schema from pydantic import BaseModel, Field, GetCoreSchemaHandler from cryptography.fernet import Fernet, InvalidToken @@ -280,10 +291,6 @@ class Valves(BaseModel): default=os.getenv("GOOGLE_CLOUD_LOCATION", "global"), description="The Google Cloud region to use with Vertex AI.", ) - VERTEX_AI_RAG_STORE: str | None = Field( - default=os.getenv("GOOGLE_VERTEX_AI_RAG_STORE"), - 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.", - ) USE_PERMISSIVE_SAFETY: bool = Field( default=os.getenv("GOOGLE_USE_PERMISSIVE_SAFETY", "false").lower() == "true", @@ -2187,12 +2194,11 @@ def _get_user_valve_value( return value return None - def _configure_generation( + async def _configure_generation( self, body: Dict[str, Any], system_instruction: Optional[str], __metadata__: Dict[str, Any], - __tools__: dict[str, Any] | None = None, __user__: Optional[dict] = None, enable_image_generation: bool = False, model_id: str = "", @@ -2403,57 +2409,47 @@ def _configure_generation( gen_config_params |= {"safety_settings": safety_settings} # Add various tools to Gemini as required - features = __metadata__.get("features", {}) - params = __metadata__.get("params", {}) tools = [] - if features.get("google_search_tool", False): - if self.valves.USE_ENTERPRISE_WEB_SEARCH: - self.log.debug("Enabling Enterprise Web Search grounding") + # metadata['tools'] is populated only in native tool calling mode, + # and contains all tools, not only user-defined tools, contrarily to __tools__ + for name, tool_def in __metadata__.get("tools", {}).items(): + if name == "search_web" and self.valves.USE_ENTERPRISE_WEB_SEARCH: + self.log.debug( + "Replacing 'search_web' with Enterprise Web Search grounding" + ) tools.append( types.Tool(enterprise_web_search=types.EnterpriseWebSearch()) ) - else: - self.log.debug("Enabling Google search grounding") + elif name == "search_web": + self.log.debug("Replacing 'search_web' with Google search grounding") tools.append(types.Tool(google_search=types.GoogleSearch())) - self.log.debug("Enabling URL context grounding") - tools.append(types.Tool(url_context=types.UrlContext())) - - if features.get("vertex_ai_search", False) or ( - self.valves.USE_VERTEX_AI - and (self.valves.VERTEX_AI_RAG_STORE or os.getenv("VERTEX_AI_RAG_STORE")) - ): - vertex_rag_store = ( - params.get("vertex_rag_store") - or self.valves.VERTEX_AI_RAG_STORE - or os.getenv("VERTEX_AI_RAG_STORE") - ) - if vertex_rag_store: - self.log.debug( - f"Enabling Vertex AI Search grounding: {vertex_rag_store}" - ) - tools.append( - types.Tool( - retrieval=types.Retrieval( - vertex_ai_search=types.VertexAISearch( - datastore=vertex_rag_store - ) - ) - ) + elif name == "fetch_url": + self.log.debug("Replacing 'fetch_url' with URL context grounding") + tools.append(types.Tool(url_context=types.UrlContext())) + elif tool_def.get("type") == "mcp": + self.log.debug(f"Adding MCP tool '{name}'") + mcp_tool = self._create_callable_from_spec( + name, tool_def["spec"], tool_def["callable"] ) + tools.append(mcp_tool) + elif ( + inspect.signature(tool_def["callable"]).return_annotation is types.Tool + ): + try: + self.log.debug(f"Getting native Gemini tool: {name}") + native_tool = await tool_def["callable"]() + if isinstance(native_tool, types.Tool): + self.log.debug(f"Adding tool '{name}'") + tools.append(native_tool) + else: + self.log.warning(f"'{name}' is not a 'types.Tool'. Skipping.") + continue + except Exception as e: + self.log.warning(f"Failed to check/execute native tool {name}: {e}") else: - self.log.warning( - "Vertex AI Search requested but vertex_rag_store not provided in params, valves, or env" - ) - - if __tools__ is not None and params.get("function_calling") == "native": - for name, tool_def in __tools__.items(): - if not name.startswith("_"): - tool = tool_def["callable"] - self.log.debug( - f"Adding tool '{name}' with signature {tool.__signature__}" - ) - tools.append(tool) + self.log.debug(f"Adding tool '{name}'") + tools.append(tool_def["callable"]) if tools: gen_config_params["tools"] = tools @@ -2462,6 +2458,84 @@ def _configure_generation( filtered_params = {k: v for k, v in gen_config_params.items() if v is not None} return types.GenerateContentConfig(**filtered_params) + @staticmethod + def _create_callable_from_spec( + name: str, spec: dict, callable_func: Callable[..., Awaitable[Any]] + ) -> Callable[..., Awaitable[Any]]: + """ + Dynamically creates a well-typed async function from an MCP-style tool specification. + This satisfies inspection-based SDKs (like Gemini) by providing proper + signatures, docstrings, and unique function names. + """ + import inspect + + description = spec.get("description", "") + parameters_spec = spec.get("parameters", spec.get("inputSchema", {})) + properties = parameters_spec.get("properties", {}) + required_params = parameters_spec.get("required", []) + + # Type mapping from JSON schema to Python + type_map = { + "string": str, + "number": float, + "integer": int, + "boolean": bool, + "object": dict, + "array": list, + } + + params = [] + doc_params = [] + + # Sort properties so required parameters come first to avoid "non-default argument follows default argument" + sorted_properties = sorted( + properties.items(), + key=lambda item: item[0] not in required_params, + ) + + for param_name, param_info in sorted_properties: + if param_name.startswith("__"): + continue + + param_type = type_map.get(param_info.get("type"), Any) + param_desc = param_info.get("description", "") + + default = inspect.Parameter.empty + if param_name not in required_params: + # If not required, default to None or a provided default + default = param_info.get("default", None) + + params.append( + inspect.Parameter( + name=param_name, + kind=inspect.Parameter.KEYWORD_ONLY, + default=default, + annotation=param_type, + ) + ) + + if param_desc: + doc_params.append(f":param {param_name}: {param_desc}") + + # Build the docstring + docstring = description + if doc_params: + docstring += "\n\n" + "\n".join(doc_params) + + # The actual wrapper function + async def wrapped_func(*args, **kwargs): + return await callable_func(*args, **kwargs) + + # Set metadata to satisfy SDK inspection + wrapped_func.__name__ = name + wrapped_func.__qualname__ = name + wrapped_func.__doc__ = docstring + wrapped_func.__signature__ = inspect.Signature( + parameters=params, return_annotation=Any + ) + + return wrapped_func + @staticmethod def _format_grounding_chunks_as_sources( grounding_chunks: list[types.GroundingChunk], @@ -3089,7 +3163,6 @@ async def pipe( body: Dict[str, Any], __metadata__: dict[str, Any], __event_emitter__: Callable, - __tools__: dict[str, Any] | None, __request__: Optional[Request] = None, __user__: Optional[dict] = None, ) -> Union[str, Dict[str, Any], AsyncIterator[Union[str, Dict[str, Any]]]]: @@ -3100,7 +3173,6 @@ async def pipe( body: The request body containing messages and other parameters. __metadata__: Request metadata __event_emitter__: Event emitter for status updates - __tools__: Available tools __request__: FastAPI request object (for image upload) __user__: User information (for image upload) @@ -3171,11 +3243,10 @@ async def pipe( # Configure generation parameters and safety settings self.log.debug(f"Supports image generation: {supports_image_generation}") - generation_config = self._configure_generation( + generation_config = await self._configure_generation( body, system_instruction, __metadata__, - __tools__, __user__, supports_image_generation, model_id, diff --git a/tools/vertex_ai_search.py b/tools/vertex_ai_search.py new file mode 100644 index 0000000..1f11a57 --- /dev/null +++ b/tools/vertex_ai_search.py @@ -0,0 +1,28 @@ +import os +from google.genai import types +from pydantic import BaseModel, Field + +class Tools: + class Valves(BaseModel): + VERTEX_AI_RAG_STORE: str = Field( + default=os.getenv("GOOGLE_VERTEX_AI_RAG_STORE", ""), + description="Vertex AI RAG Store path for grounding (e.g., projects/PROJECT/locations/LOCATION/ragCorpora/DATA_STORE_ID).", + ) + + def __init__(self): + self.valves = self.Valves() + + def vertex_ai_search(self) -> types.Tool: + """ + Enable Vertex AI Search grounding for RAG. + """ + if not self.valves.VERTEX_AI_RAG_STORE: + raise ValueError("VERTEX_AI_RAG_STORE valve is not set.") + + return types.Tool( + retrieval=types.Retrieval( + vertex_ai_search=types.VertexAISearch( + datastore=self.valves.VERTEX_AI_RAG_STORE + ) + ) + )