-SKIP
-DATASET CSV llm/embeddings/googlevertex/

--

-CASE CreateEmbeddingGoogleVertex

-STATEMENT load extension "${KUZU_ROOT_DIRECTORY}/extension/llm/build/libllm.kuzu_extension"
---- ok

-STATEMENT MATCH (e:embedding) RETURN array_cosine_similarity(e.embedding, create_embedding(e.text, 'google-vertex', 'gemini-embedding-001', 'us-central1')) >= 0.9 AS similarity;
---- 6
True
True
True
True
True
True

-CASE TryConfiguration

-STATEMENT load extension "${KUZU_ROOT_DIRECTORY}/extension/llm/build/libllm.kuzu_extension"
---- ok

-STATEMENT RETURN CREATE_EMBEDDING("Hello world", "google-vertex", "gemini-embedding-001", "us-central1");
---- ok

-STATEMENT return create_embedding('Hello World', 'google-vertex', 'gemini-embedding-001', 256, 'us-east1');
---- ok

-CASE SemanticDistanceCheck

-STATEMENT load extension "${KUZU_ROOT_DIRECTORY}/extension/llm/build/libllm.kuzu_extension"
---- ok

-LOG Similar Text 1
-STATEMENT WITH create_embedding('Who is the author of 1984?', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t1, create_embedding('Who wrote the novel 1984?', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t2 RETURN array_cosine_similarity(t1, t2) >= 0.8;
---- 1
True

-LOG Similar Text 2
-STATEMENT WITH create_embedding('What is the capital of France?', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t1, create_embedding('Which city is the capital of France?', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t2 RETURN array_cosine_similarity(t1, t2) >= 0.8;
---- 1
True

-LOG Dissimilar Text 1
-STATEMENT WITH create_embedding('How to train a dog', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t1, create_embedding('What is the capital of France', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t2 RETURN array_cosine_similarity(t1, t2) >= 0.8;
---- 1
False

-LOG Dissimilar Text 2
-STATEMENT WITH create_embedding('How do Unix file permissions work?', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t1, create_embedding('History of the Roman Empire', 'google-vertex', 'gemini-embedding-001', 'us-central1') AS t2 RETURN array_cosine_similarity(t1, t2) >= 0.8;
---- 1
False

-CASE UseWithVectorExtension

-STATEMENT load extension "${KUZU_ROOT_DIRECTORY}/extension/llm/build/libllm.kuzu_extension"
---- ok
-STATEMENT load extension "${KUZU_ROOT_DIRECTORY}/extension/vector/build/libvector.kuzu_extension"
---- ok

-STATEMENT CREATE NODE TABLE Book(id SERIAL PRIMARY KEY, title STRING, title_embedding FLOAT[3072], published_year INT64);
---- ok
-STATEMENT CREATE
    (:Book {title: 'The Quantum World', published_year: 2004}),
    (:Book {title: 'Chronicles of the Universe', published_year: 2022}),
    (:Book {title: 'Learning Machines', published_year: 2019}),
    (:Book {title: 'Echoes of the Past', published_year: 2010}),
    (:Book {title: 'The Dragon Call', published_year: 2015});
---- ok
-STATEMENT MATCH (b:Book) WITH b, create_embedding(b.title, 'google-vertex', 'gemini-embedding-001', 'us-central1') AS emb SET b.title_embedding = emb;
---- ok
-STATEMENT CALL CREATE_VECTOR_INDEX('Book', 'title_vec_index', 'title_embedding');
---- ok

-STATEMENT CALL QUERY_VECTOR_INDEX(
        'Book',
        'title_vec_index',
        create_embedding('quantum machine learning', 'google-vertex', 'gemini-embedding-001', 'us-central1'),
        2
    )
    RETURN node.title
    ORDER BY distance;
-CHECK_ORDER
---- 2
Learning Machines
The Quantum World
