diff --git a/embedding_cluster/server/models.py b/embedding_cluster/server/models.py index 73946fe..4b35092 100644 --- a/embedding_cluster/server/models.py +++ b/embedding_cluster/server/models.py @@ -8,6 +8,8 @@ class CollectionInfo(BaseModel): name: str count: int + model_name: str | None = None + model_type: str | None = None class CollectionDetail(BaseModel): diff --git a/embedding_cluster/server/routes/collections.py b/embedding_cluster/server/routes/collections.py index 8b887bc..049ccd9 100644 --- a/embedding_cluster/server/routes/collections.py +++ b/embedding_cluster/server/routes/collections.py @@ -25,10 +25,19 @@ def _get_chromadb_client() -> ClientAPI: async def list_collections() -> list[CollectionInfo]: client = _get_chromadb_client() collection_names = client.list_collections() - return [ - CollectionInfo(name=c, count=client.get_collection(c).count()) - for c in collection_names - ] + results: list[CollectionInfo] = [] + for name in collection_names: + collection = client.get_collection(name) + metadata = collection.metadata or {} + results.append( + CollectionInfo( + name=name, + count=collection.count(), + model_name=metadata.get("model_name"), + model_type=metadata.get("model_type"), + ) + ) + return results @router.get("/{name}", response_model=CollectionDetail) diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index b2de06e..493badd 100644 --- a/frontend/src/App.tsx +++ b/frontend/src/App.tsx @@ -1,6 +1,6 @@ import { QueryClient, QueryClientProvider } from '@tanstack/react-query' import { BrowserRouter, Link, NavLink, Route, Routes } from 'react-router-dom' -import CollectionsPage from './pages/CollectionsPage' +import HomePage from './pages/HomePage' import IndexPage from './pages/IndexPage' import PlotPage from './pages/PlotPage' @@ -25,14 +25,14 @@ function NavBar() {
+ Home + + Index Plot - - Collections -
@@ -49,9 +49,9 @@ export default function App() {
- } /> + } /> + } /> } /> - } />
diff --git a/frontend/src/components/collections/CollectionList.tsx b/frontend/src/components/collections/CollectionList.tsx deleted file mode 100644 index d2a929b..0000000 --- a/frontend/src/components/collections/CollectionList.tsx +++ /dev/null @@ -1,79 +0,0 @@ -import { useNavigate } from 'react-router-dom' -import { useMutation, useQueryClient } from '@tanstack/react-query' -import type { CollectionInfo } from '../../types' -import { deleteCollection } from '../../api/collections' - -interface CollectionListProps { - collections: CollectionInfo[] -} - -export default function CollectionList({ collections }: CollectionListProps) { - const navigate = useNavigate() - const queryClient = useQueryClient() - - const deleteMutation = useMutation({ - mutationFn: deleteCollection, - onSuccess: () => { - queryClient.invalidateQueries({ queryKey: ['collections'] }) - }, - }) - - const handleDelete = (name: string) => { - if (window.confirm(`Are you sure you want to delete collection "${name}"?`)) { - deleteMutation.mutate(name) - } - } - - return ( -
- - - - - - - - - - {collections.map((collection, index) => ( - - - - - - ))} - -
- Name - - Item Count - - Actions -
- {collection.name} - - {collection.count.toLocaleString()} - - - | - -
-
- ) -} diff --git a/frontend/src/components/home/CollectionCard.tsx b/frontend/src/components/home/CollectionCard.tsx new file mode 100644 index 0000000..45581be --- /dev/null +++ b/frontend/src/components/home/CollectionCard.tsx @@ -0,0 +1,122 @@ +import { useState } from 'react' +import { useNavigate } from 'react-router-dom' +import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query' +import { fetchCollection, deleteCollection } from '../../api/collections' +import type { CollectionInfo } from '../../types' + +interface CollectionCardProps { + collection: CollectionInfo +} + +export default function CollectionCard({ collection }: CollectionCardProps) { + const [isExpanded, setIsExpanded] = useState(false) + const navigate = useNavigate() + const queryClient = useQueryClient() + + const { data: detail, isLoading: isDetailLoading } = useQuery({ + queryKey: ['collection', collection.name], + queryFn: () => fetchCollection(collection.name), + enabled: isExpanded, + }) + + const deleteMutation = useMutation({ + mutationFn: deleteCollection, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: ['collections'] }) + }, + }) + + const handleDelete = () => { + if ( + window.confirm( + `Are you sure you want to delete collection "${collection.name}"?`, + ) + ) { + deleteMutation.mutate(collection.name) + } + } + + return ( +
+
+

