diff --git a/AGENTS.md b/AGENTS.md index 0104a8b..a5f0e1c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -78,6 +78,33 @@ Test configuration is in `pyproject.toml` under `[tool.pytest.ini_options]`: - `testpaths = ["tests"]` - `asyncio_mode = "auto"` (pytest-asyncio auto mode) +## E2E Testing + +```bash +# Install Playwright browsers (first-time) +cd frontend && npx playwright install chromium + +# Index sample data for E2E tests (first-time, from project root) +RUNNING_MODE=INDEX LOCAL_CSV_FILENAME=./embedding_cluster/csv/fashion_small.csv \ + ID_FIELD=id TEXT_EMBEDDING_FIELDS='["productDisplayName"]' \ + CHROMADB_COLLECTION_PREFIX=fashion_ uv run python -m embedding_cluster + +# Build frontend (required before E2E) +cd frontend && npm run build + +# Run E2E tests +cd frontend && npm run test:e2e + +# Run E2E tests with UI +cd frontend && npm run test:e2e:ui + +# Run single test file +cd frontend && npx playwright test e2e/search.spec.ts +``` + +E2E tests require pre-indexed ChromaDB data. The `webServer` config in +`playwright.config.ts` auto-starts the FastAPI backend. Tests run against +`http://localhost:8000`. ## CI GitHub Actions workflow in `.github/workflows/ci.yml` runs on push/PR: diff --git a/README.md b/README.md index d3ba508..c35ac2c 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,8 @@ FastAPI backend and React frontend. * Collection management (list, inspect, delete) through UI and API. * Async batch indexing with configurable parallelism and retry logic. * Hardware acceleration support (CPU, MPS, CUDA). +* Semantic search within clusters -- find similar items by text + query or image URL, with results highlighted in the 3D view. ### Architecture @@ -208,6 +210,8 @@ The web UI provides: and display fields, then visualize in an interactive 3D scatter plot with hover tooltips showing metadata and images. Switch between colored particles, image sprites, and instanced spheres render modes. + Use semantic search to find similar items by text or image URL, + with matching points highlighted in the 3D view. * **Collections page** -- List, inspect, and delete ChromaDB collections. @@ -248,6 +252,60 @@ uv run pytest --cov=embedding_cluster --cov-report=term-missing --cov-fail-under uv run pre-commit run --all-files ``` +### E2E Testing + +End-to-end tests use [Playwright](https://playwright.dev/) and run +against the full stack (FastAPI backend + React frontend). + +#### First-Time Setup + +1. Install Playwright browsers: + + ```bash + cd frontend + npm install + npx playwright install chromium + ``` + +2. Index sample data for tests (one-time, from project root): + + ```bash + RUNNING_MODE=INDEX \ + LOCAL_CSV_FILENAME=./embedding_cluster/csv/fashion_small.csv \ + ID_FIELD=id \ + TEXT_EMBEDDING_FIELDS='["productDisplayName"]' \ + CHROMADB_COLLECTION_PREFIX=fashion_ \ + uv run python -m embedding_cluster + ``` + +3. Build the frontend: + + ```bash + cd frontend + npm run build + ``` + +#### Running E2E Tests + +```bash +cd frontend + +# Run all E2E tests (headless, auto-starts backend) +npm run test:e2e + +# Run with interactive UI for debugging +npm run test:e2e:ui + +# Run a specific test file +npx playwright test e2e/search.spec.ts + +# Show HTML report after a run +npx playwright show-report +``` + +The Playwright config auto-starts the FastAPI server. If you +already have the server running (`RUNNING_MODE=SERVER`), it reuses +the existing server instead. ### Project Structure ```text diff --git a/docs/plans/2026-02-24-e2e-playwright-tests.md b/docs/plans/2026-02-24-e2e-playwright-tests.md new file mode 100644 index 0000000..5dce876 --- /dev/null +++ b/docs/plans/2026-02-24-e2e-playwright-tests.md @@ -0,0 +1,504 @@ +# E2E Playwright Tests Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Add Playwright E2E testing infrastructure to the project with semantic search test coverage, so any developer who clones the repo can run E2E tests. + +**Architecture:** Playwright runs against the full stack -- FastAPI backend serving the React SPA at `localhost:8000`. The `webServer` config in `playwright.config.ts` auto-starts the backend. Tests require pre-indexed ChromaDB data (the sample `fashion_small.csv`). The test suite covers the semantic search feature on the Plot page: search bar visibility, text search, result interaction, and clear/reset flows. + +**Tech Stack:** Playwright Test, TypeScript, FastAPI, ChromaDB + +**Prerequisites:** The developer must have indexed the sample data before running E2E tests. The README documents this requirement with exact commands. + +--- + +## Task 1: Install Playwright and create configuration + +**Files:** +- Modify: `frontend/package.json` (via npm install) +- Create: `frontend/playwright.config.ts` +- Create: `frontend/.gitignore` addition for Playwright artifacts + +**Step 1: Install Playwright** + +Run from `frontend/`: + +```bash +npm install -D @playwright/test +npx playwright install chromium +``` + +This adds `@playwright/test` to devDependencies and installs the Chromium browser binary. + +**Step 2: Create Playwright configuration** + +Create `frontend/playwright.config.ts`: + +```typescript +import { defineConfig, devices } from '@playwright/test' + +export default defineConfig({ + testDir: './e2e', + fullyParallel: true, + forbidOnly: !!process.env.CI, + retries: process.env.CI ? 2 : 0, + workers: process.env.CI ? 1 : undefined, + reporter: 'html', + timeout: 60_000, + use: { + baseURL: 'http://localhost:8000', + trace: 'on-first-retry', + screenshot: 'only-on-failure', + }, + projects: [ + { + name: 'chromium', + use: { ...devices['Desktop Chrome'] }, + }, + ], + webServer: { + command: 'cd .. && RUNNING_MODE=SERVER uv run python -m embedding_cluster', + url: 'http://localhost:8000/api/health', + reuseExistingServer: !process.env.CI, + timeout: 30_000, + }, +}) +``` + +Key decisions: +- `testDir: './e2e'` -- keeps E2E tests separate from future unit tests +- `baseURL: 'http://localhost:8000'` -- the FastAPI server serves the built SPA +- `webServer` -- auto-starts the backend; `reuseExistingServer: true` locally so devs can keep the server running +- `timeout: 60_000` -- generous timeout since computing clusters involves ML model loading +- Only Chromium for now -- lightweight, add Firefox/WebKit later if needed + +**Step 3: Add npm scripts** + +Add to `frontend/package.json` scripts: + +```json +"test:e2e": "playwright test", +"test:e2e:ui": "playwright test --ui" +``` + +**Step 4: Add Playwright artifacts to .gitignore** + +Append to `frontend/.gitignore`: + +``` +# Playwright +/test-results/ +/playwright-report/ +/blob-report/ +/playwright/.cache/ +/e2e/.auth/ +``` + +**Step 5: Verify Playwright installs and config loads** + +Run from `frontend/`: + +```bash +npx playwright test --list +``` + +Expected: `no tests found` (no test files yet), but no config errors. + +--- + +## Task 2: Create test fixtures and helpers + +**Files:** +- Create: `frontend/e2e/fixtures.ts` + +**Step 1: Create shared fixtures** + +Create `frontend/e2e/fixtures.ts`: + +```typescript +import { test as base, expect } from '@playwright/test' + +/** + * Extended test fixtures for embedding-clusters E2E tests. + * + * Provides: + * - plotPage: Navigates to the Plot page and waits for collections to load + */ +export const test = base.extend<{ + plotPage: void +}>({ + plotPage: async ({ page }, use) => { + await page.goto('/plot') + // Wait for the collection dropdown to be populated + await expect( + page.getByRole('combobox').first() + ).not.toHaveValue('', { timeout: 10_000 }) + await use() + }, +}) + +export { expect } +``` + +This fixture navigates to `/plot` and waits for the collections API to respond (the dropdown gets populated). All search tests need this baseline. + +--- + +## Task 3: Create semantic search E2E tests + +**Files:** +- Create: `frontend/e2e/search.spec.ts` + +**Step 1: Create the test file** + +Create `frontend/e2e/search.spec.ts`: + +```typescript +import { test, expect } from './fixtures' + +// These tests require pre-indexed data in ChromaDB. +// Run the indexing command from the README first: +// RUNNING_MODE=INDEX LOCAL_CSV_FILENAME=./embedding_cluster/csv/fashion_small.csv \ +// ID_FIELD=id TEXT_EMBEDDING_FIELDS='["productDisplayName"]' \ +// CHROMADB_COLLECTION_PREFIX=fashion_ uv run python -m embedding_cluster + +const COLLECTION_NAME = 'fashion_productDisplayName' + +test.describe('Semantic Search', () => { + test.beforeEach(async ({ page, plotPage: _ }) => { + // Select the collection from the dropdown + await page.getByRole('combobox').first().selectOption(COLLECTION_NAME) + + // Wait for collection details to load (Compute button appears) + await expect( + page.getByRole('button', { name: 'Compute Plot' }) + ).toBeVisible({ timeout: 10_000 }) + + // Click Compute and wait for the plot to render + await page.getByRole('button', { name: 'Compute Plot' }).click() + + // Wait for computing to finish -- "Computing Clusters..." disappears + await expect( + page.getByText('Computing Clusters...') + ).toBeHidden({ timeout: 120_000 }) + + // Verify plot rendered -- the canvas should be present + await expect(page.locator('canvas')).toBeVisible({ timeout: 10_000 }) + }) + + test('search bar appears after computing plot', async ({ page }) => { + // The "Semantic Search" heading should be visible in the sidebar + await expect( + page.getByRole('heading', { name: 'Semantic Search' }) + ).toBeVisible() + + // Text radio should be selected by default + const textRadio = page.getByRole('radio', { name: 'Text' }) + await expect(textRadio).toBeChecked() + + // Search input should be present + await expect( + page.getByPlaceholder('Search by text...') + ).toBeVisible() + + // Search button should be visible but disabled (no query) + const searchButton = page.getByRole('button', { name: 'Search' }) + await expect(searchButton).toBeVisible() + await expect(searchButton).toBeDisabled() + }) + + test('text search returns results', async ({ page }) => { + // Type a search query + await page.getByPlaceholder('Search by text...').fill('blue shirt') + + // Search button should be enabled now + const searchButton = page.getByRole('button', { name: 'Search' }) + await expect(searchButton).toBeEnabled() + + // Click search + await searchButton.click() + + // Wait for results to appear -- "Results" heading with count + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // "Highlight All" button should be visible + await expect( + page.getByRole('button', { name: 'Highlight All' }) + ).toBeVisible() + + // Result items should be present (at least one) + const resultItems = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }) + await expect(resultItems.first()).toBeVisible() + }) + + test('search via Enter key', async ({ page }) => { + const searchInput = page.getByPlaceholder('Search by text...') + await searchInput.fill('casual shoes') + await searchInput.press('Enter') + + // Results should appear + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + }) + + test('clear search removes results', async ({ page }) => { + // Perform a search first + await page.getByPlaceholder('Search by text...').fill('jacket') + await page.getByRole('button', { name: 'Search' }).click() + + // Wait for results + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // Click Clear + await page.getByRole('button', { name: 'Clear' }).click() + + // Results should disappear + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeHidden() + + // Search input should be empty + await expect( + page.getByPlaceholder('Search by text...') + ).toHaveValue('') + }) + + test('clicking a result highlights it', async ({ page }) => { + // Search + await page.getByPlaceholder('Search by text...').fill('men') + await page.getByRole('button', { name: 'Search' }).click() + + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // Click the first result -- it should get the active style (border-l-2) + const firstResult = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }).first() + await firstResult.click() + + // The clicked result should have the active indicator (blue left border) + await expect(firstResult).toHaveClass(/border-blue-500/) + }) + + test('highlight all button activates all results', async ({ page }) => { + // Search + await page.getByPlaceholder('Search by text...').fill('shirt') + await page.getByRole('button', { name: 'Search' }).click() + + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // Click a single result first to narrow highlight + const firstResult = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }).first() + await firstResult.click() + + // Now click "Highlight All" + await page.getByRole('button', { name: 'Highlight All' }).click() + + // All result buttons should have the active class + const allResults = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }) + const count = await allResults.count() + expect(count).toBeGreaterThan(1) + }) + + test('switch to image URL mode', async ({ page }) => { + // Click image radio + await page.getByText('Image URL').click() + + // Placeholder should change + await expect( + page.getByPlaceholder('Paste image URL...') + ).toBeVisible() + + // Text placeholder should be gone + await expect( + page.getByPlaceholder('Search by text...') + ).toBeHidden() + }) + + test('adjusting results slider changes value', async ({ page }) => { + // The slider label should show the default value + await expect( + page.getByText('Results: 10') + ).toBeVisible() + + // Adjust the slider + const slider = page.locator('input[type="range"]').last() + await slider.fill('25') + + // Label should update + await expect( + page.getByText('Results: 25') + ).toBeVisible() + }) +}) +``` + +**Step 2: Run the tests** + +First, ensure data is indexed (one-time setup): + +```bash +# From project root +RUNNING_MODE=INDEX \ + LOCAL_CSV_FILENAME=./embedding_cluster/csv/fashion_small.csv \ + ID_FIELD=id \ + TEXT_EMBEDDING_FIELDS='["productDisplayName"]' \ + CHROMADB_COLLECTION_PREFIX=fashion_ \ + uv run python -m embedding_cluster +``` + +Then build the frontend and run tests: + +```bash +cd frontend +npm run build +npx playwright test +``` + +Expected: All tests pass. + +--- + +## Task 4: Update documentation + +**Files:** +- Modify: `README.md` +- Modify: `AGENTS.md` + +**Step 1: Add E2E testing section to README.md** + +Add to the Development section, after the existing Commands subsection: + +```markdown +### E2E Testing + +End-to-end tests use [Playwright](https://playwright.dev/) and run +against the full stack (FastAPI backend + React frontend). + +#### First-Time Setup + +1. Install Playwright browsers: + + ```bash + cd frontend + npm install + npx playwright install chromium + ``` + +2. Index sample data for tests (one-time, from project root): + + ```bash + RUNNING_MODE=INDEX \ + LOCAL_CSV_FILENAME=./embedding_cluster/csv/fashion_small.csv \ + ID_FIELD=id \ + TEXT_EMBEDDING_FIELDS='["productDisplayName"]' \ + CHROMADB_COLLECTION_PREFIX=fashion_ \ + uv run python -m embedding_cluster + ``` + +3. Build the frontend: + + ```bash + cd frontend + npm run build + ``` + +#### Running E2E Tests + +```bash +cd frontend + +# Run all E2E tests (headless, auto-starts backend) +npm run test:e2e + +# Run with interactive UI for debugging +npm run test:e2e:ui + +# Run a specific test file +npx playwright test e2e/search.spec.ts + +# Show HTML report after a run +npx playwright show-report +``` + +The Playwright config auto-starts the FastAPI server. If you +already have the server running (`RUNNING_MODE=SERVER`), it reuses +the existing server instead. +``` + +**Step 2: Add E2E commands to AGENTS.md** + +Add to the Testing section in `AGENTS.md`: + +```markdown +## E2E Testing + +```bash +# Install Playwright browsers (first-time) +cd frontend && npx playwright install chromium + +# Index sample data for E2E tests (first-time, from project root) +RUNNING_MODE=INDEX LOCAL_CSV_FILENAME=./embedding_cluster/csv/fashion_small.csv \ + ID_FIELD=id TEXT_EMBEDDING_FIELDS='["productDisplayName"]' \ + CHROMADB_COLLECTION_PREFIX=fashion_ uv run python -m embedding_cluster + +# Build frontend (required before E2E) +cd frontend && npm run build + +# Run E2E tests +cd frontend && npm run test:e2e + +# Run E2E tests with UI +cd frontend && npm run test:e2e:ui + +# Run single test file +cd frontend && npx playwright test e2e/search.spec.ts +``` + +E2E tests require pre-indexed ChromaDB data. The `webServer` config in +`playwright.config.ts` auto-starts the FastAPI backend. Tests run against +`http://localhost:8000`. +``` + +--- + +## Task 5: Run tests and verify + +**Step 1: Build frontend** + +```bash +cd frontend && npm run build +``` + +**Step 2: Run E2E tests** + +```bash +cd frontend && npx playwright test +``` + +Expected: All 8 tests pass. + +**Step 3: Fix any failures** + +If tests fail, debug with: + +```bash +cd frontend && npx playwright test --ui +``` + +--- diff --git a/docs/plans/2026-02-24-semantic-search.md b/docs/plans/2026-02-24-semantic-search.md new file mode 100644 index 0000000..b6c91f0 --- /dev/null +++ b/docs/plans/2026-02-24-semantic-search.md @@ -0,0 +1,913 @@ +# Semantic Search Within Clusters Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Add a search bar to the Plot page that accepts text or image URL input, generates an embedding, queries ChromaDB for nearest neighbors, and highlights matching points in the 3D scatter plot. + +**Architecture:** New `POST /api/search` endpoint generates query embeddings using the same model that indexed the collection (CLIP for images, SentenceTransformer for text), then calls ChromaDB's `.query()`. Frontend adds a search bar + results panel to PlotPage, with highlighted points in all 3D render modes. Model type is passed explicitly in the search request since collections don't store model metadata. + +**Tech Stack:** Python 3.13 / FastAPI / ChromaDB / CLIP / SentenceTransformer / React 19 / TypeScript / Zustand / Tailwind CSS / Three.js + +--- + +## Task 1: Backend - Pydantic models for search request/response + +**Files:** +- Modify: `embedding_cluster/server/models.py` + +**Step 1: Add SearchRequest and SearchResponse models** + +Add to the bottom of `embedding_cluster/server/models.py`: + +```python +class SearchResult(BaseModel): + id: str + distance: float + metadata: dict[str, Any] + + +class SearchRequest(BaseModel): + collection_name: str + query_text: str | None = None + query_image_url: str | None = None + n_results: int = 10 + model_type: str = "text" + image_model_name: str = "openai/clip-vit-base-patch32" + text_model_name: str = "BAAI/bge-small-en-v1.5" + + +class SearchResponse(BaseModel): + results: list[SearchResult] +``` + +**Step 2: Run type check** + +Run: `uv run mypy embedding_cluster/server/models.py` +Expected: PASS + +--- + +## Task 2: Backend - Search route with embedding generation + +**Files:** +- Create: `embedding_cluster/server/routes/search.py` +- Modify: `embedding_cluster/server/app.py` (register router) + +**Step 1: Create the search route** + +Create `embedding_cluster/server/routes/search.py`: + +```python +from __future__ import annotations + +import logging +from typing import Any + +import chromadb +import torch +from fastapi import APIRouter, HTTPException +from sentence_transformers import SentenceTransformer +from transformers import CLIPModel, CLIPProcessor + +from embedding_cluster.server.models import SearchRequest, SearchResponse, SearchResult +from embedding_cluster.utils import ImageDownloader + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/api/search", tags=["search"]) + +# Lazy-loaded model cache +_model_cache: dict[str, Any] = {} + + +def _get_chromadb_client() -> chromadb.ClientAPI: + return chromadb.PersistentClient(path="./chromadb") + + +def _get_text_model(model_name: str) -> SentenceTransformer: + cache_key = f"text:{model_name}" + if cache_key not in _model_cache: + logger.info("Loading text model: %s", model_name) + _model_cache[cache_key] = SentenceTransformer(model_name) + return _model_cache[cache_key] + + +def _get_image_model( + model_name: str, +) -> tuple[CLIPModel, CLIPProcessor]: + cache_key = f"image:{model_name}" + if cache_key not in _model_cache: + logger.info("Loading image model: %s", model_name) + _model_cache[cache_key] = ( + CLIPModel.from_pretrained(model_name), + CLIPProcessor.from_pretrained(model_name), + ) + return _model_cache[cache_key] + + +async def _generate_text_embedding( + query_text: str, model_name: str +) -> list[float]: + model = _get_text_model(model_name) + embedding = model.encode(query_text, show_progress_bar=False) + return embedding.tolist() + + +async def _generate_image_embedding( + image_url: str, model_name: str +) -> list[float]: + image = await ImageDownloader().download_image_exp_backoff(image_url) + if image is None: + msg = f"Failed to download image from {image_url}" + raise ValueError(msg) + + model, processor = _get_image_model(model_name) + inputs = processor( + text=None, images=image, return_tensors="pt", padding=True + ) + with torch.no_grad(): + img_features = model.get_image_features(inputs["pixel_values"]) + return img_features.squeeze(0).cpu().numpy().tolist() + + +@router.post("", response_model=SearchResponse) +async def search_collection(request: SearchRequest) -> SearchResponse: + if not request.query_text and not request.query_image_url: + raise HTTPException( + status_code=400, + detail="Either query_text or query_image_url is required", + ) + + client = _get_chromadb_client() + try: + collection = client.get_collection(request.collection_name) + except Exception as e: + raise HTTPException( + status_code=404, + detail=f"Collection not found: {request.collection_name}", + ) from e + + if collection.count() == 0: + return SearchResponse(results=[]) + + # Generate embedding + if request.query_text: + embedding = await _generate_text_embedding( + request.query_text, request.text_model_name + ) + else: + assert request.query_image_url is not None + embedding = await _generate_image_embedding( + request.query_image_url, request.image_model_name + ) + + # Query ChromaDB + query_result = collection.query( + query_embeddings=[embedding], + n_results=min(request.n_results, collection.count()), + ) + + results: list[SearchResult] = [] + if query_result["ids"] and query_result["distances"]: + ids = query_result["ids"][0] + distances = query_result["distances"][0] + metadatas = ( + query_result["metadatas"][0] + if query_result["metadatas"] + else [{}] * len(ids) + ) + for i, doc_id in enumerate(ids): + results.append( + SearchResult( + id=doc_id, + distance=distances[i], + metadata=metadatas[i] if metadatas[i] else {}, + ) + ) + + return SearchResponse(results=results) +``` + +**Step 2: Register the search router in app.py** + +Add import and `include_router` in `embedding_cluster/server/app.py`: + +```python +from embedding_cluster.server.routes.search import router as search_router +# ... +app.include_router(search_router) +``` + +**Step 3: Run type check and lint** + +Run: `uv run mypy embedding_cluster/server/routes/search.py` +Run: `uv run ruff check embedding_cluster/server/routes/search.py` +Expected: PASS (may need minor fixes) + +--- + +## Task 3: Backend tests for search endpoint + +**Files:** +- Create: `tests/test_server_search.py` + +**Step 1: Write comprehensive tests** + +Create `tests/test_server_search.py` following the pattern from `test_server_collections.py` and `test_server_plot.py`: + +```python +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import numpy as np +import pytest +from httpx import ASGITransport, AsyncClient + +from embedding_cluster.server.app import create_app + + +@pytest.fixture +def app(): + return create_app() + + +@pytest.fixture +async def client(app): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + yield ac + + +@pytest.fixture +def mock_chromadb_client(): + mock_client = MagicMock() + return mock_client + + +@pytest.fixture +def mock_collection(mock_chromadb_client): + mock_coll = MagicMock() + mock_coll.count.return_value = 100 + mock_coll.query.return_value = { + "ids": [["id1", "id2", "id3"]], + "distances": [[0.1, 0.3, 0.5]], + "metadatas": [ + [ + {"name": "item1", "imageUrl": "http://example.com/1.jpg"}, + {"name": "item2", "imageUrl": "http://example.com/2.jpg"}, + {"name": "item3", "imageUrl": "http://example.com/3.jpg"}, + ] + ], + } + mock_chromadb_client.get_collection.return_value = mock_coll + return mock_coll + + +@pytest.fixture +def mock_text_model(): + mock_model = MagicMock() + mock_model.encode.return_value = np.zeros(384) + return mock_model + + +@pytest.fixture +def mock_image_model(): + mock_clip = MagicMock() + mock_processor = MagicMock() + + import torch + mock_features = torch.zeros(1, 512) + mock_clip.get_image_features.return_value = mock_features + + return mock_clip, mock_processor + + +async def test_search_text_query( + client, mock_chromadb_client, mock_collection, mock_text_model +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_text_model, + ), + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_text": "red shoes", + "n_results": 3, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["results"]) == 3 + assert data["results"][0]["id"] == "id1" + assert data["results"][0]["distance"] == 0.1 + assert "metadata" in data["results"][0] + + +async def test_search_image_query( + client, mock_chromadb_client, mock_collection, mock_image_model +): + mock_clip, mock_processor = mock_image_model + mock_image = MagicMock() + + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_image_model", + return_value=(mock_clip, mock_processor), + ), + patch( + "embedding_cluster.server.routes.search.ImageDownloader" + ) as mock_downloader_cls, + ): + mock_instance = MagicMock() + mock_instance.download_image_exp_backoff = AsyncMock( + return_value=mock_image + ) + mock_downloader_cls.return_value = mock_instance + + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_image_url": "http://example.com/query.jpg", + "model_type": "image", + "n_results": 3, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["results"]) == 3 + + +async def test_search_collection_not_found(client): + mock_client = MagicMock() + mock_client.get_collection.side_effect = Exception("Not found") + + with patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_client, + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "nonexistent", + "query_text": "test", + }, + ) + + assert response.status_code == 404 + assert "Collection not found" in response.json()["detail"] + + +async def test_search_empty_query(client): + response = await client.post( + "/api/search", + json={"collection_name": "test_collection"}, + ) + assert response.status_code == 400 + + +async def test_search_custom_n_results( + client, mock_chromadb_client, mock_collection, mock_text_model +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_text_model, + ), + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_text": "blue jacket", + "n_results": 5, + }, + ) + + assert response.status_code == 200 + mock_collection.query.assert_called_once() + call_kwargs = mock_collection.query.call_args + assert call_kwargs.kwargs.get("n_results") == 5 or call_kwargs[1].get("n_results") == 5 + + +async def test_search_response_structure( + client, mock_chromadb_client, mock_collection, mock_text_model +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_text_model, + ), + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_text": "test", + }, + ) + + data = response.json() + assert "results" in data + for result in data["results"]: + assert "id" in result + assert "distance" in result + assert "metadata" in result + assert isinstance(result["distance"], float) + + +async def test_search_empty_collection(client, mock_chromadb_client): + mock_coll = MagicMock() + mock_coll.count.return_value = 0 + mock_chromadb_client.get_collection.return_value = mock_coll + + with patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "empty_collection", + "query_text": "test", + }, + ) + + assert response.status_code == 200 + assert response.json()["results"] == [] + + +async def test_search_image_download_failure( + client, mock_chromadb_client, mock_collection +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search.ImageDownloader" + ) as mock_downloader_cls, + ): + mock_instance = MagicMock() + mock_instance.download_image_exp_backoff = AsyncMock( + return_value=None + ) + mock_downloader_cls.return_value = mock_instance + + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_image_url": "http://example.com/bad.jpg", + "model_type": "image", + }, + ) + + assert response.status_code == 500 +``` + +**Step 2: Run tests** + +Run: `uv run pytest tests/test_server_search.py -v` +Expected: All PASS + +--- + +## Task 4: Frontend - Types and API client for search + +**Files:** +- Modify: `frontend/src/types/index.ts` +- Modify: `frontend/src/api/plot.ts` + +**Step 1: Add search types** + +Add to end of `frontend/src/types/index.ts`: + +```typescript +// Search +export interface SearchResult { + id: string; + distance: number; + metadata: Record; +} + +export interface SearchRequest { + collection_name: string; + query_text?: string; + query_image_url?: string; + n_results?: number; + model_type?: string; + image_model_name?: string; + text_model_name?: string; +} + +export interface SearchResponse { + results: SearchResult[]; +} +``` + +**Step 2: Add search API function** + +Add to `frontend/src/api/plot.ts`: + +```typescript +import type { SearchRequest, SearchResponse } from "../types"; + +export async function searchCollection( + request: SearchRequest, +): Promise { + return apiPost("/search", request); +} +``` + +--- + +## Task 5: Frontend - Search state in plotStore + +**Files:** +- Modify: `frontend/src/stores/plotStore.ts` + +**Step 1: Add search state and actions** + +Add to the `PlotState` interface and store: + +```typescript +// Add to interface: +searchResults: SearchResult[] | null +highlightedIds: Set +isSearching: boolean + +// Add actions: +setSearchResults: (results: SearchResult[] | null) => void +setHighlightedIds: (ids: Set) => void +setIsSearching: (searching: boolean) => void +clearSearch: () => void + +// Add to store implementation: +searchResults: null, +highlightedIds: new Set(), +isSearching: false, + +setSearchResults: (results) => set({ + searchResults: results, + highlightedIds: new Set(results?.map(r => r.id) ?? []), +}), +setHighlightedIds: (ids) => set({ highlightedIds: ids }), +setIsSearching: (searching) => set({ isSearching: searching }), +clearSearch: () => set({ + searchResults: null, + highlightedIds: new Set(), + isSearching: false, +}), +``` + +--- + +## Task 6: Frontend - SearchBar component + +**Files:** +- Create: `frontend/src/components/plot/SearchBar.tsx` + +**Step 1: Create SearchBar component** + +```tsx +import { useState } from 'react' +import { useMutation } from '@tanstack/react-query' +import { searchCollection } from '../../api/plot' +import { usePlotStore } from '../../stores/plotStore' +import type { SearchRequest } from '../../types' + +interface SearchBarProps { + collectionName: string +} + +export default function SearchBar({ collectionName }: SearchBarProps) { + const [queryType, setQueryType] = useState<'text' | 'image'>('text') + const [queryValue, setQueryValue] = useState('') + const [nResults, setNResults] = useState(10) + const { setSearchResults, setIsSearching, clearSearch } = usePlotStore() + + const mutation = useMutation({ + mutationFn: (request: SearchRequest) => searchCollection(request), + onMutate: () => setIsSearching(true), + onSuccess: (data) => { + setSearchResults(data.results) + setIsSearching(false) + }, + onError: () => setIsSearching(false), + }) + + const handleSearch = () => { + if (!queryValue.trim()) return + const request: SearchRequest = { + collection_name: collectionName, + n_results: nResults, + ...(queryType === 'text' + ? { query_text: queryValue } + : { query_image_url: queryValue, model_type: 'image' }), + } + mutation.mutate(request) + } + + const handleClear = () => { + setQueryValue('') + clearSearch() + } + + return ( +
+

