Build a reliable GenAI search system with Gemini Grounding and Vertex AI

gcloud components update
gcloud components install beta

pip install vertexai==1.71.1 google-cloud-aiplatform==1.72.0 google-auth==2.21.0 google-auth-oauthlib==1.0.0 wandb weave


import vertexai
from vertexai.generative_models import (
    GenerationConfig,
    GenerativeModel,
    Tool,
)
from google.cloud.aiplatform_v1beta1.types import tool as tool_types
import json
import weave

weave.init("gemini_grounding")

# Initialize Vertex AI with your project details
PROJECT_ID = "dsports-6ab79"
LOCATION = "us-central1"
vertexai.init(project=PROJECT_ID, location=LOCATION)

# Define your threshold for dynamic retrieval
DYNAMIC_THRESHOLD = 0.68 # Adjust this value as needed

@weave.op
def perform_inference_with_grounding(prompt):
    # Initialize the Gemini model
    model = GenerativeModel("gemini-1.5-pro-002")

    # Set up dynamic retrieval config with a threshold for Google Search grounding
    dynamic_retrieval_config = tool_types.DynamicRetrievalConfig(
        mode=tool_types.DynamicRetrievalConfig.Mode.MODE_DYNAMIC,
        dynamic_threshold=DYNAMIC_THRESHOLD
    )

    # Configure GoogleSearchRetrieval with dynamic retrieval config
    google_search_retrieval = tool_types.GoogleSearchRetrieval(
        dynamic_retrieval_config=dynamic_retrieval_config
    )

    # Create the Tool using the google_search_retrieval
    raw_tool = tool_types.Tool(
        google_search_retrieval=google_search_retrieval
    )

    # Create the Tool object using _from_gapic method
    tool = Tool._from_gapic(raw_tool=raw_tool)

    # Generate content with grounding
    response = model.generate_content(
        prompt,
        tools=[tool],
        generation_config=GenerationConfig(
            temperature=0.0,
        ),
    )

    # Structure metadata and collect output text
    output_texts = []
    flat_metadata = []
    prediction_score = None  # Initialize prediction score variable

    # Loop over all candidates in the response
    if response and response.candidates:
        for candidate in response.candidates:
            output_text = "".join(part.text for part in candidate.content.parts)
            output_texts.append(output_text)

            # Extract grounding data
            grounding_chunks = candidate.grounding_metadata.grounding_chunks
            grounding_supports = candidate.grounding_metadata.grounding_supports

            # Retrieve prediction score if available
            # prediction_score = candidate.grounding_metadata.retrieval_metadata.get("googleSearchDynamicRetrievalScore")
            prediction_score = 0.0
            if candidate.grounding_metadata.retrieval_metadata:
                prediction_score = candidate.grounding_metadata.retrieval_metadata.google_search_dynamic_retrieval_score

            for support in grounding_supports:
                # For each segment, create a flat structure with source details
                for i, index in enumerate(support.grounding_chunk_indices):
                    source_info = grounding_chunks[index].web
                    flat_metadata.append({
                        "text_segment": support.segment.text,
                        "source_title": source_info.title,
                        "source_url": source_info.uri,
                        "confidence": support.confidence_scores[i]
                    })

    return output_texts, flat_metadata, prediction_score

# Example usage
prompt = "what is a fibonacci sequence?"
output_texts, flat_metadata, prediction_score = perform_inference_with_grounding(prompt)

# Print output texts, structured flat metadata, and prediction score as JSON
print("Output Texts:")
print(json.dumps(output_texts, indent=2))

print("\nFlat Metadata:")
print(json.dumps(flat_metadata, indent=2))

print("\nPrediction Score:")
print(prediction_score)


import os
import subprocess
from google.cloud import storage
import time
# Set your Google Cloud project details and bucket name
PROJECT_ID = "dsports-6ab79"
BUCKET_NAME = "groundinggemini"

# Initialize the Cloud Storage client
storage_client = storage.Client(project=PROJECT_ID)
bucket = storage_client.bucket(BUCKET_NAME)


def download_weave_docs(repo_dir="weave_docs", retries=3):
    """Download Weave documentation using sparse checkout with shallow clone and retry logic."""
    success = False
    attempt = 0

    while not success and attempt < retries:
        attempt += 1
        print(f"Attempt {attempt} to download Weave documentation...")
        if not os.path.exists(repo_dir):
            os.makedirs(repo_dir)
        try:
            # Initialize the Git repository with sparse checkout
            subprocess.run(["git", "init"], cwd=repo_dir, check=True)
            subprocess.run(["git", "remote", "add", "origin", "https://github.com/wandb/weave.git"], cwd=repo_dir, check=True)
            subprocess.run(["git", "config", "core.sparseCheckout", "true"], cwd=repo_dir, check=True)
            subprocess.run(["git", "sparse-checkout", "set", "docs/docs/guides/tracking"], cwd=repo_dir, check=True)
            
            # Use a shallow clone by fetching only the latest commit
            subprocess.run(["git", "fetch", "--depth", "1", "origin", "master"], cwd=repo_dir, check=True)
            subprocess.run(["git", "checkout", "master"], cwd=repo_dir, check=True)
            
            success = True
            print("Weave documentation downloaded.")
        except subprocess.CalledProcessError:
            print("Download failed, retrying...")
            time.sleep(3)  # Wait a bit before retrying
            # Clean up any partial downloads to start fresh
            subprocess.run(["rm", "-rf", repo_dir])
    
    if not success:
        raise Exception("Failed to download Weave documentation after multiple attempts.")


