diff --git a/examples/bookstore.py b/examples/bookstore.py new file mode 100644 index 00000000..b228c940 --- /dev/null +++ b/examples/bookstore.py @@ -0,0 +1,116 @@ +# N::Chapter { +# chapter_index: I64 +# } +# +# N::SubChapter { +# title: String, +# content: String +# } +# +# E::Contains { +# From: Chapter, +# To: SubChapter, +# Properties: { +# } +# } +# +# V::Embedding { +# chunk: String +# } +# +# E::EmbeddingOf { +# From: SubChapter, +# To: Embedding, +# Properties: { +# chunk: String +# } +# } + +# QUERY loaddocs_rag(chapters: [{ id: I64, subchapters: [{ title: String, content: String, chunks: [{chunk: String, vector: [F64]}]}] }]) => +# FOR {id, subchapters} IN chapters { +# chapter_node <- AddN({ chapter_index: id }) +# FOR {title, content, chunks} IN subchapters { +# subchapter_node <- AddN({ title: title, content: content }) +# AddE::From(chapter_node)::To(subchapter_node) +# FOR {chunk, vector} IN chunks { +# vec <- AddV(vector) +# AddE({chunk: chunk})::From(subchapter_node)::To(vec) +# } +# } +# } +# RETURN "Success" + +# QUERY searchdocs_rag(query: [F64], k: I32) => +# vecs <- SearchV(query, k) +# subchapters <- vecs::In +# RETURN subchapters::{title, content} + +# QUERY edge_node() => +# e <- N::OutE +# RETURN e + +# QUERY edge_node_id(id: ID) => +# e <- N::OutE(id) +# RETURN e +# + +import helix + +db = helix.Db() + +class Chapter(db.Node): + index: helix.I64 + +db.index(Chapter.index, unique=True) + +class SubChapter(db.Node): + title: helix.String + content: helix.String + + embedding: SubChapterEmbedding + +class SubChapterEmbedding(db.Vector(dimensions=1536, hnsw=helix.cosine)): + pass + +class Contains(db.Edge[Chapter, SubChapter]): + pass + +class ArgChapter(helix.Struct): + id: helix.I64 + subchapters: helix.List[ArgSubchapter] + +class ArgSubchapter(helix.Struct): + title: helix.String + content: helix.String + chunk: helix.Vector + +@db.query +def loaddocs_rag(chapters: helix.List[ArgChapter]) -> helix.String: + for c in chapters: + c_node = db.add_node(Chapter(index=c.id)) + + for sc in c.subchapters: + sc_node = db.add_node(SubChapter( + title=sc.title, + content=sc.content, + embedding=SubChapterEmbedding(sc.chunk) + )) + + db.add_edge(Contains(from=c_node, to=sc_node)) + + return "Success" + +@db.query +def searchdocs_rag(query: helix.Vector, k: helix.I32) -> helix.Iterator[helix.Map[helix.String, helix.Value]]: + # TODO + vecs = db.search_vector(query, k) + chapters = vecs.incoming_nodes[Contains] + return chapters.map(lambda c: {"index": c.index}) + +@db.query +def edge_node() -> helix.Iterator[Contains]: + return db.nodes[Chapter].outgoing_edges[Contains] + +@db.query +def edge_node_id(id: helix.Id[Chapter]) -> helix.Iterator[Contains]: + return db.nodes[Chapter](id=id).outgoing_edges[Contains] diff --git a/examples/bookstore.ts b/examples/bookstore.ts new file mode 100644 index 00000000..26458c9a --- /dev/null +++ b/examples/bookstore.ts @@ -0,0 +1,72 @@ +import hx from "helix"; + +const schema = hx.schema(); + +const Chapter = schema.defineNode({ + index: I64, +}); + +schema.index(Chapter.index, { unique: true }); + +const SubChapter = schema.defineNode({ + title: hx.String, + content: hx.String, + + embedding: SubChapterEmbedding, +}); + +const SubChapterEmbedding = schema.defineVector({ + dimensions: 1536, + hnsw: hx.cosine, +}); + +const Contains = schema.defineEdge({ from: Chapter, to: SubChapter }); + +const ArgChapter = hx.Struct({ + id: hx.I64, + subchapters: hx.List(ArgSubchapter), +}); + +const ArgSubChapter = hx.Struct({ + title: hx.String, + content: hx.String, + chunk: hx.Vector, +}); + +const loadDocsRag = schema.query({ + name: "loaddocs_rag", + arguments: [hx.List(ArgChapter)], + returns: hx.String, +}, (db, [chapters]) => { + chapters.forEach((c) => { + const cNode = db.addNode(Chapter({ index: c.id })); + + c.subchapters.forEach((sc) => { + const scNode = db.addNode(SubChapter({ + title: sc.title, + content: sc.content, + embedding: SubChapterEmbedding(sc.chunk), + })); + + db.addEdge(Contains({ from: cNode, to: scNode })); + }); + }); + + return "Success"; +}); + +const edgeNode = schema.query({ + name: "edge_node", + arguments: [], + returns: hx.Iterator(Contains), +}, (db, []) => { + return db.nodes[Chapter].outgoingEdges[Contains]; +}); + +const edgeNodeId = schema.query({ + name: "edge_node_id", + arguments: [hx.Id(Chapter)], + returns: hx.Iterator(Contains), +}, (db, [id]) => { + return db.nodes[Chapter]({ id }).outgoingEdges[Contains]; +});