From 576afa20ee8a4312f1004524e46b7f7ceb55f74e Mon Sep 17 00:00:00 2001 From: William Guss Date: Mon, 5 Aug 2024 15:14:03 -0700 Subject: [PATCH] added websockets --- ell-studio/src/App.js | 38 +++++++--- .../src/components/HierarchicalTable.js | 13 +++- .../invocations/InvocationsTable.js | 1 - ell-studio/src/hooks/useBackend.js | 70 ++++++++++++++---- ell-studio/src/pages/LMP.js | 17 +++-- ell-studio/src/pages/Traces.js | 18 ++--- examples/sqlite_example/ell.db-shm | Bin 0 -> 32768 bytes examples/sqlite_example/ell.db-wal | 0 src/ell/stores/sql.py | 1 + src/ell/studio/__main__.py | 23 +++++- src/ell/studio/data_server.py | 58 ++++++++++----- tailwind.config.js | 9 +++ 12 files changed, 183 insertions(+), 65 deletions(-) create mode 100644 examples/sqlite_example/ell.db-shm create mode 100644 examples/sqlite_example/ell.db-wal diff --git a/ell-studio/src/App.js b/ell-studio/src/App.js index 083d4fb8..ec74947b 100644 --- a/ell-studio/src/App.js +++ b/ell-studio/src/App.js @@ -8,22 +8,41 @@ import Traces from './pages/Traces'; import { ThemeProvider } from './contexts/ThemeContext'; import './styles/globals.css'; import './styles/sourceCode.css'; +import { useWebSocketConnection } from './hooks/useBackend'; +import { Toaster, toast } from 'react-hot-toast'; + +const WebSocketConnectionProvider = ({children}) => { + const { isConnected } = useWebSocketConnection(); + + React.useEffect(() => { + if (isConnected) { + toast.success('Store connected', { + duration: 1000, + }); + } else { + toast('Connecting to store...', { + icon: '🔄', + duration: 500, + }); + } + }, [isConnected]); + + return ( + <> + {children} + + + ); +}; // Create a client -const queryClient = new QueryClient({ - defaultOptions: { - queries: { - refetchOnWindowFocus: false, // default: true - retry: false, // default: 3 - staleTime: 5 * 60 * 1000, // 5 minutes - }, - }, -}); +const queryClient = new QueryClient(); function App() { return ( +
@@ -38,6 +57,7 @@ function App() {
+ ); diff --git a/ell-studio/src/components/HierarchicalTable.js b/ell-studio/src/components/HierarchicalTable.js index 2be14b29..194d8acf 100644 --- a/ell-studio/src/components/HierarchicalTable.js +++ b/ell-studio/src/components/HierarchicalTable.js @@ -9,13 +9,24 @@ const TableRow = ({ item, schema, level = 0, onRowClick, columnWidths, updateWid const hasChildren = item.children && item.children.length > 0; const isExpanded = expandedRows[item.id]; const isSelected = isItemSelected(item); + const [isNew, setIsNew] = useState(true); const customRowClassName = rowClassName ? rowClassName(item) : ''; + useEffect(() => { + if (isNew) { + const timer = setTimeout(() => setIsNew(false), 200); + return () => clearTimeout(timer); + } + }, [isNew]); + return ( { if (onRowClick) onRowClick(item); }} diff --git a/ell-studio/src/components/invocations/InvocationsTable.js b/ell-studio/src/components/invocations/InvocationsTable.js index 239de446..eee89f84 100644 --- a/ell-studio/src/components/invocations/InvocationsTable.js +++ b/ell-studio/src/components/invocations/InvocationsTable.js @@ -11,7 +11,6 @@ const InvocationsTable = ({ invocations, currentPage, setCurrentPage, pageSize, const navigate = useNavigate(); - const onClickLMP = useCallback(({lmp, id : invocationId}) => { navigate(`/lmp/${lmp.name}/${lmp.lmp_id}?i=${invocationId}`); }, [navigate]); diff --git a/ell-studio/src/hooks/useBackend.js b/ell-studio/src/hooks/useBackend.js index 0553117a..a300396a 100644 --- a/ell-studio/src/hooks/useBackend.js +++ b/ell-studio/src/hooks/useBackend.js @@ -1,8 +1,45 @@ -import { useQuery, useQueries } from '@tanstack/react-query'; +import { useQuery, useQueryClient, useQueries } from '@tanstack/react-query'; import axios from 'axios'; - +import { useEffect, useState } from 'react'; const API_BASE_URL = "http://localhost:8080"; +const WS_URL = "ws://localhost:8080/ws"; + +export const useWebSocketConnection = () => { + const queryClient = useQueryClient(); + const [isConnected, setIsConnected] = useState(false); + useEffect(() => { + const socket = new WebSocket(WS_URL); + + socket.onopen = () => { + console.log('WebSocket connected'); + setIsConnected(true); + }; + + socket.onmessage = (event) => { + const data = JSON.parse(event.data); + if (data.entity === 'database_updated') { + // Invalidate relevant queries + queryClient.invalidateQueries({queryKey: ['traces']}); + queryClient.invalidateQueries({queryKey: ['latestLMPs']}); + queryClient.invalidateQueries({queryKey: ['invocations']}) ; + queryClient.invalidateQueries({queryKey: ['lmpDetails']}); + console.log('Database updated, invalidating queries'); + } + }; + + socket.onclose = () => { + console.log('WebSocket disconnected'); + setIsConnected(false); + }; + + return () => { + console.log('WebSocket connection closed'); + socket.close(); + }; + }, [queryClient]); + return { isConnected }; +}; export const useLMPs = (name, id) => { return useQuery({ @@ -21,7 +58,7 @@ export const useLMPs = (name, id) => { }); }; -export const useInvocations = (name, id, page = 0, pageSize = 50) => { +export const useInvocationsFromLMP = (name, id, page = 0, pageSize = 50) => { return useQuery({ queryKey: ['invocations', name, id, page, pageSize], queryFn: async () => { @@ -39,6 +76,18 @@ export const useInvocations = (name, id, page = 0, pageSize = 50) => { }); }; +export const useInvocation = (id) => { + return useQuery({ + queryKey: ['invocation', id], + queryFn: async () => { + const response = await axios.get(`${API_BASE_URL}/api/invocation/${id}`); + return response.data; + }, + enabled: !!id, + }); +} + + export const useMultipleLMPs = (usesIds) => { const multipleLMPs = useQueries({ queries: (usesIds || []).map(use => ({ @@ -55,26 +104,17 @@ export const useMultipleLMPs = (usesIds) => { return { isLoading, data }; }; - - - export const useLatestLMPs = (page = 0, pageSize = 100) => { return useQuery({ - queryKey: ['allLMPs', page, pageSize], + queryKey: ['latestLMPs', page, pageSize], queryFn: async () => { const skip = page * pageSize; const response = await axios.get(`${API_BASE_URL}/api/latest/lmps?skip=${skip}&limit=${pageSize}`); - const lmps = response.data; - - return lmps; + return response.data; } }); }; - - - - export const useTraces = (lmps) => { return useQuery({ queryKey: ['traces', lmps], @@ -103,4 +143,4 @@ export const useTraces = (lmps) => { }, enabled: !!lmps && lmps.length > 0, }); - }; + }; \ No newline at end of file diff --git a/ell-studio/src/pages/LMP.js b/ell-studio/src/pages/LMP.js index 01eb4821..6be317a5 100644 --- a/ell-studio/src/pages/LMP.js +++ b/ell-studio/src/pages/LMP.js @@ -1,6 +1,6 @@ import React, { useState, useEffect, useMemo } from "react"; import { useParams, useSearchParams, useNavigate, Link } from "react-router-dom"; -import { useLMPs, useInvocations, useMultipleLMPs } from "../hooks/useBackend"; +import { useLMPs, useInvocationsFromLMP, useMultipleLMPs, useInvocation } from "../hooks/useBackend"; import InvocationsTable from "../components/invocations/InvocationsTable"; import DependencyGraphPane from "../components/DependencyGraphPane"; import LMPSourceView from "../components/source/LMPSourceView"; @@ -35,6 +35,7 @@ function LMP() { const requestedInvocationId = searchParams.get("i"); const [currentPage, setCurrentPage] = useState(0); + const pageSize = 50; // TODO: Could be expensive if you have a funct on of versions. const { data: versionHistory, isLoading: isLoadingLMP } = useLMPs(name); @@ -47,7 +48,7 @@ function LMP() { } }, [versionHistory, id]); - const { data: invocations } = useInvocations(name, id); + const { data: invocations } = useInvocationsFromLMP(name, id, currentPage, pageSize); const { data: uses } = useMultipleLMPs(lmp?.uses); @@ -65,9 +66,14 @@ function LMP() { : null; }, [versionHistory, lmp]); - const requestedInvocation = useMemo(() => invocations?.find( - (invocation) => invocation.id === requestedInvocationId - ), [invocations, requestedInvocationId]); + const {data: requestedInvocationQueryData} = useInvocation(requestedInvocationId); + const requestedInvocation = useMemo(() => { + if (!requestedInvocationQueryData) + return invocations?.find(i => i.id === requestedInvocationId); + else + return requestedInvocationQueryData; + + }, [requestedInvocationQueryData, invocations, requestedInvocationId]); useEffect(() => { setSelectedTrace(requestedInvocation); @@ -233,6 +239,7 @@ function LMP() { { diff --git a/ell-studio/src/pages/Traces.js b/ell-studio/src/pages/Traces.js index f0570ca2..05ce620a 100644 --- a/ell-studio/src/pages/Traces.js +++ b/ell-studio/src/pages/Traces.js @@ -3,7 +3,7 @@ import { FiCopy, FiZap, FiEdit2, FiFilter, FiClock, FiColumns, FiPause, FiPlay } import InvocationsTable from '../components/invocations/InvocationsTable'; import InvocationsLayout from '../components/invocations/InvocationsLayout'; import { useNavigate, useLocation } from 'react-router-dom'; -import { useInvocations } from '../hooks/useBackend'; +import { useInvocationsFromLMP } from '../hooks/useBackend'; const Traces = () => { const [selectedTrace, setSelectedTrace] = useState(null); @@ -11,21 +11,13 @@ const Traces = () => { const navigate = useNavigate(); const location = useLocation(); + + // TODO Unify invocation search behaviour with the LMP page. const [currentPage, setCurrentPage] = useState(0); - const pageSize = 10; + const pageSize = 50; - const { data: invocations, refetch , isLoading } = useInvocations(null, null, currentPage, pageSize); + const { data: invocations , isLoading } = useInvocationsFromLMP(null, null, currentPage, pageSize); - useEffect(() => { - let intervalId; - if (isPolling) { - intervalId = setInterval(refetch, 200); // Poll every 200ms - } - - return () => { - if (intervalId) clearInterval(intervalId); - }; - }, [isPolling, refetch]); useEffect(() => { const searchParams = new URLSearchParams(location.search); diff --git a/examples/sqlite_example/ell.db-shm b/examples/sqlite_example/ell.db-shm new file mode 100644 index 0000000000000000000000000000000000000000..fe9ac2845eca6fe6da8a63cd096d9cf9e24ece10 GIT binary patch literal 32768 zcmeIuAr62r3