@trustgraph/react-state 1.4.3 → 1.4.5

package/dist/index.cjs

@@ -1399,53 +1399,36 @@ const useGraphSubgraph = ({ entityUri, flow, collection, }) => {
 };
 
 /**
- * Custom hook for managing token cost operations
- * Provides functionality for fetching, deleting, and updating token costs
- * for AI models
- * @returns {Object} Token cost state and operations
+ * Custom hook for querying graph embeddings
+ * Finds graph entities similar to the provided embedding vectors
  */
-const useGraphEmbeddings = ({ flow, vecs, limit, collection }) => {
-    // WebSocket connection for communicating with the configuration service
+const useGraphEmbeddings = ({ flow, vecs, limit = 10, collection }) => {
     const socket = reactProvider.useSocket();
-    // Hook for displaying user notifications
     const notify = useNotification();
-    // Settings for default collection
     const { settings } = useSettings();
-    // Session state for default flow ID
     const sessionFlowId = useSessionStore((state) => state.flowId);
-    // Use explicit param if provided, otherwise fall back to session state
     const effectiveFlow = flow ?? sessionFlowId;
-    /**
-     * Query for fetching graph embeddings
-     * Uses React Query for caching and background refetching
-     */
+    const effectiveCollection = collection ?? settings.collection;
     const query = reactQuery.useQuery({
-        queryKey: ["graph-embeddings", { flow: effectiveFlow, vecs, limit }],
+        queryKey: ["graph-embeddings", { flow: effectiveFlow, vecs, limit, collection: effectiveCollection }],
+        enabled: !!vecs && vecs.length > 0 && !!effectiveFlow,
         queryFn: () => {
             return socket
                 .flow(effectiveFlow)
-                .graphEmbeddingsQuery(vecs, limit, collection || settings.collection)
-                .then((x) => {
-                return x;
-            })
+                .graphEmbeddingsQuery(vecs, limit, effectiveCollection)
                 .catch((err) => {
-                console.log("Error:", err);
                 const message = err instanceof Error ? err.message : String(err);
                 notify.error(message);
                 throw err;
             });
         },
     });
-    // Show loading indicators for long-running operations
     useActivity(query.isLoading, "Loading graph embeddings");
-    // Return token cost state and operations for use in components
     return {
-        // Token cost query state
         graphEmbeddings: query.data,
         isLoading: query.isLoading,
         isError: query.isError,
         error: query.error,
-        // Manual refetch function
         refetch: query.refetch,
     };
 };
@@ -2173,71 +2156,90 @@ const useStructuredQuery = () => {
 };
 
 // @ts-nocheck
-// React Query hooks for data fetching and mutation management
 /**
- * Custom hook for managing row embeddings query operations
- * Provides functionality for executing semantic searches on structured data indexes
- * First converts query text to embeddings, then searches for similar records
+ * Custom hook for querying row embeddings using vectors.
+ * Searches for similar records in structured data indexes.
  */
 const useRowEmbeddingsQuery = ({ flow } = {}) => {
-    // Socket connection for API calls
     const socket = reactProvider.useSocket();
     const connectionState = reactProvider.useConnectionState();
-    // Notification system for user feedback
     const notify = useNotification();
-    // Session state for current flow ID
     const sessionFlowId = useSessionStore((state) => state.flowId);
-    // Settings for default collection
     const { settings } = useSettings();
-    // Use explicit param if provided, otherwise fall back to session state
     const effectiveFlow = flow ?? sessionFlowId;
-    // Only enable operations when socket is connected and ready
     const isSocketReady = connectionState?.status === "authenticated" ||
         connectionState?.status === "unauthenticated";
-    // Mutation for executing row embeddings queries
-    const rowEmbeddingsQueryMutation = reactQuery.useMutation({
-        mutationFn: async ({ query, schemaName, collection, indexName, limit, }) => {
+    const mutation = reactQuery.useMutation({
+        mutationFn: async ({ vectors, schemaName, collection, indexName, limit = 10, }) => {
             if (!isSocketReady) {
                 throw new Error("Socket connection not ready");
             }
-            const flowApi = socket.flow(effectiveFlow);
-            // First, get embeddings for the query text
-            const vectors = await flowApi.embeddings(query);
-            // Then query row embeddings with those vectors
-            return flowApi.rowEmbeddingsQuery(vectors, schemaName, collection || settings.collection, indexName, limit || 10);
+            return socket.flow(effectiveFlow).rowEmbeddingsQuery(vectors, schemaName, collection ?? settings.collection, indexName, limit);
         },
         onError: (err) => {
-            console.log("Row embeddings query error:", err);
-            const errorMessage = err instanceof Error
-                ? err.message
-                : err?.toString() || "Unknown error";
-            notify.error(`Row embeddings query failed: ${errorMessage}`);
+            const message = err instanceof Error ? err.message : String(err);
+            notify.error(`Row embeddings query failed: ${message}`);
+        },
+        onSuccess: () => {
+        },
+    });
+    useActivity(mutation.isPending, "Querying row embeddings");
+    return {
+        executeQuery: mutation.mutate,
+        executeQueryAsync: mutation.mutateAsync,
+        isExecuting: mutation.isPending,
+        error: mutation.error,
+        matches: mutation.data ?? [],
+        hasResults: (mutation.data?.length ?? 0) > 0,
+        reset: mutation.reset,
+        isReady: isSocketReady,
+    };
+};
+
+// @ts-nocheck
+/**
+ * Custom hook for querying document chunks using vectors.
+ * Searches for document chunks with similar embeddings.
+ */
+const useDocumentEmbeddingsQuery = ({ flow } = {}) => {
+    const socket = reactProvider.useSocket();
+    const connectionState = reactProvider.useConnectionState();
+    const notify = useNotification();
+    const sessionFlowId = useSessionStore((state) => state.flowId);
+    const { settings } = useSettings();
+    const effectiveFlow = flow ?? sessionFlowId;
+    const isSocketReady = connectionState?.status === "authenticated" ||
+        connectionState?.status === "unauthenticated";
+    const mutation = reactQuery.useMutation({
+        mutationFn: async ({ vectors, user, collection, limit = 10, }) => {
+            if (!isSocketReady) {
+                throw new Error("Socket connection not ready");
+            }
+            return socket.flow(effectiveFlow).documentEmbeddingsQuery(vectors, user ?? settings.user, collection ?? settings.collection, limit);
+        },
+        onError: (err) => {
+            const message = err instanceof Error ? err.message : String(err);
+            notify.error(`Document embeddings query failed: ${message}`);
         },
         onSuccess: (data) => {
-            if (data && data.length > 0) {
-                notify.success(`Found ${data.length} matching record${data.length !== 1 ? 's' : ''}`);
+            const count = Array.isArray(data) ? data.length : 0;
+            if (count > 0) {
+                notify.success(`Found ${count} matching chunk${count !== 1 ? "s" : ""}`);
             }
             else {
-                notify.info("No matching records found");
+                notify.info("No matching document chunks found");
             }
         },
     });
-    // Show loading indicator for row embeddings query operations
-    useActivity(rowEmbeddingsQueryMutation.isPending, "Executing row embeddings query");
-    // Return the public API for the hook
+    useActivity(mutation.isPending, "Searching document chunks");
     return {
-        // Query execution
-        executeQuery: rowEmbeddingsQueryMutation.mutate,
-        executeQueryAsync: rowEmbeddingsQueryMutation.mutateAsync,
-        // Query state
-        isExecuting: rowEmbeddingsQueryMutation.isPending,
-        error: rowEmbeddingsQueryMutation.error,
-        // Results
-        matches: rowEmbeddingsQueryMutation.data || [],
-        hasResults: (rowEmbeddingsQueryMutation.data?.length || 0) > 0,
-        // Reset function to clear previous results
-        reset: rowEmbeddingsQueryMutation.reset,
-        // Socket readiness
+        executeQuery: mutation.mutate,
+        executeQueryAsync: mutation.mutateAsync,
+        isExecuting: mutation.isPending,
+        error: mutation.error,
+        results: mutation.data ?? [],
+        hasResults: Array.isArray(mutation.data) && mutation.data.length > 0,
+        reset: mutation.reset,
         isReady: isSocketReady,
     };
 };
@@ -2278,7 +2280,6 @@ const useRowsQuery = ({ flow } = {}) => {
             notify.error(`GraphQL query failed: ${errorMessage}`);
         },
         onSuccess: () => {
-            notify.success("GraphQL query executed successfully");
         },
     });
     // Show loading indicator for query operations
@@ -4690,6 +4691,7 @@ exports.useChat = useChat;
 exports.useChatSession = useChatSession;
 exports.useCollections = useCollections;
 exports.useConversation = useConversation;
+exports.useDocumentEmbeddingsQuery = useDocumentEmbeddingsQuery;
 exports.useEmbeddings = useEmbeddings;
 exports.useEntityDetail = useEntityDetail;
 exports.useFlowBlueprints = useFlowBlueprints;