-- Find products similar to a user's last purchase

-- Quick function to generate vector arrays
DEFINE FUNCTION fn::rand_array() {
    [[],[],[],[],[]].map(|$i| rand::int(0, 10));
};

-- Create a user that purchased two products, and some
-- vector representing its location in semantic space
CREATE user:one;
CREATE product:one SET vector = [1,2,3,4,5];
CREATE product:two SET vector = [6,5,4,3,2];

RELATE user:one->purchased->product:one SET at = time::now();
RELATE user:one->purchased->product:two SET at = time::now();

-- Make some other random users and products
FOR $_ IN 0..10 {
    LET $user = UPSERT ONLY user;
    LET $product = UPSERT ONLY type::thing("product", rand::string(7)) SET vector = fn::rand_array();
    RELATE $user->purchased->$product SET at = time::now();
};

-- Get the last user's purchase info by datetime
LET $last_purchase = user:one->purchased.at.last();
-- And the most recently purchased product from that
LET $last_product = (user:one->purchased[WHERE at = $last_purchase]->product)[0];

-- Then get the most similar products
-- $last_product will show up as most similar so
-- slice starting at index 1
(SELECT 
    id, 
    vector, 
    vector::similarity::cosine($last_product.vector, vector) AS similarity
FROM product
ORDER BY similarity DESC LIMIT 3)[1..];

SELECT VALUE
    (
        SELECT *,
            vector::similarity::cosine(embedding, $parent.embedding) AS similarity
        FROM array::distinct(->tagged_with->$tag<-tagged_with<-document)
    ) AS siblings
FROM ONLY $record
FETCH siblings

# DB connection
conn = Surreal(url)
conn.signin({"username": user, "password": password})
conn.use(ns, db)

# Vector Store
vector_store = SurrealDBVectorStore(
    OllamaEmbeddings(model="llama3.2"),
    conn
)

# Graph Store
graph_store = SurrealDBGraph(conn)

documents.append(Document(page_content=chunk, metadata=...))

# This calculates the embeddings and inserts the documents into the DB
vector_store.add_documents(documents)

# Find nodes and edges (Product -> tagged_with -> Tag)
tag_node = Node(id=tag.name, type="Tag", properties=asdict(tag))
product_node = Node(id=product.id, type="Product", properties=addict(product)

# Edges
relationships.append(
    Relationship(source=product_node, target=symptom_node, type="tagged_with")
)
graph_documents.append(
    GraphDocument(nodes=nodes, relationships=relationships, source=doc)
)

# Store the graph
graph_store.add_graph_documents(graph_documents, include_source=True)

chat_model = ChatOllama(model="llama3.2", temperature=0)

# -- Find relevant docs
docs = vector_search(query, vector_store, k=3)

# -- Query the graph
chain = SurrealDBGraphQAChain.from_llm(chat_model, graph=graph_store)
response = chain.invoke({"query": docs_into_str(docs))
# docs_into_str is a pseudo-function to concatente document information
# in a single string to send to the LLM