92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
"""Create and deploy a Vertex AI Vector Search index."""
|
|
|
|
from typing import Annotated
|
|
|
|
import typer
|
|
from rich.console import Console
|
|
|
|
from rag_eval.config import settings as config
|
|
from vector_search.vertex_ai import GoogleCloudVectorSearch
|
|
|
|
app = typer.Typer()
|
|
|
|
|
|
@app.callback(invoke_without_command=True)
|
|
def create(
|
|
path: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--path",
|
|
"-p",
|
|
help="The GCS URI (gs://...) to the directory containing your embedding JSON file(s).",
|
|
),
|
|
],
|
|
agent_name: Annotated[
|
|
str,
|
|
typer.Option(
|
|
"--agent",
|
|
"-a",
|
|
help="The name of the agent to create the index for.",
|
|
),
|
|
],
|
|
):
|
|
"""Create and deploy a Vertex AI Vector Search index for a specific agent."""
|
|
console = Console()
|
|
|
|
try:
|
|
console.print(
|
|
f"[bold green]Looking up configuration for agent '{agent_name}'...[/bold green]"
|
|
)
|
|
agent_config = config.agents.get(agent_name)
|
|
if not agent_config:
|
|
console.print(
|
|
f"[bold red]Agent '{agent_name}' not found in settings.[/bold red]"
|
|
)
|
|
raise typer.Exit(code=1)
|
|
|
|
if not agent_config.index:
|
|
console.print(
|
|
f"[bold red]Index configuration not found for agent '{agent_name}'.[/bold red]"
|
|
)
|
|
raise typer.Exit(code=1)
|
|
|
|
index_config = agent_config.index
|
|
|
|
console.print(
|
|
f"[bold green]Initializing Vertex AI client for project '{config.project_id}' in '{config.location}'...[/bold green]"
|
|
)
|
|
vector_search = GoogleCloudVectorSearch(
|
|
project_id=config.project_id,
|
|
location=config.location,
|
|
bucket=config.bucket,
|
|
index_name=index_config.name,
|
|
)
|
|
|
|
console.print(
|
|
f"[bold green]Starting creation of index '{index_config.name}'...[/bold green]"
|
|
)
|
|
console.print("This may take a while.")
|
|
vector_search.create_index(
|
|
name=index_config.name,
|
|
content_path=f"gs://{config.bucket}/{path}",
|
|
dimensions=index_config.dimensions,
|
|
)
|
|
console.print(
|
|
f"[bold green]Index '{index_config.name}' created successfully.[/bold green]"
|
|
)
|
|
|
|
console.print("[bold green]Deploying index to a new endpoint...[/bold green]")
|
|
console.print("This will also take some time.")
|
|
vector_search.deploy_index(
|
|
index_name=index_config.name, machine_type=index_config.machine_type
|
|
)
|
|
console.print("[bold green]Index deployed successfully![/bold green]")
|
|
console.print(f"Endpoint name: {vector_search.index_endpoint.display_name}")
|
|
console.print(
|
|
f"Endpoint resource name: {vector_search.index_endpoint.resource_name}"
|
|
)
|
|
|
|
except Exception as e:
|
|
console.print(f"[bold red]An error occurred: {e}[/bold red]")
|
|
raise typer.Exit(code=1)
|