upsert.py

# 1. IMPORTS
from dotenv import load_dotenv
load_dotenv()

from transformers import logging as hf_logging
hf_logging.set_verbosity_error()
hf_logging.disable_progress_bar()

from sentence_transformers import SentenceTransformer
from databricks.vector_search.client import VectorSearchClient
import time

# 2. CONFIG
INDEX_NAME = "dv_projc.live_chat.messages_feelings_index"
ENDPOINT_NAME = "hate-or-love-sv"

# 3. MODELO
model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')

# 4. DADOS
texts = [
    {"id": "1", "text": "Estou muito feliz hoje"},
    {"id": "2", "text": "Que dia maravilhoso"},
    {"id": "3", "text": "A vida é linda"},
    {"id": "4", "text": "Odeio tudo isso"},
    {"id": "5", "text": "Isso me deixa com raiva"},
    {"id": "6", "text": "Estou muito irritado"}
]

# 5. EMBEDDINGS
embeddings = model.encode([t["text"] for t in texts])
print("Dimensão do embedding:", len(embeddings[0]))

data = [
    {
        "id": t["id"],
        "text": t["text"],
        "embedding": embeddings[i].tolist()
    }
    for i, t in enumerate(texts)
]

# 6. CLIENTE
client = VectorSearchClient(disable_notice=True)

# 7. AGUARDAR ENDPOINT
print("\nAguardando endpoint ficar ONLINE...")

MAX_WAIT = 15  # minutos
for i in range(MAX_WAIT * 6):  # checa a cada 10s
    endpoint = client.get_endpoint(ENDPOINT_NAME)
    state = endpoint.get("endpoint_status", {}).get("state", "")
    
    print(f"[{i*10}s] Endpoint state: {state}")
    
    if state == "ONLINE":
        print("Endpoint pronto!")
        break
    
    time.sleep(10)
else:
    raise TimeoutError("Endpoint não ficou ONLINE a tempo")

# 8. CRIAR ÍNDICE (SE NÃO EXISTIR)
print("\nCriando índice...")

try:
    client.create_direct_access_index(
        endpoint_name=ENDPOINT_NAME,
        index_name=INDEX_NAME,
        primary_key="id",
        embedding_dimension=384,
        embedding_vector_column="embedding",
        schema={
            "id": "string",
            "text": "string",
            "embedding": "array<float>"
        }
    )
    print("Índice criado!")
except Exception as e:
    print("Índice pode já existir:", e)

# 9. OBTER ÍNDICE
index = client.get_index(ENDPOINT_NAME, INDEX_NAME)

# 10. AGUARDAR INDEX
print("\nAguardando índice ficar ONLINE...")

for i in range(MAX_WAIT * 6):
    info = index.describe()
    status = info.get("status", {})
    state = status.get("detailed_state", "")
    msg = status.get("message", "")
    
    print(f"[{i*10}s] Estado: {state} | {msg}")
    
    if "ONLINE" in state:
        print("Índice pronto!")
        break
    
    if "FAILED" in state:
        raise RuntimeError(f"Falha no índice: {msg}")
    
    time.sleep(10)
else:
    raise TimeoutError("Índice não ficou ONLINE a tempo")

# 11. UPSERT
print("\nInserindo dados...")

index.upsert(data)

print("Dados inseridos!")