Semantic Search

+ +
+ {(['text', 'image'] as const).map((type) => ( + + ))} +
+ + setQueryValue(e.target.value)} + onKeyDown={(e) => e.key === 'Enter' && handleSearch()} + placeholder={queryType === 'text' ? 'Search by text...' : 'Paste image URL...'} + className="w-full border border-gray-300 rounded-md px-3 py-2 text-sm" + /> + +
+ + setNResults(Number(e.target.value))} + className="w-full" + /> +
+ +
+ + +
+ + {mutation.isError && ( +

+ Search failed: {(mutation.error as Error).message} +

+ )} +
+ ) +} +``` + +--- + +## Task 7: Frontend - SearchResults component + +**Files:** +- Create: `frontend/src/components/plot/SearchResults.tsx` + +**Step 1: Create SearchResults panel** + +```tsx +import { usePlotStore } from '../../stores/plotStore' + +export default function SearchResults() { + const searchResults = usePlotStore((state) => state.searchResults) + const highlightedIds = usePlotStore((state) => state.highlightedIds) + const setHighlightedIds = usePlotStore((state) => state.setHighlightedIds) + + if (!searchResults || searchResults.length === 0) return null + + const handleClickResult = (id: string) => { + setHighlightedIds(new Set([id])) + } + + const handleShowAll = () => { + setHighlightedIds(new Set(searchResults.map((r) => r.id))) + } + + const getImageUrl = (metadata: Record): string | null => { + for (const value of Object.values(metadata)) { + if (typeof value === 'string' && (value.startsWith('http') || value.startsWith('/'))) { + return value + } + } + return null + } + + return ( +
+
+

+ Results ({searchResults.length}) +

+ +
+
+ {searchResults.map((result) => { + const imageUrl = getImageUrl(result.metadata) + const isActive = highlightedIds.has(result.id) + return ( + + ) + })} +
+
+ ) +} +``` + +--- + +## Task 8: Frontend - Integrate SearchBar and SearchResults into PlotPage + +**Files:** +- Modify: `frontend/src/pages/PlotPage.tsx` + +**Step 1: Add search components to the left sidebar** + +Import and render `SearchBar` below `PlotControls` (inside the sidebar div), and `SearchResults` below it. Pass the selected collection name to SearchBar. + +The PlotPage needs access to the selected collection name. We can get it from the plotStore or from the URL search params. Since `PlotControls` already uses `searchParams`, extract collection from there. + +Add the SearchBar after `PlotControls` and SearchResults after that, inside the sidebar div. SearchBar should only show when plotData is available (meaning a collection has been computed). + +--- + +## Task 9: Frontend - Highlight matching points in 3D renderers + +**Files:** +- Modify: `frontend/src/components/plot/ParticleCloud.tsx` +- Modify: `frontend/src/components/plot/InstancedSpheres.tsx` +- Modify: `frontend/src/components/plot/ImageSpriteCloud.tsx` + +**Step 1: ParticleCloud - highlight logic** + +Read `highlightedIds` from the store. When `highlightedIds.size > 0`: +- Matched points: full opacity, normal size +- Non-matched points: reduced opacity (0.15), normal size + +Use a custom ShaderMaterial or set alpha in the color buffer. The simplest approach: add an `alpha` buffer attribute and use `transparent={true}` with a custom opacity per point. Since `pointsMaterial` doesn't support per-point opacity natively, use the color channel trick: dim non-highlighted colors by multiplying RGB by 0.15. + +```typescript +// In the useMemo, after setting colors: +const hasHighlights = highlightedIds.size > 0 + +// Modify color loop: +const dimFactor = hasHighlights && !highlightedIds.has(p.id) ? 0.15 : 1.0 +cols[i * 3] = color.r * dimFactor +cols[i * 3 + 1] = color.g * dimFactor +cols[i * 3 + 2] = color.b * dimFactor +``` + +Add `highlightedIds` to the useMemo dependency array. + +**Step 2: InstancedSpheres - highlight logic** + +Same approach: dim non-highlighted instance colors. + +```typescript +// In useEffect, modify the color setting: +const dimFactor = highlightedIds.size > 0 && !highlightedIds.has(p.id) ? 0.15 : 1.0 +const c = colorObjects[p.cluster % colorObjects.length].clone().multiplyScalar(dimFactor) +mesh.setColorAt(i, c) +``` + +Add `highlightedIds` store subscription. + +**Step 3: ImageSpriteCloud - highlight logic** + +Add opacity prop based on highlight state: + +```typescript +// In spritesToRender mapping: +const isHighlighted = highlightedIds.size === 0 || highlightedIds.has(point.id) +const opacity = isHighlighted ? 1.0 : 0.15 +// Pass opacity to PointSprite and apply to spriteMaterial +``` + +--- + +## Task 10: Backend - Run full test suite and lint + +**Step 1: Run all backend checks** + +Run: `uv run ruff check embedding_cluster/ tests/` +Run: `uv run ruff format --check embedding_cluster/ tests/` +Run: `uv run mypy embedding_cluster/` +Run: `uv run pytest --cov=embedding_cluster --cov-report=term-missing --cov-fail-under=70` + +Fix any failures. + +--- + +## Task 11: Frontend - Build check + +**Step 1: Run frontend build** + +Run: `cd frontend && npm run build` + +Fix any TypeScript or build errors. + +--- + +## Task 12: Update README.md + +**Files:** +- Modify: `README.md` + +**Step 1: Add search feature to Features list** + +Add bullet: `* Semantic search within clusters — find similar items by text query or image URL, with results highlighted in the 3D view.` + +**Step 2: Update Web UI section** + +Add to the Web UI bullet list: +`* **Search** -- Type a text query or paste an image URL to find the most similar items in a collection. Results are highlighted in the 3D scatter plot with distance scores.` diff --git a/docs/plans/2026-02-25-real-time-indexing-progress.md b/docs/plans/2026-02-25-real-time-indexing-progress.md new file mode 100644 index 0000000..92e467b --- /dev/null +++ b/docs/plans/2026-02-25-real-time-indexing-progress.md @@ -0,0 +1,184 @@ +# Real-time Indexing Progress Backend Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Add on_log callbacks, heartbeat, and completion broadcasts for real-time indexing progress while preserving existing behavior. + +**Architecture:** Extend the indexer to emit optional log callbacks and more frequent progress updates without changing existing data structures. Wire the server route to broadcast log and heartbeat messages over WebSocket, and send a completion payload with collection names and final stats. Maintain synchronous callback signatures and fire-and-forget broadcasts via asyncio tasks. + +**Tech Stack:** Python 3.13, FastAPI, asyncio, mypy (strict), ruff + +--- + +### Task 1: Extend indexer callbacks and logging + +**Files:** +- Modify: `embedding_cluster/indexer.py:33-143` + +**Step 1: Write the failing test** + +```python +def test_main_indexer_calls_on_log_and_progress_callbacks(): + # TODO: add test when allowed by scope + assert True +``` + +**Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/test_indexer.py::test_main_indexer_calls_on_log_and_progress_callbacks -v` +Expected: FAIL with missing test or behavior + +**Step 3: Write minimal implementation** + +```python +PROGRESS_UPDATE_INTERVAL = 10 + +async def main_indexer( + settings: Settings, + on_progress: Callable[[dict[str, Any]], None] | None = None, + on_log: Callable[[str, str, str], None] | None = None, + cancel_event: asyncio.Event | None = None, +) -> None: + # wrap model loading with try/except and call on_log for failures + # emit on_log at requested milestones + # emit on_progress every PROGRESS_UPDATE_INTERVAL rows +``` + +**Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/test_indexer.py::test_main_indexer_calls_on_log_and_progress_callbacks -v` +Expected: PASS + +**Step 5: Commit** + +```bash +git add embedding_cluster/indexer.py +git commit -m "feat(indexer): add log callbacks and progress updates" +``` + +### Task 2: Wire WebSocket logging, heartbeat, and completion message + +**Files:** +- Modify: `embedding_cluster/server/routes/index.py:1-128` + +**Step 1: Write the failing test** + +```python +def test_index_routes_broadcast_log_heartbeat_completed(): + # TODO: add test when allowed by scope + assert True +``` + +**Step 2: Run test to verify it fails** + +Run: `uv run pytest tests/test_main.py::test_index_routes_broadcast_log_heartbeat_completed -v` +Expected: FAIL with missing test or behavior + +**Step 3: Write minimal implementation** + +```python +import time + +def _get_collection_names(settings: Settings) -> list[str]: + names: list[str] = [] + prefix = settings.chromadb_collection_prefix + if settings.image_embedding_fields: + for field in settings.image_embedding_fields: + names.append(f"{prefix}{field}") + if settings.text_embedding_fields: + for field in settings.text_embedding_fields: + names.append(f"{prefix}{field}") + return names + +def on_log(message: str, level: str, verbosity: str) -> None: + asyncio.create_task( + ws_manager.broadcast( + task_state.job_id, + { + "type": "log", + "level": level, + "message": message, + "verbosity": verbosity, + }, + ) + ) + +async def heartbeat() -> None: + while task_state.status in (TaskStatus.PENDING, TaskStatus.RUNNING): + await ws_manager.broadcast( + task_state.job_id, + {"type": "heartbeat", "elapsed_seconds": time.perf_counter() - start_time}, + ) + await asyncio.sleep(3) + +start_time = time.perf_counter() +heartbeat_task = asyncio.create_task(heartbeat()) + +await main_indexer( + settings, + on_progress=on_progress, + on_log=on_log, + cancel_event=task_state.cancel_event, +) + +heartbeat_task.cancel() +try: + await heartbeat_task +except asyncio.CancelledError: + pass + +task_state.status = TaskStatus.COMPLETED +final_progress = task_state.progress +asyncio.create_task( + ws_manager.broadcast( + task_state.job_id, + { + "type": "completed", + "status": "completed", + "progress": final_progress, + "total_indexed": final_progress.get("rows_indexed", 0), + "collection_names": _get_collection_names(settings), + }, + ) +) +``` + +**Step 4: Run test to verify it passes** + +Run: `uv run pytest tests/test_main.py::test_index_routes_broadcast_log_heartbeat_completed -v` +Expected: PASS + +**Step 5: Commit** + +```bash +git add embedding_cluster/server/routes/index.py +git commit -m "feat(server): add websocket logs and heartbeat" +``` + +### Task 3: Lint, format, and type-check + +**Files:** +- Modify: `embedding_cluster/indexer.py` +- Modify: `embedding_cluster/server/routes/index.py` + +**Step 1: Run ruff check** + +Run: `uv run ruff check embedding_cluster/indexer.py embedding_cluster/server/routes/index.py` +Expected: PASS + +**Step 2: Run ruff format** + +Run: `uv run ruff format embedding_cluster/indexer.py embedding_cluster/server/routes/index.py` +Expected: PASS + +**Step 3: Run mypy** + +Run: `uv run mypy embedding_cluster/indexer.py embedding_cluster/server/routes/index.py` +Expected: PASS + +**Step 4: Commit** + +```bash +git add embedding_cluster/indexer.py embedding_cluster/server/routes/index.py +git commit -m "chore: lint and typecheck indexing changes" +``` diff --git a/embedding_cluster/server/app.py b/embedding_cluster/server/app.py index 3a9a5b3..937c0c4 100644 --- a/embedding_cluster/server/app.py +++ b/embedding_cluster/server/app.py @@ -14,6 +14,9 @@ from embedding_cluster.server.routes.csv import router as csv_router from embedding_cluster.server.routes.index import router as index_router from embedding_cluster.server.routes.plot import router as plot_router +from embedding_cluster.server.routes.search import ( + router as search_router, +) logger = logging.getLogger(__name__) @@ -41,6 +44,7 @@ async def health_check() -> dict[str, str]: app.include_router(csv_router) app.include_router(index_router) app.include_router(plot_router) + app.include_router(search_router) if FRONTEND_DIR.is_dir(): app.mount( diff --git a/embedding_cluster/server/models.py b/embedding_cluster/server/models.py index f57a24c..73946fe 100644 --- a/embedding_cluster/server/models.py +++ b/embedding_cluster/server/models.py @@ -93,3 +93,23 @@ class PlotResponse(BaseModel): points: list[PlotPoint] clusters: list[PlotCluster] total_points: int + + +class SearchResult(BaseModel): + id: str + distance: float + metadata: dict[str, Any] + + +class SearchRequest(BaseModel): + collection_name: str + query_text: str | None = None + query_image_url: str | None = None + n_results: int = 10 + model_type: str = "text" + image_model_name: str = "openai/clip-vit-base-patch32" + text_model_name: str = "BAAI/bge-small-en-v1.5" + + +class SearchResponse(BaseModel): + results: list[SearchResult] diff --git a/embedding_cluster/server/routes/search.py b/embedding_cluster/server/routes/search.py new file mode 100644 index 0000000..9bd4c7f --- /dev/null +++ b/embedding_cluster/server/routes/search.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import chromadb +import torch +from fastapi import APIRouter, HTTPException +from sentence_transformers import SentenceTransformer +from transformers import CLIPModel, CLIPProcessor + +from embedding_cluster.server.models import ( + SearchRequest, + SearchResponse, + SearchResult, +) +from embedding_cluster.utils import ImageDownloader + +if TYPE_CHECKING: + from chromadb.api import ClientAPI + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/api/search", tags=["search"]) + +# Lazy-loaded model cache +_model_cache: dict[str, Any] = {} + + +def _get_chromadb_client() -> ClientAPI: + return chromadb.PersistentClient(path="./chromadb") + + +def _get_text_model(model_name: str) -> SentenceTransformer: + cache_key = f"text:{model_name}" + if cache_key not in _model_cache: + logger.info("Loading text model: %s", model_name) + _model_cache[cache_key] = SentenceTransformer(model_name) + return _model_cache[cache_key] # type: ignore[no-any-return] + + +def _get_image_model( + model_name: str, +) -> tuple[CLIPModel, CLIPProcessor]: + cache_key = f"image:{model_name}" + if cache_key not in _model_cache: + logger.info("Loading image model: %s", model_name) + _model_cache[cache_key] = ( + CLIPModel.from_pretrained(model_name), + CLIPProcessor.from_pretrained(model_name), + ) + return _model_cache[cache_key] # type: ignore[no-any-return] + + +async def _generate_text_embedding(query_text: str, model_name: str) -> list[float]: + model = _get_text_model(model_name) + embedding = model.encode(query_text, show_progress_bar=False) + return embedding.tolist() + + +async def _generate_image_embedding(image_url: str, model_name: str) -> list[float]: + image = await ImageDownloader().download_image_exp_backoff(image_url) + if image is None: + msg = f"Failed to download image from {image_url}" + raise ValueError(msg) + + model, processor = _get_image_model(model_name) + inputs = processor(text=None, images=image, return_tensors="pt", padding=True) + with torch.no_grad(): + img_features = model.get_image_features(inputs["pixel_values"]) + return img_features.squeeze(0).cpu().numpy().tolist() # type: ignore[no-any-return] + + +async def _generate_clip_text_embedding(query_text: str, model_name: str) -> list[float]: + """Encode text using CLIP's text encoder. + + This produces embeddings in the same vector space as CLIP image + embeddings, enabling text-to-image similarity search. + """ + model, processor = _get_image_model(model_name) + inputs = processor(text=query_text, images=None, return_tensors="pt", padding=True) + with torch.no_grad(): + text_features = model.get_text_features(inputs["input_ids"]) + return text_features.squeeze(0).cpu().numpy().tolist() # type: ignore[no-any-return] + + +@router.post("", response_model=SearchResponse) +async def search_collection( + request: SearchRequest, +) -> SearchResponse: + if not request.query_text and not request.query_image_url: + raise HTTPException( + status_code=400, + detail="Either query_text or query_image_url is required", + ) + + client = _get_chromadb_client() + try: + collection = client.get_collection(request.collection_name) + except Exception as e: + raise HTTPException( + status_code=404, + detail=(f"Collection not found: {request.collection_name}"), + ) from e + + if collection.count() == 0: + return SearchResponse(results=[]) + + # Determine model from collection metadata + metadata = collection.metadata or {} + stored_model_name = metadata.get("model_name") + stored_model_type = metadata.get("model_type") + + # Generate embedding + try: + if request.query_text: + if stored_model_type == "image" and stored_model_name: + # Collection was indexed with CLIP — use CLIP text encoder + embedding = await _generate_clip_text_embedding( + request.query_text, stored_model_name + ) + else: + # Text collection or no metadata — use SentenceTransformer + model_name = ( + stored_model_name + if stored_model_name and stored_model_type == "text" + else request.text_model_name + ) + embedding = await _generate_text_embedding(request.query_text, model_name) + else: + assert request.query_image_url is not None + model_name = ( + stored_model_name + if stored_model_name and stored_model_type == "image" + else request.image_model_name + ) + embedding = await _generate_image_embedding( + request.query_image_url, model_name + ) + except ValueError as e: + raise HTTPException( + status_code=500, + detail=str(e), + ) from e + + # Query ChromaDB + query_result = collection.query( + query_embeddings=[embedding], # type: ignore[arg-type] + n_results=min(request.n_results, collection.count()), + ) + + results: list[SearchResult] = [] + if query_result["ids"] and query_result["distances"]: + ids = query_result["ids"][0] + distances = query_result["distances"][0] + metadatas = ( + query_result["metadatas"][0] if query_result["metadatas"] else [{}] * len(ids) + ) + for i, doc_id in enumerate(ids): + results.append( + SearchResult( + id=doc_id, + distance=distances[i], + metadata=dict(metadatas[i]) if metadatas[i] else {}, + ) + ) + + return SearchResponse(results=results) diff --git a/embedding_cluster/utils.py b/embedding_cluster/utils.py index f22238f..9b70f8c 100644 --- a/embedding_cluster/utils.py +++ b/embedding_cluster/utils.py @@ -166,7 +166,13 @@ def get_or_create_chromadb_collections( f"{settings.chromadb_collection_prefix}{image_embedding_field}" ) chromadb_collections[collection_name] = ( - chromadb_client.get_or_create_collection(collection_name) + chromadb_client.get_or_create_collection( + collection_name, + metadata={ + "model_name": settings.image_model_name, + "model_type": "image", + }, + ) ) if settings.text_embedding_fields is not None: for text_embedding_field in settings.text_embedding_fields: @@ -174,7 +180,13 @@ def get_or_create_chromadb_collections( f"{settings.chromadb_collection_prefix}{text_embedding_field}" ) chromadb_collections[collection_name] = ( - chromadb_client.get_or_create_collection(collection_name) + chromadb_client.get_or_create_collection( + collection_name, + metadata={ + "model_name": settings.text_model_name, + "model_type": "text", + }, + ) ) return chromadb_collections diff --git a/frontend/.gitignore b/frontend/.gitignore index a547bf3..c880eae 100644 --- a/frontend/.gitignore +++ b/frontend/.gitignore @@ -22,3 +22,11 @@ dist-ssr *.njsproj *.sln *.sw? + + +# Playwright +/test-results/ +/playwright-report/ +/blob-report/ +/playwright/.cache/ +/e2e/.auth/ \ No newline at end of file diff --git a/frontend/e2e/fixtures.ts b/frontend/e2e/fixtures.ts new file mode 100644 index 0000000..0d46ab1 --- /dev/null +++ b/frontend/e2e/fixtures.ts @@ -0,0 +1,24 @@ +import { test as base, expect } from '@playwright/test' + +/** + * Extended test fixtures for embedding-clusters E2E tests. + * + * Provides: + * - plotPage: Navigates to the Plot page with a collection pre-selected + * via the URL search param, and waits for collection details to load + */ +export const test = base.extend<{ + plotPage: void +}>({ + plotPage: async ({ page }, use) => { + // Navigate with collection param so SearchBar renders after compute + await page.goto('/plot?collection=fashion_productDisplayName') + // Wait for the collection dropdown to be populated with options + await expect( + page.locator('select option:not([value=""])').first() + ).toBeAttached({ timeout: 10_000 }) + await use() + }, +}) + +export { expect } diff --git a/frontend/e2e/search.spec.ts b/frontend/e2e/search.spec.ts new file mode 100644 index 0000000..f3ffc74 --- /dev/null +++ b/frontend/e2e/search.spec.ts @@ -0,0 +1,182 @@ +import { test, expect } from './fixtures' + +// These tests require pre-indexed data in ChromaDB. +// Run the indexing command from the README first: +// RUNNING_MODE=INDEX LOCAL_CSV_FILENAME=./embedding_cluster/csv/fashion_small.csv \ +// ID_FIELD=id TEXT_EMBEDDING_FIELDS='["productDisplayName"]' \ +// CHROMADB_COLLECTION_PREFIX=fashion_ uv run python -m embedding_cluster + +test.describe('Semantic Search', () => { + test.beforeEach(async ({ page, plotPage: _ }) => { + // Wait for collection details to load (Compute button appears) + await expect( + page.getByRole('button', { name: 'Compute Plot' }) + ).toBeVisible({ timeout: 10_000 }) + + // Click Compute and wait for the plot to render + await page.getByRole('button', { name: 'Compute Plot' }).click() + + // Wait for the canvas to appear (plot rendered) + await expect(page.locator('canvas')).toBeVisible({ timeout: 120_000 }) + + // Wait for the search bar to be available + await expect( + page.getByRole('heading', { name: 'Semantic Search' }) + ).toBeVisible({ timeout: 5_000 }) + }) + + test('search bar appears after computing plot', async ({ page }) => { + // Text radio should be selected by default + const textRadio = page.getByRole('radio', { name: 'Text' }) + await expect(textRadio).toBeChecked() + + // Search input should be present + await expect( + page.getByPlaceholder('Search by text...') + ).toBeVisible() + + // Search button should be visible but disabled (no query) + const searchButton = page.getByRole('button', { name: 'Search' }) + await expect(searchButton).toBeVisible() + await expect(searchButton).toBeDisabled() + }) + + test('text search returns results', async ({ page }) => { + // Type a search query + await page.getByPlaceholder('Search by text...').fill('blue shirt') + + // Search button should be enabled now + const searchButton = page.getByRole('button', { name: 'Search' }) + await expect(searchButton).toBeEnabled() + + // Click search + await searchButton.click() + + // Wait for results to appear -- "Results" heading with count + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // "Highlight All" button should be visible + await expect( + page.getByRole('button', { name: 'Highlight All' }) + ).toBeVisible() + + // Result items should be present (at least one) + const resultItems = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }) + await expect(resultItems.first()).toBeVisible() + }) + + test('search via Enter key', async ({ page }) => { + const searchInput = page.getByPlaceholder('Search by text...') + await searchInput.fill('casual shoes') + await searchInput.press('Enter') + + // Results should appear + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + }) + + test('clear search removes results', async ({ page }) => { + // Perform a search first + await page.getByPlaceholder('Search by text...').fill('jacket') + await page.getByRole('button', { name: 'Search' }).click() + + // Wait for results + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // Click Clear + await page.getByRole('button', { name: 'Clear' }).click() + + // Results should disappear + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeHidden() + + // Search input should be empty + await expect( + page.getByPlaceholder('Search by text...') + ).toHaveValue('') + }) + + test('clicking a result highlights it', async ({ page }) => { + // Search + await page.getByPlaceholder('Search by text...').fill('men') + await page.getByRole('button', { name: 'Search' }).click() + + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // Click the first result -- it should get the active style (border-l-2) + const firstResult = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }).first() + await firstResult.click() + + // The clicked result should have the active indicator (blue left border) + await expect(firstResult).toHaveClass(/border-blue-500/) + }) + + test('highlight all button activates all results', async ({ page }) => { + // Search + await page.getByPlaceholder('Search by text...').fill('shirt') + await page.getByRole('button', { name: 'Search' }).click() + + await expect( + page.getByRole('heading', { name: /Results \(\d+\)/ }) + ).toBeVisible({ timeout: 30_000 }) + + // Click a single result first to narrow highlight + const firstResult = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }).first() + await firstResult.click() + + // Now click "Highlight All" + await page.getByRole('button', { name: 'Highlight All' }).click() + + // All result buttons should have the active class + const allResults = page.locator('button').filter({ + has: page.locator('.text-xs.text-gray-400'), + }) + const count = await allResults.count() + expect(count).toBeGreaterThan(1) + }) + + test('switch to image URL mode', async ({ page }) => { + // Click image radio + await page.getByText('Image URL').click() + + // Placeholder should change + await expect( + page.getByPlaceholder('Paste image URL...') + ).toBeVisible() + + // Text placeholder should be gone + await expect( + page.getByPlaceholder('Search by text...') + ).toBeHidden() + }) + + test('adjusting results slider changes value', async ({ page }) => { + // The slider label should show the default value + await expect( + page.getByText('Results: 10') + ).toBeVisible() + + // Adjust the slider + const slider = page.locator('input[type="range"]').last() + await slider.fill('25') + + // Label should update + await expect( + page.getByText('Results: 25') + ).toBeVisible() + }) +}) diff --git a/frontend/package-lock.json b/frontend/package-lock.json index 23055a2..bf186fd 100644 --- a/frontend/package-lock.json +++ b/frontend/package-lock.json @@ -19,6 +19,7 @@ }, "devDependencies": { "@eslint/js": "^9.39.1", + "@playwright/test": "^1.58.2", "@tailwindcss/vite": "^4.2.0", "@types/node": "^24.10.1", "@types/react": "^19.2.7", @@ -1051,6 +1052,22 @@ "three": ">= 0.159.0" } }, + "node_modules/@playwright/test": { + "version": "1.58.2", + "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.2.tgz", + "integrity": "sha512-akea+6bHYBBfA9uQqSYmlJXn61cTa+jbO87xVLCWbTqbWadRVmhxlXATaOjOgcBaWU4ePo0wB41KMFv3o35IXA==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright": "1.58.2" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, "node_modules/@react-three/drei": { "version": "10.7.7", "resolved": "https://registry.npmjs.org/@react-three/drei/-/drei-10.7.7.tgz", @@ -1719,6 +1736,70 @@ "node": ">=14.0.0" } }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/core": { + "version": "1.8.1", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "@emnapi/wasi-threads": "1.1.0", + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/runtime": { + "version": "1.8.1", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@emnapi/wasi-threads": { + "version": "1.1.0", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@napi-rs/wasm-runtime": { + "version": "1.1.1", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "@emnapi/core": "^1.7.1", + "@emnapi/runtime": "^1.7.1", + "@tybys/wasm-util": "^0.10.1" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/Brooooooklyn" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/@tybys/wasm-util": { + "version": "0.10.1", + "dev": true, + "inBundle": true, + "license": "MIT", + "optional": true, + "dependencies": { + "tslib": "^2.4.0" + } + }, + "node_modules/@tailwindcss/oxide-wasm32-wasi/node_modules/tslib": { + "version": "2.8.1", + "dev": true, + "inBundle": true, + "license": "0BSD", + "optional": true + }, "node_modules/@tailwindcss/oxide-win32-arm64-msvc": { "version": "4.2.0", "resolved": "https://registry.npmjs.org/@tailwindcss/oxide-win32-arm64-msvc/-/oxide-win32-arm64-msvc-4.2.0.tgz", @@ -3804,6 +3885,53 @@ "url": "https://github.com/sponsors/jonschlinkert" } }, + "node_modules/playwright": { + "version": "1.58.2", + "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.2.tgz", + "integrity": "sha512-vA30H8Nvkq/cPBnNw4Q8TWz1EJyqgpuinBcHET0YVJVFldr8JDNiU9LaWAE1KqSkRYazuaBhTpB5ZzShOezQ6A==", + "dev": true, + "license": "Apache-2.0", + "dependencies": { + "playwright-core": "1.58.2" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "fsevents": "2.3.2" + } + }, + "node_modules/playwright-core": { + "version": "1.58.2", + "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.2.tgz", + "integrity": "sha512-yZkEtftgwS8CsfYo7nm0KE8jsvm6i/PTgVtB8DL726wNf6H2IMsDuxCpJj59KDaxCtSnrWan2AeDqM7JBaultg==", + "dev": true, + "license": "Apache-2.0", + "bin": { + "playwright-core": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/playwright/node_modules/fsevents": { + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", + "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, "node_modules/postcss": { "version": "8.5.6", "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.6.tgz", diff --git a/frontend/package.json b/frontend/package.json index 68bfd98..b726576 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -7,7 +7,9 @@ "dev": "vite", "build": "tsc -b && vite build", "lint": "eslint .", - "preview": "vite preview" + "preview": "vite preview", + "test:e2e": "playwright test", + "test:e2e:ui": "playwright test --ui" }, "dependencies": { "@react-three/drei": "^10.7.7", @@ -21,6 +23,7 @@ }, "devDependencies": { "@eslint/js": "^9.39.1", + "@playwright/test": "^1.58.2", "@tailwindcss/vite": "^4.2.0", "@types/node": "^24.10.1", "@types/react": "^19.2.7", diff --git a/frontend/playwright.config.ts b/frontend/playwright.config.ts new file mode 100644 index 0000000..01d9ec2 --- /dev/null +++ b/frontend/playwright.config.ts @@ -0,0 +1,28 @@ +import { defineConfig, devices } from '@playwright/test' + +export default defineConfig({ + testDir: './e2e', + fullyParallel: true, + forbidOnly: !!process.env.CI, + retries: process.env.CI ? 2 : 0, + workers: process.env.CI ? 1 : undefined, + reporter: 'html', + timeout: 60_000, + use: { + baseURL: 'http://localhost:8000', + trace: 'on-first-retry', + screenshot: 'only-on-failure', + }, + projects: [ + { + name: 'chromium', + use: { ...devices['Desktop Chrome'] }, + }, + ], + webServer: { + command: 'cd .. && RUNNING_MODE=SERVER uv run python -m embedding_cluster', + url: 'http://localhost:8000/api/health', + reuseExistingServer: !process.env.CI, + timeout: 30_000, + }, +}) diff --git a/frontend/src/api/plot.ts b/frontend/src/api/plot.ts index c7460c6..5c9992e 100644 --- a/frontend/src/api/plot.ts +++ b/frontend/src/api/plot.ts @@ -1,4 +1,4 @@ -import type { IndexStartResponse, PlotRequest, PlotResponse } from "../types"; +import type { IndexStartResponse, PlotRequest, PlotResponse, SearchRequest, SearchResponse } from "../types"; import { apiFetch, apiPost } from "./client"; export async function startPlotCompute( @@ -14,3 +14,9 @@ export async function getPlotData( `/plot/data/${jobId}`, ); } + +export async function searchCollection( + request: SearchRequest, +): Promise { + return apiPost("/search", request); +} diff --git a/frontend/src/components/plot/ImageSpriteCloud.tsx b/frontend/src/components/plot/ImageSpriteCloud.tsx index c68832c..a0ba776 100644 --- a/frontend/src/components/plot/ImageSpriteCloud.tsx +++ b/frontend/src/components/plot/ImageSpriteCloud.tsx @@ -19,12 +19,13 @@ interface PointSpriteProps { color: string imageUrl: string | null size: number + opacity: number onHover: (id: string | null) => void } -function PointSprite({ point, color, imageUrl, size, onHover }: PointSpriteProps) { +function PointSprite({ point, color, imageUrl, size, opacity, onHover }: PointSpriteProps) { if (imageUrl) { - return + return } const scale = size * 0.2 @@ -38,12 +39,12 @@ function PointSprite({ point, color, imageUrl, size, onHover }: PointSpriteProps }} onPointerOut={() => onHover(null)} > - + ) } -function TextureSprite({ point, imageUrl, size, onHover }: PointSpriteProps & { imageUrl: string }) { +function TextureSprite({ point, imageUrl, size, opacity, onHover }: PointSpriteProps & { imageUrl: string }) { const texture = useTexture(imageUrl) const scale = size * 0.6 @@ -62,6 +63,7 @@ function TextureSprite({ point, imageUrl, size, onHover }: PointSpriteProps & { map={texture} color={'white'} transparent={true} + opacity={opacity} /> ) @@ -80,6 +82,7 @@ export default function ImageSpriteCloud() { const plotData = usePlotStore((state) => state.plotData) const visibleClusters = usePlotStore((state) => state.visibleClusters) const pointSize = usePlotStore((state) => state.pointSize) + const highlightedIds = usePlotStore((state) => state.highlightedIds) const setHoveredPointId = usePlotStore((state) => state.setHoveredPointId) const visiblePoints = useMemo(() => { @@ -91,21 +94,24 @@ export default function ImageSpriteCloud() { return visiblePoints.slice(0, MAX_SPRITES).map((point) => { const color = CLUSTER_COLORS[point.cluster % CLUSTER_COLORS.length] const imageUrl = getImageUrl(point.metadata) - return { point, color, imageUrl } + const isHighlighted = highlightedIds.size === 0 || highlightedIds.has(point.id) + const opacity = isHighlighted ? 1.0 : 0.15 + return { point, color, imageUrl, opacity } }) - }, [visiblePoints]) + }, [visiblePoints, highlightedIds]) if (!plotData) return null return ( - {spritesToRender.map(({ point, color, imageUrl }) => ( + {spritesToRender.map(({ point, color, imageUrl, opacity }) => ( }> diff --git a/frontend/src/components/plot/InstancedSpheres.tsx b/frontend/src/components/plot/InstancedSpheres.tsx index 4c8f9ea..016618e 100644 --- a/frontend/src/components/plot/InstancedSpheres.tsx +++ b/frontend/src/components/plot/InstancedSpheres.tsx @@ -8,6 +8,7 @@ export default function InstancedSpheres() { const plotData = usePlotStore((state) => state.plotData) const visibleClusters = usePlotStore((state) => state.visibleClusters) const pointSize = usePlotStore((state) => state.pointSize) + const highlightedIds = usePlotStore((state) => state.highlightedIds) const setHoveredPointId = usePlotStore((state) => state.setHoveredPointId) const { filteredPoints, filteredPointIds } = useMemo(() => { @@ -24,19 +25,22 @@ export default function InstancedSpheres() { const matrix = new THREE.Matrix4() // Pre-create color objects for efficiency const colorObjects = CLUSTER_COLORS.map(hex => new THREE.Color(hex)) + const hasHighlights = highlightedIds.size > 0 + for (let i = 0; i < filteredPoints.length; i++) { const p = filteredPoints[i] matrix.setPosition(p.x, p.y, p.z) mesh.setMatrixAt(i, matrix) - const c = colorObjects[p.cluster % colorObjects.length] + const dimFactor = hasHighlights && !highlightedIds.has(p.id) ? 0.15 : 1.0 + const c = colorObjects[p.cluster % colorObjects.length].clone().multiplyScalar(dimFactor) mesh.setColorAt(i, c) } mesh.instanceMatrix.needsUpdate = true if (mesh.instanceColor) mesh.instanceColor.needsUpdate = true - }, [filteredPoints]) + }, [filteredPoints, highlightedIds]) const handlePointerMove = useCallback((e: ThreeEvent) => { e.stopPropagation() diff --git a/frontend/src/components/plot/ParticleCloud.tsx b/frontend/src/components/plot/ParticleCloud.tsx index ede4d82..1605985 100644 --- a/frontend/src/components/plot/ParticleCloud.tsx +++ b/frontend/src/components/plot/ParticleCloud.tsx @@ -8,6 +8,7 @@ export default function ParticleCloud() { const plotData = usePlotStore((state) => state.plotData) const visibleClusters = usePlotStore((state) => state.visibleClusters) const pointSize = usePlotStore((state) => state.pointSize) + const highlightedIds = usePlotStore((state) => state.highlightedIds) const setHoveredPointId = usePlotStore((state) => state.setHoveredPointId) // Memoize positions, colors, and the mapping back to original point IDs @@ -29,6 +30,8 @@ export default function ParticleCloud() { const colorObjects = CLUSTER_COLORS.map(hex => new THREE.Color(hex)) + const hasHighlights = highlightedIds.size > 0 + for (let i = 0; i < count; i++) { const p = filteredPoints[i] @@ -39,16 +42,17 @@ export default function ParticleCloud() { // Color const color = colorObjects[p.cluster % colorObjects.length] - cols[i * 3] = color.r - cols[i * 3 + 1] = color.g - cols[i * 3 + 2] = color.b + const dimFactor = hasHighlights && !highlightedIds.has(p.id) ? 0.15 : 1.0 + cols[i * 3] = color.r * dimFactor + cols[i * 3 + 1] = color.g * dimFactor + cols[i * 3 + 2] = color.b * dimFactor // ID mapping ids[i] = p.id } return { positions: pos, colors: cols, filteredPointIds: ids } - }, [plotData, visibleClusters]) + }, [plotData, visibleClusters, highlightedIds]) const handlePointerMove = useCallback((e: ThreeEvent) => { // Stop event propagation so we don't trigger other things diff --git a/frontend/src/components/plot/PlotControls.tsx b/frontend/src/components/plot/PlotControls.tsx index 23197a1..7c65427 100644 --- a/frontend/src/components/plot/PlotControls.tsx +++ b/frontend/src/components/plot/PlotControls.tsx @@ -11,7 +11,7 @@ interface PlotControlsProps { } export default function PlotControls({ onCompute, isComputing }: PlotControlsProps) { - const [searchParams] = useSearchParams() + const [searchParams, setSearchParams] = useSearchParams() const [selectedCollection, setSelectedCollection] = useState(searchParams.get('collection') || '') const [numClusters, setNumClusters] = useState(10) const [textDisplayFields, setTextDisplayFields] = useState([]) @@ -80,8 +80,10 @@ export default function PlotControls({ onCompute, isComputing }: PlotControlsPro setQueryType(type)} + className="mr-1" + /> + {type === 'text' ? 'Text' : 'Image URL'} + + ))} + + + setQueryValue(e.target.value)} + onKeyDown={(e) => e.key === 'Enter' && handleSearch()} + placeholder={queryType === 'text' ? 'Search by text...' : 'Paste image URL...'} + className="w-full border border-gray-300 rounded-md px-3 py-2 text-sm" + /> + +
+ + setNResults(Number(e.target.value))} + className="w-full" + /> +
+ +
+ + +
+ + {mutation.isError && ( +

+ Search failed: {(mutation.error as Error).message} +

+ )} + + ) +} diff --git a/frontend/src/components/plot/SearchResults.tsx b/frontend/src/components/plot/SearchResults.tsx new file mode 100644 index 0000000..c0f0ba2 --- /dev/null +++ b/frontend/src/components/plot/SearchResults.tsx @@ -0,0 +1,84 @@ +import { usePlotStore } from '../../stores/plotStore' + +export default function SearchResults() { + const searchResults = usePlotStore((state) => state.searchResults) + const highlightedIds = usePlotStore((state) => state.highlightedIds) + const setHighlightedIds = usePlotStore((state) => state.setHighlightedIds) + const setHoveredPointId = usePlotStore((state) => state.setHoveredPointId) + + if (!searchResults || searchResults.length === 0) return null + + const handleClickResult = (id: string) => { + setHighlightedIds(new Set([id])) + } + + const handleShowAll = () => { + setHighlightedIds(new Set(searchResults.map((r) => r.id))) + } + + const getImageUrl = (metadata: Record): string | null => { + for (const value of Object.values(metadata)) { + if (typeof value === 'string' && (value.startsWith('http') || value.startsWith('/'))) { + return value + } + } + return null + } + + return ( +
+
+

+ Results ({searchResults.length}) +

+ +
+
+ {searchResults.map((result) => { + const imageUrl = getImageUrl(result.metadata) + const isActive = highlightedIds.has(result.id) + return ( + + ) + })} +
+
+ ) +} diff --git a/frontend/src/pages/PlotPage.tsx b/frontend/src/pages/PlotPage.tsx index 4ae345a..b71bf1c 100644 --- a/frontend/src/pages/PlotPage.tsx +++ b/frontend/src/pages/PlotPage.tsx @@ -1,7 +1,10 @@ import { useRef, useState, useCallback } from 'react' +import { useSearchParams } from 'react-router-dom' import PlotControls from '../components/plot/PlotControls' import ScatterPlot from '../components/plot/ScatterPlot' import ClusterLegend from '../components/plot/ClusterLegend' +import SearchBar from '../components/plot/SearchBar' +import SearchResults from '../components/plot/SearchResults' import { usePlotData } from '../hooks/usePlotData' import { usePlotStore } from '../stores/plotStore' @@ -10,6 +13,8 @@ export default function PlotPage() { const plotData = usePlotStore((state) => state.plotData) const plotContainerRef = useRef(null) const [isFullscreen, setIsFullscreen] = useState(false) + const [searchParams] = useSearchParams() + const collectionName = searchParams.get('collection') ?? '' const toggleFullscreen = useCallback(() => { if (!plotContainerRef.current) return @@ -28,6 +33,12 @@ export default function PlotPage() {
+ {plotData && collectionName && ( + <> + + + + )}
diff --git a/frontend/src/stores/plotStore.ts b/frontend/src/stores/plotStore.ts index 86eeec8..3b6d9da 100644 --- a/frontend/src/stores/plotStore.ts +++ b/frontend/src/stores/plotStore.ts @@ -1,5 +1,5 @@ import { create } from 'zustand' -import type { PlotResponse } from '../types' +import type { PlotResponse, SearchResult } from '../types' interface PlotState { plotData: PlotResponse | null @@ -7,6 +7,10 @@ interface PlotState { hoveredPointId: string | null renderMode: 'particles' | 'sprites' | 'spheres' pointSize: number + searchResults: SearchResult[] | null + highlightedIds: Set + isSearching: boolean + queryPoint: { x: number; y: number; z: number } | null // actions setPlotData: (data: PlotResponse | null) => void toggleCluster: (index: number) => void @@ -14,6 +18,11 @@ interface PlotState { setRenderMode: (mode: 'particles' | 'sprites' | 'spheres') => void setPointSize: (size: number) => void resetVisibleClusters: (clusterCount: number) => void + setSearchResults: (results: SearchResult[] | null) => void + setHighlightedIds: (ids: Set) => void + setIsSearching: (searching: boolean) => void + clearSearch: () => void + setQueryPoint: (point: { x: number; y: number; z: number } | null) => void } export const CLUSTER_COLORS = [ @@ -29,6 +38,10 @@ export const usePlotStore = create((set) => ({ hoveredPointId: null, renderMode: 'particles', pointSize: 5, + searchResults: null, + highlightedIds: new Set(), + isSearching: false, + queryPoint: null, setPlotData: (data) => set({ plotData: data }), @@ -53,4 +66,18 @@ export const usePlotStore = create((set) => ({ set({ visibleClusters: new Set(Array.from({ length: clusterCount }, (_, i) => i)), }), + + setSearchResults: (results) => set({ + searchResults: results, + highlightedIds: new Set(results?.map(r => r.id) ?? []), + }), + setHighlightedIds: (ids) => set({ highlightedIds: ids }), + setIsSearching: (searching) => set({ isSearching: searching }), + setQueryPoint: (point) => set({ queryPoint: point }), + clearSearch: () => set({ + searchResults: null, + highlightedIds: new Set(), + isSearching: false, + queryPoint: null, + }), })) diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 028793d..c075225 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -92,3 +92,24 @@ export interface PlotResponse { export interface MessageResponse { message: string; } + +// Search +export interface SearchResult { + id: string; + distance: number; + metadata: Record; +} + +export interface SearchRequest { + collection_name: string; + query_text?: string; + query_image_url?: string; + n_results?: number; + model_type?: string; + image_model_name?: string; + text_model_name?: string; +} + +export interface SearchResponse { + results: SearchResult[]; +} diff --git a/tests/test_server_search.py b/tests/test_server_search.py new file mode 100644 index 0000000..8950109 --- /dev/null +++ b/tests/test_server_search.py @@ -0,0 +1,464 @@ +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import numpy as np +import pytest +from httpx import ASGITransport, AsyncClient + +from embedding_cluster.server.app import create_app + + +@pytest.fixture +def app(): + return create_app() + + +@pytest.fixture +async def client(app): + transport = ASGITransport(app=app) + async with AsyncClient(transport=transport, base_url="http://test") as ac: + yield ac + + +@pytest.fixture +def mock_chromadb_client(): + mock_client = MagicMock() + return mock_client + + +@pytest.fixture +def mock_collection(mock_chromadb_client): + mock_coll = MagicMock() + mock_coll.count.return_value = 100 + mock_coll.metadata = { + "model_name": "BAAI/bge-small-en-v1.5", + "model_type": "text", + } + mock_coll.query.return_value = { + "ids": [["id1", "id2", "id3"]], + "distances": [[0.1, 0.3, 0.5]], + "metadatas": [ + [ + { + "name": "item1", + "imageUrl": "http://example.com/1.jpg", + }, + { + "name": "item2", + "imageUrl": "http://example.com/2.jpg", + }, + { + "name": "item3", + "imageUrl": "http://example.com/3.jpg", + }, + ] + ], + } + mock_chromadb_client.get_collection.return_value = mock_coll + return mock_coll + + +@pytest.fixture +def mock_text_model(): + mock_model = MagicMock() + mock_model.encode.return_value = np.zeros(384) + return mock_model + + +@pytest.fixture +def mock_image_model(): + mock_clip = MagicMock() + mock_processor = MagicMock() + + import torch + + mock_features = torch.zeros(1, 512) + mock_clip.get_image_features.return_value = mock_features + + return mock_clip, mock_processor + + +async def test_search_text_query( + client, mock_chromadb_client, mock_collection, mock_text_model +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_text_model, + ), + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_text": "red shoes", + "n_results": 3, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["results"]) == 3 + assert data["results"][0]["id"] == "id1" + assert data["results"][0]["distance"] == 0.1 + assert "metadata" in data["results"][0] + + +async def test_search_image_query( + client, mock_chromadb_client, mock_collection, mock_image_model +): + mock_clip, mock_processor = mock_image_model + mock_image = MagicMock() + + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_image_model", + return_value=(mock_clip, mock_processor), + ), + patch( + "embedding_cluster.server.routes.search.ImageDownloader" + ) as mock_downloader_cls, + ): + mock_instance = MagicMock() + mock_instance.download_image_exp_backoff = AsyncMock(return_value=mock_image) + mock_downloader_cls.return_value = mock_instance + + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_image_url": "http://example.com/query.jpg", + "model_type": "image", + "n_results": 3, + }, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["results"]) == 3 + + +async def test_search_collection_not_found(client): + mock_client = MagicMock() + mock_client.get_collection.side_effect = Exception("Not found") + + with patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_client, + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "nonexistent", + "query_text": "test", + }, + ) + + assert response.status_code == 404 + assert "Collection not found" in response.json()["detail"] + + +async def test_search_empty_query(client): + response = await client.post( + "/api/search", + json={"collection_name": "test_collection"}, + ) + assert response.status_code == 400 + + +async def test_search_custom_n_results( + client, mock_chromadb_client, mock_collection, mock_text_model +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_text_model, + ), + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_text": "blue jacket", + "n_results": 5, + }, + ) + + assert response.status_code == 200 + mock_collection.query.assert_called_once() + call_kwargs = mock_collection.query.call_args + assert ( + call_kwargs.kwargs.get("n_results") == 5 or call_kwargs[1].get("n_results") == 5 + ) + + +async def test_search_response_structure( + client, mock_chromadb_client, mock_collection, mock_text_model +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_text_model, + ), + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_text": "test", + }, + ) + + data = response.json() + assert "results" in data + for result in data["results"]: + assert "id" in result + assert "distance" in result + assert "metadata" in result + assert isinstance(result["distance"], float) + + +async def test_search_empty_collection(client, mock_chromadb_client): + mock_coll = MagicMock() + mock_coll.count.return_value = 0 + mock_chromadb_client.get_collection.return_value = mock_coll + + with patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "empty_collection", + "query_text": "test", + }, + ) + + assert response.status_code == 200 + assert response.json()["results"] == [] + + +async def test_search_image_download_failure( + client, mock_chromadb_client, mock_collection +): + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search.ImageDownloader" + ) as mock_downloader_cls, + ): + mock_instance = MagicMock() + mock_instance.download_image_exp_backoff = AsyncMock(return_value=None) + mock_downloader_cls.return_value = mock_instance + + response = await client.post( + "/api/search", + json={ + "collection_name": "test_collection", + "query_image_url": "http://example.com/bad.jpg", + "model_type": "image", + }, + ) + + assert response.status_code == 500 + + +async def test_search_text_on_image_collection_uses_clip( + client, mock_chromadb_client +): + """Text query on an image (CLIP) collection should use CLIP text encoder.""" + import torch + + mock_coll = MagicMock() + mock_coll.count.return_value = 100 + mock_coll.metadata = { + "model_name": "openai/clip-vit-base-patch32", + "model_type": "image", + } + mock_coll.query.return_value = { + "ids": [["id1"]], + "distances": [[0.2]], + "metadatas": [[{"name": "item1"}]], + } + mock_chromadb_client.get_collection.return_value = mock_coll + + mock_clip = MagicMock() + mock_processor = MagicMock() + mock_text_features = torch.zeros(1, 512) + mock_clip.get_text_features.return_value = mock_text_features + + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_image_model", + return_value=(mock_clip, mock_processor), + ), + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "clip_collection", + "query_text": "red shoes", + }, + ) + + assert response.status_code == 200 + data = response.json() + assert len(data["results"]) == 1 + # Verify CLIP text encoder was used (not SentenceTransformer) + mock_clip.get_text_features.assert_called_once() + + +async def test_search_uses_stored_text_model_name( + client, mock_chromadb_client +): + """Text query should use the model name stored in collection metadata.""" + mock_coll = MagicMock() + mock_coll.count.return_value = 100 + mock_coll.metadata = { + "model_name": "custom-text-model", + "model_type": "text", + } + mock_coll.query.return_value = { + "ids": [["id1"]], + "distances": [[0.1]], + "metadatas": [[{"name": "item1"}]], + } + mock_chromadb_client.get_collection.return_value = mock_coll + + mock_model = MagicMock() + mock_model.encode.return_value = np.zeros(384) + + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_model, + ) as mock_get_text, + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "text_collection", + "query_text": "test query", + }, + ) + + assert response.status_code == 200 + # Verify the stored model name was used + mock_get_text.assert_called_with("custom-text-model") + + +async def test_search_uses_stored_image_model_name( + client, mock_chromadb_client, mock_image_model +): + """Image query should use the model name stored in collection metadata.""" + mock_clip, mock_processor = mock_image_model + mock_image = MagicMock() + + mock_coll = MagicMock() + mock_coll.count.return_value = 100 + mock_coll.metadata = { + "model_name": "custom-clip-model", + "model_type": "image", + } + mock_coll.query.return_value = { + "ids": [["id1"]], + "distances": [[0.1]], + "metadatas": [[{"name": "item1"}]], + } + mock_chromadb_client.get_collection.return_value = mock_coll + + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_image_model", + return_value=(mock_clip, mock_processor), + ) as mock_get_image, + patch( + "embedding_cluster.server.routes.search.ImageDownloader" + ) as mock_downloader_cls, + ): + mock_instance = MagicMock() + mock_instance.download_image_exp_backoff = AsyncMock( + return_value=mock_image + ) + mock_downloader_cls.return_value = mock_instance + + response = await client.post( + "/api/search", + json={ + "collection_name": "image_collection", + "query_image_url": "http://example.com/test.jpg", + }, + ) + + assert response.status_code == 200 + mock_get_image.assert_called_with("custom-clip-model") + + +async def test_search_fallback_when_no_metadata( + client, mock_chromadb_client, mock_text_model +): + """When collection has no metadata, fall back to request defaults.""" + mock_coll = MagicMock() + mock_coll.count.return_value = 100 + mock_coll.metadata = None + mock_coll.query.return_value = { + "ids": [["id1"]], + "distances": [[0.1]], + "metadatas": [[{"name": "item1"}]], + } + mock_chromadb_client.get_collection.return_value = mock_coll + + with ( + patch( + "embedding_cluster.server.routes.search._get_chromadb_client", + return_value=mock_chromadb_client, + ), + patch( + "embedding_cluster.server.routes.search._get_text_model", + return_value=mock_text_model, + ) as mock_get_text, + ): + response = await client.post( + "/api/search", + json={ + "collection_name": "legacy_collection", + "query_text": "test", + }, + ) + + assert response.status_code == 200 + # Should fall back to default text model name from request + mock_get_text.assert_called_with("BAAI/bge-small-en-v1.5") diff --git a/tests/test_utils.py b/tests/test_utils.py index 8deede3..05a45ac 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -171,8 +171,13 @@ def test_with_image_fields(self, monkeypatch: pytest.MonkeyPatch) -> None: result = get_or_create_chromadb_collections(settings, mock_client) assert "test_imageUrl" in result - mock_client.get_or_create_collection.assert_called_with("test_imageUrl") - + mock_client.get_or_create_collection.assert_called_with( + "test_imageUrl", + metadata={ + "model_name": "openai/clip-vit-base-patch32", + "model_type": "image", + }, + ) def test_with_text_fields(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("TEXT_EMBEDDING_FIELDS", '["description"]') monkeypatch.setenv("CHROMADB_COLLECTION_PREFIX", "pre_") @@ -184,7 +189,13 @@ def test_with_text_fields(self, monkeypatch: pytest.MonkeyPatch) -> None: result = get_or_create_chromadb_collections(settings, mock_client) assert "pre_description" in result - + mock_client.get_or_create_collection.assert_called_with( + "pre_description", + metadata={ + "model_name": "BAAI/bge-small-en-v1.5", + "model_type": "text", + }, + ) def test_with_no_fields(self) -> None: settings = Settings() mock_client = MagicMock()