Add similarity API endpoint

I've refactored `compute_similarity_metadata()` to get query tag by name
and not by ID, so that `compute_similarity_metadata()` can create the
metadata name from the query name if it's not provided.

I've also introduce TagNotFoundError exception.

Author: lukas-lightly

lightly_studio/src/lightly_studio/api/routes/api/metadata.py

@@ -5,13 +5,15 @@
 from typing import List
 from uuid import UUID
 
-from fastapi import APIRouter, Depends, Path
+from fastapi import APIRouter, Depends, HTTPException, Path
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
 
 from lightly_studio.api.routes.api.dataset import get_and_validate_dataset_id
+from lightly_studio.api.routes.api.status import HTTP_STATUS_NOT_FOUND
 from lightly_studio.db_manager import SessionDep
-from lightly_studio.metadata import compute_typicality
+from lightly_studio.errors import TagNotFoundError
+from lightly_studio.metadata import compute_similarity, compute_typicality
 from lightly_studio.models.dataset import DatasetTable
 from lightly_studio.models.metadata import MetadataInfoView
 from lightly_studio.resolvers import embedding_model_resolver
@@ -89,3 +91,73 @@ def compute_typicality_metadata(
         embedding_model_id=embedding_model.embedding_model_id,
         metadata_name=request.metadata_name,
     )
+
+
+class ComputeSimilarityRequest(BaseModel):
+    """Request model for computing typicality metadata."""
+
+    embedding_model_name: str | None = Field(
+        default=None,
+        description="Embedding model name (uses default if not specified)",
+    )
+    query_tag_name: str = Field(
+        description="The name of the tag to use for the query",
+    )
+    metadata_name: str | None = Field(
+        default=None,
+        description="Metadata field name (defaults to None)",
+    )
+
+
+@metadata_router.post(
+    "/metadata/similarity",
+    status_code=200,
+    response_model=str,
+)
+def compute_similarity_metadata(
+    session: SessionDep,
+    dataset: Annotated[
+        DatasetTable,
+        Depends(get_and_validate_dataset_id),
+    ],
+    request: ComputeSimilarityRequest,
+) -> str:
+    """Compute similarity metadata for a dataset.
+
+    Args:
+        session: The database session.
+        dataset: The dataset to compute similarity for.
+        request: Request parameters including optional embedding model name
+            and metadata field name.
+
+    Returns:
+        Metadata name used for the similarity.
+
+    Raises:
+        HTTPException: 404 if invalid embedding model or query tag is given.
+    """
+    try:
+        embedding_model = embedding_model_resolver.get_by_name(
+            session=session,
+            dataset_id=dataset.dataset_id,
+            embedding_model_name=request.embedding_model_name,
+        )
+    except ValueError as e:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"embedding model {request.embedding_model_name} not found",
+        ) from e
+
+    try:
+        return compute_similarity.compute_similarity_metadata(
+            session=session,
+            key_dataset_id=dataset.dataset_id,
+            query_tag_name=request.query_tag_name,
+            embedding_model_id=embedding_model.embedding_model_id,
+            metadata_name=request.metadata_name,
+        )
+    except TagNotFoundError as e:
+        raise HTTPException(
+            status_code=HTTP_STATUS_NOT_FOUND,
+            detail=f"Query tag {request.query_tag_name} not found",
+        ) from e

lightly_studio/src/lightly_studio/core/dataset.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import datetime
 from pathlib import Path
 from typing import Iterable, Iterator
 from uuid import UUID
