# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0
"""Vertex AI Vector Search with Firestore sample.
This sample demonstrates how to use Vertex AI Vector Search with Firestore
as the document store for enterprise-scale vector similarity search.
Key Concepts (ELI5)::
┌─────────────────────┬────────────────────────────────────────────────────┐
│ Concept │ ELI5 Explanation │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Vertex AI Vector │ Google's enterprise vector search. Handles │
│ Search │ billions of items with fast nearest-neighbor. │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Firestore │ Stores the actual document content. Vector Search │
│ │ returns IDs, Firestore returns full text. │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Deployed Index │ Your vector index running on Google's servers. │
│ │ Ready to answer similarity queries 24/7. │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Endpoint │ The address where your index listens for queries. │
│ │ Like a phone number for your search service. │
├─────────────────────┼────────────────────────────────────────────────────┤
│ Distance │ How "far apart" two vectors are. Lower = more │
│ │ similar. 0 = identical match. │
└─────────────────────┴────────────────────────────────────────────────────┘
Data Flow (Enterprise RAG)::
┌─────────────────────────────────────────────────────────────────────────┐
│ HOW VERTEX VECTOR SEARCH + FIRESTORE WORK TOGETHER │
│ │
│ Query: "What movies are about time travel?" │
│ │ │
│ │ (1) Convert to embedding │
│ ▼ │
│ ┌─────────────────┐ │
│ │ Embedder │ Query → [0.3, -0.2, 0.7, ...] │
│ └────────┬────────┘ │
│ │ │
│ │ (2) Send to Vector Search endpoint │
│ ▼ │
│ ┌─────────────────┐ │
│ │ Vertex AI │ Returns: ["doc_123", "doc_456", ...] │
│ │ Vector Search │ (IDs only, ranked by similarity) │
│ └────────┬────────┘ │
│ │ │
│ │ (3) Fetch full documents from Firestore │
│ ▼ │
│ ┌─────────────────┐ │
│ │ Firestore │ Returns: full document content │
│ │ (by ID) │ "Back to the Future is a 1985 film..." │
│ └────────┬────────┘ │
│ │ │
│ │ (4) Sorted results returned │
│ ▼ │
│ ┌─────────────────┐ │
│ │ Your App │ Complete docs, ranked by relevance │
│ └─────────────────┘ │
└─────────────────────────────────────────────────────────────────────────┘
Key Features
============
| Feature Description | Example Function / Code Snippet |
|-----------------------------------------|-------------------------------------|
| Firestore Vector Search Definition | `define_vertex_vector_search_firestore`|
| Firestore Async Client Integration | `firestore.AsyncClient()` |
| Document Retrieval | `ai.retrieve()` |
| Result Ranking | Custom sorting by distance |
Testing This Demo
=================
1. **Prerequisites** - Set up GCP resources:
```bash
# Required environment variables
export LOCATION=us-central1
export PROJECT_ID=your_project_id
export FIRESTORE_COLLECTION=your_collection_name
export VECTOR_SEARCH_DEPLOYED_INDEX_ID=your_deployed_index_id
export VECTOR_SEARCH_INDEX_ENDPOINT_PATH=your_endpoint_path
export VECTOR_SEARCH_API_ENDPOINT=your_api_endpoint
# Authenticate with GCP
gcloud auth application-default login
```
2. **GCP Setup Required**:
- Create Vertex AI Vector Search index
- Deploy index to an endpoint
- Create Firestore collection with documents
- Ensure documents have matching IDs in both services
3. **Run the demo**:
```bash
cd py/samples/vertex-ai-vector-search-firestore
./run.sh
```
4. **Open DevUI** at http://localhost:4000
5. **Test the flows**:
- [ ] `retrieve_documents` - Vector similarity search
- [ ] Check results are ranked by distance
- [ ] Verify Firestore document metadata is returned
6. **Expected behavior**:
- Query is embedded and sent to Vector Search
- Similar vectors are found and IDs returned
- Firestore is queried for full document content
- Results sorted by similarity distance
"""
import os
import time
from google.cloud import aiplatform, firestore
from pydantic import BaseModel, Field
from rich.traceback import install as install_rich_traceback
from genkit.ai import Genkit
from genkit.blocks.document import Document
from genkit.core.logging import get_logger
from genkit.core.typing import RetrieverResponse
from genkit.plugins.google_genai import VertexAI
from genkit.plugins.vertex_ai import define_vertex_vector_search_firestore
install_rich_traceback(show_locals=True, width=120, extra_lines=3)
LOCATION = os.environ['LOCATION']
PROJECT_ID = os.environ['PROJECT_ID']
FIRESTORE_COLLECTION = os.environ['FIRESTORE_COLLECTION']
VECTOR_SEARCH_DEPLOYED_INDEX_ID = os.environ['VECTOR_SEARCH_DEPLOYED_INDEX_ID']
VECTOR_SEARCH_INDEX_ENDPOINT_PATH = os.environ['VECTOR_SEARCH_INDEX_ENDPOINT_PATH']
VECTOR_SEARCH_API_ENDPOINT = os.environ['VECTOR_SEARCH_API_ENDPOINT']
firestore_client = firestore.AsyncClient(project=PROJECT_ID)
aiplatform.init(project=PROJECT_ID, location=LOCATION)
logger = get_logger(__name__)
ai = Genkit(plugins=[VertexAI()])
# Define Vertex AI Vector Search with Firestore
define_vertex_vector_search_firestore(
ai,
name='my-vector-search',
embedder='vertexai/text-embedding-004',
embedder_options={
'task': 'RETRIEVAL_DOCUMENT',
'output_dimensionality': 128,
},
firestore_client=firestore_client,
collection_name=FIRESTORE_COLLECTION,
)
class QueryFlowInputSchema(BaseModel):
"""Input schema."""
query: str = Field(default='document 1', description='Search query text')
k: int = Field(default=5, description='Number of results to return')
class QueryFlowOutputSchema(BaseModel):
"""Output schema."""
result: list[dict[str, object]]
length: int
time: int
@ai.flow(name='queryFlow')
async def query_flow(_input: QueryFlowInputSchema) -> QueryFlowOutputSchema:
"""Executes a vector search with VertexAI Vector Search."""
start_time = time.time()
query_document = Document.from_text(text=_input.query)
query_document.metadata = {
'api_endpoint': VECTOR_SEARCH_API_ENDPOINT,
'index_endpoint_path': VECTOR_SEARCH_INDEX_ENDPOINT_PATH,
'deployed_index_id': VECTOR_SEARCH_DEPLOYED_INDEX_ID,
}
result: RetrieverResponse = await ai.retrieve(
retriever='my-vector-search',
query=query_document,
options={'limit': 10},
)
end_time = time.time()
duration = int(end_time - start_time)
result_data = []
for doc in result.documents:
metadata = doc.metadata or {}
result_data.append({
'id': metadata.get('id'),
'text': doc.content[0].root.text if doc.content and doc.content[0].root.text else '',
'distance': metadata.get('distance', 0.0),
})
result_data = sorted(result_data, key=lambda x: x['distance'])
return QueryFlowOutputSchema(
result=result_data,
length=len(result_data),
time=duration,
)
async def main() -> None:
"""Main function."""
query_input = QueryFlowInputSchema(
query='Content for doc',
k=3,
)
result = await query_flow(query_input)
await logger.ainfo(str(result))
if __name__ == '__main__':
ai.run_main(main())
