Merge pull request #94 from eun2ce/feat/vertex-ai-search-grounding

Add Vertex AI Search grounding

Author: owndev

README.md

@@ -159,6 +159,7 @@ The functions include a built-in encryption mechanism for sensitive information:
 - **Advanced Image Processing**: Optimized image handling with configurable compression, resizing, and quality settings.
 - **Configurable Parameters**: Environment variables for image optimization (quality, max dimensions, format conversion).
 - Grounding with Google search with [google_search_tool.py filter](./filters/google_search_tool.py)
+- Grounding with Vertex AI Search with [vertex_ai_search_tool.py filter](./filters/vertex_ai_search_tool.py)
 - Native tool calling support
 - Configurable API version support
 

docs/google-gemini-integration.md

@@ -157,6 +157,11 @@ GOOGLE_CLOUD_PROJECT="your-gcp-project-id"
 # The Google Cloud region for Vertex AI (e.g., "us-central1").
 # Defaults to "global" if not set.
 GOOGLE_CLOUD_LOCATION="your-gcp-location"
+
+# Vertex AI RAG Store path for grounding (e.g., projects/PROJECT/locations/global/collections/default_collection/dataStores/DATA_STORE_ID)
+# Optional: Can also be set via metadata params or filter
+# Auto-enabled when USE_VERTEX_AI is true and this is set
+VERTEX_AI_RAG_STORE="projects/your-project/locations/global/collections/default_collection/dataStores/your-data-store-id"
 ```
 
 > [!IMPORTANT]
@@ -179,6 +184,46 @@ For instance, the following [Filter (google_search_tool.py)](../filters/google_s
 
 When enabled, sources and google queries used by Gemini will be displayed with the response.
 
+## Grounding with Vertex AI Search
+
+Improve the accuracy and recency of Gemini responses by grounding them with your own data in Vertex AI Search.
+
+### 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).
+
+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.
+
+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
+
+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:
+
+```python
+# filters/vertex_ai_search_tool.py
+# ... (filter code) ...
+```
+
+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 tool calling support
 
 Native tool calling is enabled/disabled via the standard 'Function calling' Open Web UI toggle.

filters/vertex_ai_search_tool.py

@@ -0,0 +1,44 @@
+"""
+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
+

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.6.8
+version: 1.7.0
 required_open_webui_version: 0.6.26
 license: Apache License 2.0
 description: Highly optimized Google Gemini pipeline with advanced image generation capabilities, intelligent compression, and streamlined processing workflows.
@@ -26,6 +26,7 @@
   - Configurable safety settings with environment variable support
   - Military-grade encrypted storage of sensitive API keys
   - Intelligent grounding with Google search integration
+  - Vertex AI Search grounding for RAG
   - Native tool calling support with automatic signature management
   - Unified image processing with consolidated helper methods
   - Optimized payload creation for image generation models
@@ -177,6 +178,10 @@ 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("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("USE_PERMISSIVE_SAFETY", "false").lower() == "true",
             description="Use permissive safety settings for content generation.",
@@ -1414,6 +1419,28 @@ def _configure_generation(
             )
 
         params = __metadata__.get("params", {})
+        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}")
+                gen_config_params.setdefault("tools", []).append(
+                    types.Tool(
+                        retrieval=types.Retrieval(
+                            vertex_ai_search=types.VertexAISearch(datastore=vertex_rag_store)
+                        )
+                    )
+                )
+            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("_"):
@@ -1433,24 +1460,35 @@ def _format_grounding_chunks_as_sources(
     ):
         formatted_sources = []
         for chunk in grounding_chunks:
-            context = chunk.web or chunk.retrieved_context
-            if not context:
-                continue
-
-            uri = context.uri
-            title = context.title or "Source"
+            if hasattr(chunk, "retrieved_context") and chunk.retrieved_context:
+                context = chunk.retrieved_context
+                formatted_sources.append(
+                    {
+                        "source": {
+                            "name": getattr(context, "title", None) or "Document",
+                            "type": "vertex_ai_search",
+                            "uri": getattr(context, "uri", None),
+                        },
+                        "document": [getattr(context, "chunk_text", None) or ""],
+                        "metadata": [{"source": getattr(context, "title", None) or "Document"}],
+                    }
+                )
+            elif hasattr(chunk, "web") and chunk.web:
+                context = chunk.web
+                uri = context.uri
+                title = context.title or "Source"
 
-            formatted_sources.append(
-                {
-                    "source": {
-                        "name": title,
-                        "type": "web_search_results",
-                        "url": uri,
-                    },
-                    "document": ["Click the link to view the content."],
-                    "metadata": [{"source": title}],
-                }
-            )
+                formatted_sources.append(
+                    {
+                        "source": {
+                            "name": title,
+                            "type": "web_search_results",
+                            "url": uri,
+                        },
+                        "document": ["Click the link to view the content."],
+                        "metadata": [{"source": title}],
+                    }
+                )
         return formatted_sources
 
     async def _process_grounding_metadata(