def convert_to_text(filepath):
    """Attempts to read any file and save it as a .txt file."""
    txt_filepath = f"{filepath}.txt"
    
    try:
        # Try reading as plain text
        with open(filepath, 'r', encoding='utf-8') as file:
            text_content = file.read()

    except (UnicodeDecodeError, IOError):
        # Log and skip files that can't be read as text (e.g., binary files)
        print(f"Skipping unsupported or binary file: {filepath}")
        return None
    
    # Save the read content to a .txt file
    with open(txt_filepath, 'w', encoding='utf-8') as txt_file:
        txt_file.write(text_content)
    
    return txt_filepath

def upload_to_gcs(txt_filepath):
    """Uploads a .txt file to the specified Google Cloud Storage bucket."""
    blob_name = os.path.basename(txt_filepath)  # Use filename as blob name
    blob = bucket.blob(blob_name)
    blob.upload_from_filename(txt_filepath)
    print(f"Uploaded {txt_filepath} to gs://{BUCKET_NAME}/{blob_name}")

# Step 1: Download Weave docs
download_weave_docs()

# Step 2: Define the directory containing the downloaded files
documents_dir = "weave_docs/docs/docs/guides/tracking"  # Path to the downloaded Weave documentation

# Step 3: Process each document in the specified directory
for filename in os.listdir(documents_dir):
    filepath = os.path.join(documents_dir, filename)
    if os.path.isfile(filepath):
        txt_filepath = convert_to_text(filepath)  # Attempt to convert to .txt
        if txt_filepath:
            upload_to_gcs(txt_filepath)  # Upload the .txt file to Cloud Storage
            os.remove(txt_filepath)  # Optionally, delete the .txt file after uploading


import vertexai
from vertexai.preview.generative_models import (
    GenerationConfig,
    GenerativeModel,
    Tool,
    grounding,
)
import json
import weave
weave.init("gemini_vertex_search_grounding")

# Initialize Vertex AI with your project details
PROJECT_ID = "dsports-6ab79"
DATA_STORE_ID = "gemini-grounding_1731397909128"
LOCATION = "us-central1"

vertexai.init(project=PROJECT_ID, location=LOCATION)

@weave.op
def perform_vertex_search_inference(prompt):
    # Initialize the Gemini model
    model = GenerativeModel("gemini-1.5-pro-002")
    
    # Set up the Vertex AI Search grounding tool
    tool = Tool.from_retrieval(
        grounding.Retrieval(
            grounding.VertexAISearch(
                datastore=DATA_STORE_ID,
                project=PROJECT_ID,
                location="global",
            )
        )
    )
    
    # Generate content with grounding
    response = model.generate_content(
        prompt,
        tools=[tool],
        generation_config=GenerationConfig(
            temperature=0.0,
        ),
    )
    
    # Structure metadata and collect output text
    output_texts = []
    flat_metadata = []
    
    # Loop over all candidates in the response
    if response and response.candidates:
        for candidate in response.candidates:
            # Get content parts and join them
            content_text = candidate.content.text if hasattr(candidate.content, 'text') else ""
            output_texts.append(content_text)
            
            # Extract grounding metadata
            if hasattr(candidate, 'grounding_metadata'):
                grounding_chunks = candidate.grounding_metadata.grounding_chunks
                grounding_supports = candidate.grounding_metadata.grounding_supports
                
                for support in grounding_supports:
                    # For each segment, create a flat structure with source details
                    for i, chunk_idx in enumerate(support.grounding_chunk_indices):
                        chunk = grounding_chunks[chunk_idx]
                        # Check if chunk has retrieved_context
                        if hasattr(chunk, 'retrieved_context'):
                            source_info = chunk.retrieved_context
                            flat_metadata.append({
                                "text_segment": support.segment.text,
                                "source_uri": source_info.uri,
                                "source_title": source_info.title,
                                "confidence": support.confidence_scores[i] if hasattr(support, 'confidence_scores') else None
                            })
    
    return output_texts, flat_metadata

# Example usage
def run_vertex_search_query(prompt):
    output_texts, flat_metadata = perform_vertex_search_inference(prompt)
    
    # Print results in the same format as the example
    print("Model Response:")
    for text in output_texts:
        print(text)
    print("\nRetrieval Queries:", prompt)
    
    if flat_metadata:
        print("\nGrounding Metadata:")
        print(json.dumps(flat_metadata, indent=2))
    
    return output_texts, flat_metadata

# Example query
prompt = "What is wandb weave?"
results = run_vertex_search_query(prompt)