@@ -609,27 +608,18 @@ def compute_similarity_metadata(
         Returns:
             The name of the metadata storing the similarity values.
         """
-        query_tag = tag_resolver.get_by_name(
-            session=self.session, tag_name=query_tag_name, dataset_id=self.dataset_id
-        )
-        if query_tag is None:
-            raise ValueError("Query tag not found")
         embedding_model_id = embedding_model_resolver.get_by_name(
             session=self.session,
             dataset_id=self.dataset_id,
             embedding_model_name=embedding_model_name,
         ).embedding_model_id
-        date = datetime.datetime.now(datetime.timezone.utc)
-        if metadata_name is None:
-            metadata_name = f"similarity_{query_tag_name}_{date.isoformat()}"
-        compute_similarity.compute_similarity_metadata(
+        return compute_similarity.compute_similarity_metadata(
             session=self.session,
             key_dataset_id=self.dataset_id,
             embedding_model_id=embedding_model_id,
-            query_tag_id=query_tag.tag_id,
+            query_tag_name=query_tag_name,
             metadata_name=metadata_name,
         )
-        return metadata_name
 
 
 def _generate_embeddings(session: Session, dataset_id: UUID, sample_ids: list[UUID]) -> None:

lightly_studio/src/lightly_studio/errors.py

@@ -0,0 +1,5 @@
+"""Lightly Studio Exceptions types."""
+
+
+class TagNotFoundError(Exception):
+    """Exception signaling that a tag has not been found."""

lightly_studio/src/lightly_studio/metadata/compute_similarity.py

@@ -1,22 +1,25 @@
 """Computes similarity from embeddings."""
 
+from datetime import datetime, timezone
+from typing import Optional
 from uuid import UUID
 
 from lightly_mundig import Similarity  # type: ignore[import-untyped]
 from sqlmodel import Session
 
 from lightly_studio.dataset.env import LIGHTLY_STUDIO_LICENSE_KEY
-from lightly_studio.resolvers import metadata_resolver, sample_embedding_resolver
+from lightly_studio.errors import TagNotFoundError
+from lightly_studio.resolvers import metadata_resolver, sample_embedding_resolver, tag_resolver
 from lightly_studio.resolvers.sample_resolver.sample_filter import SampleFilter
 
 
 def compute_similarity_metadata(
     session: Session,
     key_dataset_id: UUID,
     embedding_model_id: UUID,
-    query_tag_id: UUID,
-    metadata_name: str,
-) -> None:
+    query_tag_name: str,
+    metadata_name: Optional[str] = None,
+) -> str:
     """Computes similarity for each sample in the dataset from embeddings.
 
     Similarity is a measure of how similar a sample is to its nearest neighbor
@@ -31,26 +34,38 @@ def compute_similarity_metadata(
             The ID of the dataset the similarity is computed on.
         embedding_model_id:
             The ID of the embedding model to use for the computation.
-        query_tag_id:
-            The ID of the tag describing the query.
+        query_tag_name:
+            The name of the tag describing the query.
         metadata_name:
             The name of the metadata field to store the similarity values in.
-            Defaults to "similarity".
+
+    Raises:
+        TagNotFoundError if `query_tag_name` does not exist
+
+    Returns:
+        The name of the metadata storing the similarity values.
     """
     license_key = LIGHTLY_STUDIO_LICENSE_KEY
     if license_key is None:
         raise ValueError(
             "LIGHTLY_STUDIO_LICENSE_KEY environment variable is not set. "
             "Please set it to your LightlyStudio license key."
         )
+    query_tag = tag_resolver.get_by_name(
+        session=session,
+        tag_name=query_tag_name,
+        dataset_id=key_dataset_id,
+    )
+    if query_tag is None:
+        raise TagNotFoundError(f"Query tag {query_tag_name} not found")
 
     key_samples = sample_embedding_resolver.get_all_by_dataset_id(
         session=session, dataset_id=key_dataset_id, embedding_model_id=embedding_model_id
     )
     key_embeddings = [sample.embedding for sample in key_samples]
     similarity = Similarity(key_embeddings=key_embeddings, token=license_key)
 
-    tag_filter = SampleFilter(tag_ids=[query_tag_id])
+    tag_filter = SampleFilter(tag_ids=[query_tag.tag_id])
     query_samples = sample_embedding_resolver.get_all_by_dataset_id(
         session=session,
         dataset_id=key_dataset_id,
@@ -59,10 +74,14 @@ def compute_similarity_metadata(
     )
     query_embeddings = [sample.embedding for sample in query_samples]
     similarity_values = similarity.calculate_similarity(query_embeddings=query_embeddings)
+    if metadata_name is None:
+        date = datetime.now(timezone.utc)
+        metadata_name = f"similarity_{query_tag_name}_{date.isoformat()}"
 
     metadata = [
         (sample.sample_id, {metadata_name: similarity})
         for sample, similarity in zip(key_samples, similarity_values)
     ]
 
     metadata_resolver.bulk_update_metadata(session, metadata)
+    return metadata_name

lightly_studio/tests/api/routes/api/test_metadata.py

@@ -10,8 +10,8 @@
     HTTP_STATUS_OK,
 )
 from lightly_studio.models.metadata import MetadataInfoView
-from lightly_studio.resolvers import image_resolver, metadata_resolver
-from tests import helpers_resolvers
+from lightly_studio.resolvers import image_resolver, metadata_resolver, tag_resolver
+from tests.helpers_resolvers import create_tag, fill_db_with_samples_and_embeddings
 
 
 def test_get_metadata_info(test_client: TestClient, mocker: MockerFixture) -> None:
@@ -66,7 +66,7 @@ def test_get_metadata_info__empty_response(test_client: TestClient, mocker: Mock
 def test_compute_typicality_metadata(test_client: TestClient, db_session: Session) -> None:
     """Test compute typicality metadata endpoint."""
     # Create dataset with samples and embeddings
-    dataset_id = helpers_resolvers.fill_db_with_samples_and_embeddings(
+    dataset_id = fill_db_with_samples_and_embeddings(
         test_db=db_session, n_samples=10, embedding_model_names=["test_embedding_model"]
     )
 
@@ -90,3 +90,40 @@ def test_compute_typicality_metadata(test_client: TestClient, db_session: Sessio
         )
         assert typicality_value is not None
         assert isinstance(typicality_value, float)
+
+
+def test_compute_similarity_metadata(test_client: TestClient, db_session: Session) -> None:
+    """Test compute similarity metadata endpoint."""
+    dataset_id = fill_db_with_samples_and_embeddings(
+        test_db=db_session, n_samples=10, embedding_model_names=["test_embedding_model"]
+    )
+    query_tag = create_tag(session=db_session, dataset_id=dataset_id, tag_name="query_tag")
+    samples = image_resolver.get_all_by_dataset_id(
+        session=db_session, dataset_id=dataset_id
+    ).samples
+    tag_resolver.add_sample_ids_to_tag_id(
+        session=db_session,
+        tag_id=query_tag.tag_id,
+        sample_ids=[samples[0].sample_id, samples[2].sample_id],
+    )
+
+    response = test_client.post(
+        f"/api/datasets/{dataset_id}/metadata/similarity", json={"query_tag_name": "query_tag"}
+    )
+
+    assert response.status_code == 200
+    metadata_name = response.text[1:-1]  # We strip the double-quotes
+    assert metadata_name.startswith("similarity_query_tag_20")
+
+    samples = image_resolver.get_all_by_dataset_id(
+        session=db_session, dataset_id=dataset_id
+    ).samples
+    assert len(samples) == 10
+
+    # Verify all samples have similarity metadata.
+    for sample in samples:
+        similarity_value = metadata_resolver.get_value_for_sample(
+            session=db_session, sample_id=sample.sample_id, key=metadata_name
+        )
+        assert similarity_value is not None
+        assert isinstance(similarity_value, float)

lightly_studio/tests/metadata/test_compute_similarity.py

@@ -37,18 +37,17 @@ def test_compute_similarity_metadata(test_db: Session) -> None:
     )
 
     query_tag = create_tag(session=test_db, dataset_id=dataset_id, tag_name="query_tag")
-    query_tag_id = query_tag.tag_id
     tag_resolver.add_sample_ids_to_tag_id(
         session=test_db,
-        tag_id=query_tag_id,
+        tag_id=query_tag.tag_id,
         sample_ids=[samples[0].sample_id, samples[2].sample_id],
     )
 
     compute_similarity.compute_similarity_metadata(
         session=test_db,
         key_dataset_id=dataset_id,
         embedding_model_id=embedding_model_id,
-        query_tag_id=query_tag_id,
+        query_tag_name="query_tag",
         metadata_name="similarity",
     )