vdb.py

vdb.py

import os
import numpy as np
from databricks.vector_search.client import VectorSearchClient
from dotenv import load_dotenv

load_dotenv()

DATABRICKS_WORKSPACE_URL = os.getenv("DATABRICKS_WORKSPACE_URL")
access_token_dbx = os.getenv("access_token_dbx")
VECTOR_SEARCH_ENDPOINT = os.getenv("VECTOR_SEARCH_ENDPOINT")
VECTOR_INDEX_NAME = os.getenv("VECTOR_INDEX_NAME")  # Should be "mycat.myschema.my_vector_index2"
DELTA_TABLE_NAME = os.getenv("DELTA_TABLE_NAME")
EMBEDDING_COL_NAME = "embedding"

# --- Connect to Databricks Vector Search ---
client = VectorSearchClient(
    workspace_url=DATABRICKS_WORKSPACE_URL,
    personal_access_token=access_token_dbx
)

"""
CREATE TABLE mycat.myschema.my_dt (
    id BIGINT,
    text STRING,
    embedding ARRAY<FLOAT>
);

ALTER TABLE mycat.myschema.my_dt SET TBLPROPERTIES (delta.enableChangeDataFeed = true)
"""



print(f"Attempting to get or create index: {VECTOR_INDEX_NAME}")

try:
    index = client.get_index(
        endpoint_name=VECTOR_SEARCH_ENDPOINT,
        index_name=VECTOR_INDEX_NAME,
    )
    print(f"Successfully retrieved existing index: {VECTOR_INDEX_NAME}")

except Exception as e:
    # 2. If it throws a "RESOURCE_DOES_NOT_EXIST" error, create it
    if "RESOURCE_DOES_NOT_EXIST" in str(e):
        print(f"Index '{VECTOR_INDEX_NAME}' not found. Creating it now...")
        client.create_delta_sync_index(
            endpoint_name=VECTOR_SEARCH_ENDPOINT,
            index_name=VECTOR_INDEX_NAME,
            source_table_name=DELTA_TABLE_NAME,
            pipeline_type="TRIGGERED",
            primary_key="id",
            embedding_dimension=1536,
            embedding_vector_column=EMBEDDING_COL_NAME,
        )
        print("Index creation initiated. Waiting for it to be ready...")

        index = client.get_index(
            endpoint_name=VECTOR_SEARCH_ENDPOINT,
            index_name=VECTOR_INDEX_NAME,
        )
        index.wait_for_index_to_be_ready()  # This is a blocking call

        print(f"Successfully created and retrieved index: {VECTOR_INDEX_NAME}")
    else:
        # Some other error occurred, raise it
        print(f"An unexpected error occurred: {e}")
        raise e

# --- Query the Index ---
print("Attempting to query the index...")

# ⚠️ REMINDER: Your query vector MUST have 1536 dimensions
try:
    query_vector = np.random.rand(1536).tolist()

    columns_to_return = ["id", "text"]

    results = index.similarity_search(
        query_vector=query_vector,
        columns=columns_to_return,
        num_results=3,
    )

    print("Top results:")
    if 'result' in results and 'data_array' in results['result']:
        for r in results['result']['data_array']:
            print(r)
    else:
        print(f"Could not find results in response: {results}")

except Exception as e:
    print(f"An error occurred during similarity search: {e}")
    print("This can happen if the index is not ready yet.")
    print("Please go to the Databricks UI, find your vector search endpoint, and check the status of the index.")