First commmit

packages/vector-search/src/vector_search/cli/create.py

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