@trustgraph/react-state 1.4.6 → 1.5.2

package/dist/index.cjs

@@ -1400,9 +1400,9 @@ const useGraphSubgraph = ({ entityUri, flow, collection, }) => {
 
 /**
  * Custom hook for querying graph embeddings
- * Finds graph entities similar to the provided embedding vectors
+ * Finds graph entities similar to the provided embedding vector
  */
-const useGraphEmbeddings = ({ flow, vecs, limit = 10, collection }) => {
+const useGraphEmbeddings = ({ flow, vec, limit = 10, collection }) => {
     const socket = reactProvider.useSocket();
     const notify = useNotification();
     const { settings } = useSettings();
@@ -1410,12 +1410,12 @@ const useGraphEmbeddings = ({ flow, vecs, limit = 10, collection }) => {
     const effectiveFlow = flow ?? sessionFlowId;
     const effectiveCollection = collection ?? settings.collection;
     const query = reactQuery.useQuery({
-        queryKey: ["graph-embeddings", { flow: effectiveFlow, vecs, limit, collection: effectiveCollection }],
-        enabled: !!vecs && vecs.length > 0 && !!effectiveFlow,
+        queryKey: ["graph-embeddings", { flow: effectiveFlow, vec, limit, collection: effectiveCollection }],
+        enabled: !!vec && vec.length > 0 && !!effectiveFlow,
         queryFn: () => {
             return socket
                 .flow(effectiveFlow)
-                .graphEmbeddingsQuery(vecs, limit, effectiveCollection)
+                .graphEmbeddingsQuery(vec, limit, effectiveCollection)
                 .catch((err) => {
                 const message = err instanceof Error ? err.message : String(err);
                 notify.error(message);
@@ -1433,24 +1433,23 @@ const useGraphEmbeddings = ({ flow, vecs, limit = 10, collection }) => {
     };
 };
 
-// Take the embeddings, and lookup entities using graph
+// Take the embedding vector, and lookup entities using graph
 // embeddings, add embedding to each entity row, just an easy
 // place to put it
 const getGraphEmbeddings = (socket, add, remove, limit, collection) => {
-    // Take the embeddings, and lookup entities using graph
-    // embeddings, add embedding to each entity row, just an easy
-    // place to put it
-    return (vecs) => {
+    // Take the embedding, and lookup entities using graph
+    // embeddings, add embedding to each entity row
+    return (vec) => {
         const act = "Graph embedding search";
         add(act);
         return socket
-            .graphEmbeddingsQuery(vecs, limit , collection)
-            .then((ents) => {
+            .graphEmbeddingsQuery(vec, limit , collection)
+            .then((matches) => {
             remove(act);
-            return ents
-                .filter((ent) => ent.t === "i")
-                .map((ent) => {
-                return { uri: ent.i, target: vecs[0] };
+            return matches
+                .filter((m) => m.entity !== null && m.entity.t === "i")
+                .map((m) => {
+                return { uri: m.entity.i, target: vec };
             });
         })
             .catch((err) => {
@@ -1531,7 +1530,7 @@ const addRowEmbeddings = (socket, add, remove) => (entities) => {
         const act = "Embeddings " + text.substring(0, 20);
         add(act);
         return socket
-            .embeddings(text)
+            .embeddings([text])
             .then((x) => {
             if (x && x.length > 0) {
                 remove(act);
@@ -1575,7 +1574,8 @@ const vectorSearch = (socket, flowId, addActivity, removeActivity, term, collect
     const searchAct = "Search: " + term;
     addActivity(searchAct);
     return api
-        .embeddings(term)
+        .embeddings([term])
+        .then((vecs) => vecs[0])
         .then(getGraphEmbeddings(api, addActivity, removeActivity, limit || 10, collection))
         .then(addRowLabels(api, addActivity, removeActivity, collection))
         .then(addRowDefinitions(api, addActivity, removeActivity, collection))
@@ -1756,7 +1756,8 @@ const useInference = ({ flow } = {}) => {
                 })
                 : await socket.flow(effectiveFlow).graphRag(input, options || {}, collection);
             // Get embeddings for entity discovery
-            const embeddings = await socket.flow(effectiveFlow).embeddings(input);
+            const allEmbeddings = await socket.flow(effectiveFlow).embeddings([input]);
+            const embeddings = allEmbeddings[0];
             // Query graph embeddings to find entities
             const entities = await socket
                 .flow(effectiveFlow)
@@ -1895,7 +1896,10 @@ const useChatSession = ({ flow } = {}) => {
             // Start embeddings activity
             addActivity(embActivity);
             // Get labels for each entity
-            const labelPromises = result.entities.map(async (entity) => {
+            const labelPromises = result.entities
+                .filter((match) => match.entity !== null)
+                .map(async (match) => {
+                const entity = match.entity;
                 const labelActivity = "Label " + getTermValue$2(entity);
                 addActivity(labelActivity);
                 try {
@@ -2314,7 +2318,7 @@ const useEmbeddings = ({ flow, term }) => {
         queryFn: () => {
             return socket
                 .flow(effectiveFlow)
-                .embeddings(term)
+                .embeddings([term])
                 .then((x) => {
                 if (x["error"]) {
                     console.log("Error:", x);