+ {collection.name} +

+

+ {collection.count.toLocaleString()} + items +

+ {(collection.model_type || collection.model_name) && ( +
+ {collection.model_type && ( + + {collection.model_type} + + )} + {collection.model_name && ( + + {collection.model_name} + + )} +
+ )} +
+ + {isExpanded && ( +
+ {isDetailLoading ? ( +
+
+ Loading details... +
+ ) : detail ? ( +
+

+ Metadata Fields +

+ {detail.metadata_fields.length > 0 ? ( +
+ {detail.metadata_fields.map((field) => ( + + {field} + + ))} +
+ ) : ( +

No metadata fields

+ )} +
+ ) : null} +
+ )} + +
+
+ + | + +
+ +
+
+ ) +} diff --git a/frontend/src/components/home/CollectionGrid.tsx b/frontend/src/components/home/CollectionGrid.tsx new file mode 100644 index 0000000..20ae508 --- /dev/null +++ b/frontend/src/components/home/CollectionGrid.tsx @@ -0,0 +1,23 @@ +import type { CollectionInfo } from '../../types' +import CollectionCard from './CollectionCard' + +interface CollectionGridProps { + collections: CollectionInfo[] +} + +export default function CollectionGrid({ + collections, +}: CollectionGridProps) { + return ( +
+

+ Collections +

+
+ {collections.map((collection) => ( + + ))} +
+
+ ) +} diff --git a/frontend/src/components/home/EmptyState.tsx b/frontend/src/components/home/EmptyState.tsx new file mode 100644 index 0000000..16147df --- /dev/null +++ b/frontend/src/components/home/EmptyState.tsx @@ -0,0 +1,32 @@ +import { Link } from 'react-router-dom' + +export default function EmptyState() { + return ( +
+ + + +

No collections yet

+

+ Get started by indexing your first CSV +

+ + Index New Data + +
+ ) +} diff --git a/frontend/src/components/home/StatsBar.tsx b/frontend/src/components/home/StatsBar.tsx new file mode 100644 index 0000000..9e0aebe --- /dev/null +++ b/frontend/src/components/home/StatsBar.tsx @@ -0,0 +1,26 @@ +interface StatsBarProps { + collectionCount: number + totalItems: number +} + +export default function StatsBar({ + collectionCount, + totalItems, +}: StatsBarProps) { + return ( +
+
+

Collections

+

+ {collectionCount.toLocaleString()} +

+
+
+

Total Items

+

+ {totalItems.toLocaleString()} +

+
+
+ ) +} diff --git a/frontend/src/pages/CollectionsPage.tsx b/frontend/src/pages/CollectionsPage.tsx deleted file mode 100644 index f5d5db8..0000000 --- a/frontend/src/pages/CollectionsPage.tsx +++ /dev/null @@ -1,44 +0,0 @@ -import { useQuery } from '@tanstack/react-query' -import { fetchCollections } from '../api/collections' -import CollectionList from '../components/collections/CollectionList' - -export default function CollectionsPage() { - const { data: collections, isLoading, isError } = useQuery({ - queryKey: ['collections'], - queryFn: fetchCollections, - }) - - return ( -
-
-

Collections

-

Manage your ChromaDB collections

-
- - {isLoading && ( -
-
- Loading collections... -
- )} - - {isError && ( -
- Error! - Failed to fetch collections. Please make sure the backend is running. -
- )} - - {!isLoading && !isError && collections && collections.length === 0 && ( -
-

No collections found

-

Index some data to get started

-
- )} - - {!isLoading && !isError && collections && collections.length > 0 && ( - - )} -
- ) -} diff --git a/frontend/src/pages/HomePage.tsx b/frontend/src/pages/HomePage.tsx new file mode 100644 index 0000000..0c5196c --- /dev/null +++ b/frontend/src/pages/HomePage.tsx @@ -0,0 +1,85 @@ +import { useQuery } from '@tanstack/react-query' +import { Link } from 'react-router-dom' +import { fetchCollections } from '../api/collections' +import StatsBar from '../components/home/StatsBar' +import CollectionGrid from '../components/home/CollectionGrid' +import EmptyState from '../components/home/EmptyState' + +export default function HomePage() { + const { data: collections, isLoading, isError } = useQuery({ + queryKey: ['collections'], + queryFn: fetchCollections, + }) + + const totalItems = collections?.reduce((sum, c) => sum + c.count, 0) ?? 0 + + return ( +
+
+

Dashboard

+

+ Manage your embedding collections +

+
+ + {isLoading && ( +
+
+ + Loading collections... + +
+ )} + + {isError && ( +
+ Error! + + {' '} + Failed to fetch collections. Please make sure the backend is + running. + +
+ )} + + {!isLoading && !isError && collections && ( + <> + + + + + + + Index New Data + + + {collections.length === 0 ? ( + + ) : ( + + )} + + )} +
+ ) +} diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index c075225..72a39e5 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -2,6 +2,8 @@ export interface CollectionInfo { name: string; count: number; + model_name: string | null; + model_type: string | null; } export interface CollectionDetail { diff --git a/tests/test_server_collections.py b/tests/test_server_collections.py index 0092f83..d5783c6 100644 --- a/tests/test_server_collections.py +++ b/tests/test_server_collections.py @@ -30,9 +30,11 @@ async def test_list_collections(client, mock_chromadb_client): """Test GET /api/collections returns list of collections with counts.""" mock_collection1 = MagicMock() mock_collection1.count.return_value = 10 + mock_collection1.metadata = None mock_collection2 = MagicMock() mock_collection2.count.return_value = 25 + mock_collection2.metadata = None mock_chromadb_client.list_collections.return_value = [ "collection1", @@ -52,8 +54,62 @@ async def test_list_collections(client, mock_chromadb_client): assert response.status_code == 200 data = response.json() assert len(data) == 2 - assert data[0] == {"name": "collection1", "count": 10} - assert data[1] == {"name": "collection2", "count": 25} + assert data[0] == { + "name": "collection1", + "count": 10, + "model_name": None, + "model_type": None, + } + assert data[1] == { + "name": "collection2", + "count": 25, + "model_name": None, + "model_type": None, + } + + +async def test_list_collections_includes_model_metadata(client, mock_chromadb_client): + mock_collection = MagicMock() + mock_collection.count.return_value = 10 + mock_collection.metadata = { + "model_name": "openai/clip-vit-base-patch32", + "model_type": "image", + } + + mock_chromadb_client.list_collections.return_value = ["test_col"] + mock_chromadb_client.get_collection.return_value = mock_collection + + with patch( + "embedding_cluster.server.routes.collections._get_chromadb_client", + return_value=mock_chromadb_client, + ): + response = await client.get("/api/collections") + + assert response.status_code == 200 + data = response.json() + assert len(data) == 1 + assert data[0]["model_name"] == "openai/clip-vit-base-patch32" + assert data[0]["model_type"] == "image" + + +async def test_list_collections_no_metadata(client, mock_chromadb_client): + mock_collection = MagicMock() + mock_collection.count.return_value = 5 + mock_collection.metadata = None + + mock_chromadb_client.list_collections.return_value = ["no_meta"] + mock_chromadb_client.get_collection.return_value = mock_collection + + with patch( + "embedding_cluster.server.routes.collections._get_chromadb_client", + return_value=mock_chromadb_client, + ): + response = await client.get("/api/collections") + + assert response.status_code == 200 + data = response.json() + assert data[0]["model_name"] is None + assert data[0]["model_type"] is None async def test_get_collection_success(client, mock_chromadb_client):