[hma] Patch similarity score (#1641)

hasher-matcher-actioner/src/OpenMediaMatch/blueprints/development.py

@@ -10,7 +10,10 @@
 from werkzeug.exceptions import HTTPException
 
 from OpenMediaMatch.blueprints.hashing import hash_media
-from OpenMediaMatch.blueprints.matching import lookup_signal
+from OpenMediaMatch.blueprints.matching import (
+    lookup_signal,
+    lookup_signal_with_distance,
+)
 from OpenMediaMatch.utils.flask_utils import api_error_handler
 
 from OpenMediaMatch.utils import dev_utils
@@ -50,6 +53,11 @@ def query_media():
             return signal_type_to_signal_map
         abort(500, "Something went wrong while hashing the provided media.")
 
+    include_distance = bool(request.args.get("include_distance", False)) == True
+    lookup_signal_func = (
+        lookup_signal_with_distance if include_distance else lookup_signal
+    )
+
     # Check if signal_type is an option in the map of hashes
     signal_type_name = request.args.get("signal_type")
     if signal_type_name is not None:
@@ -59,14 +67,14 @@ def query_media():
                 f"Requested signal type '{signal_type_name}' is not supported for the provided "
                 "media.",
             )
-        return lookup_signal(
+        return lookup_signal_func(
             signal_type_to_signal_map[signal_type_name], signal_type_name
         )
     return {
         "matches": list(
             itertools.chain(
                 *map(
-                    lambda x: lookup_signal(x[1], x[0])["matches"],
+                    lambda x: lookup_signal_func(x[1], x[0])["matches"],
                     signal_type_to_signal_map.items(),
                 ),
             )

hasher-matcher-actioner/src/OpenMediaMatch/blueprints/matching.py

@@ -7,6 +7,7 @@
 from dataclasses import dataclass
 import datetime
 import random
+import sys
 import typing as t
 import time
 
@@ -15,18 +16,31 @@
 from werkzeug.exceptions import HTTPException
 
 from threatexchange.signal_type.signal_base import SignalType
-from threatexchange.signal_type.index import SignalTypeIndex
+from threatexchange.signal_type.index import (
+    IndexMatchUntyped,
+    SignalSimilarityInfo,
+    SignalTypeIndex,
+)
 
 from OpenMediaMatch.background_tasks.development import get_apscheduler
 from OpenMediaMatch.storage import interface
 from OpenMediaMatch.blueprints import hashing
-from OpenMediaMatch.utils.flask_utils import require_request_param, api_error_handler
+from OpenMediaMatch.utils.flask_utils import (
+    api_error_handler,
+    require_request_param,
+    str_to_bool,
+)
 from OpenMediaMatch.persistence import get_storage
 
 bp = Blueprint("matching", __name__)
 bp.register_error_handler(HTTPException, api_error_handler)
 
 
+class MatchWithDistance(t.TypedDict):
+    content_id: int
+    distance: str
+
+
 @dataclass
 class _SignalIndexInMemoryCache:
     signal_type: t.Type[SignalType]
@@ -94,15 +108,23 @@ def raw_lookup():
      * Signal type (hash type)
      * Signal value (the hash)
      * Optional list of banks to restrict search to
+     * Optional include_distance (bool) wether or not to return distance values on match
     Output:
-     * List of matching content items
+     * List of matching with content_id and, if included, distance values
     """
     signal = require_request_param("signal")
     signal_type_name = require_request_param("signal_type")
-    return lookup_signal(signal, signal_type_name)
+    include_distance = str_to_bool(request.args.get("include_distance", "false"))
+    lookup_signal_func = (
+        lookup_signal_with_distance if include_distance else lookup_signal
+    )
 
+    return {"matches": lookup_signal_func(signal, signal_type_name)}
 
-def lookup_signal(signal: str, signal_type_name: str) -> dict[str, list[int]]:
+
+def query_index(
+    signal: str, signal_type_name: str
+) -> t.Sequence[IndexMatchUntyped[SignalSimilarityInfo, int]]:
     storage = get_storage()
     signal_type = _validate_and_transform_signal_type(signal_type_name, storage)
 
@@ -118,7 +140,25 @@ def lookup_signal(signal: str, signal_type_name: str) -> dict[str, list[int]]:
     current_app.logger.debug("[lookup_signal] querying index")
     results = index.query(signal)
     current_app.logger.debug("[lookup_signal] query complete")
-    return {"matches": [m.metadata for m in results]}
+    return results
+
+
+def lookup_signal(signal: str, signal_type_name: str) -> list[int]:
+    results = query_index(signal, signal_type_name)
+    return [m.metadata for m in results]
+
+
+def lookup_signal_with_distance(
+    signal: str, signal_type_name: str
+) -> list[MatchWithDistance]:
+    results = query_index(signal, signal_type_name)
+    return [
+        {
+            "content_id": m.metadata,
+            "distance": m.similarity_info.pretty_str(),
+        }
+        for m in results
+    ]
 
 
 def _validate_and_transform_signal_type(
@@ -300,6 +340,7 @@ def index_cache_is_stale() -> bool:
 
 def _get_index(signal_type: t.Type[SignalType]) -> SignalTypeIndex[int] | None:
     entry = _get_index_cache().get(signal_type.get_name())
+
     if entry is None:
         current_app.logger.debug("[lookup_signal] no cache, loading index")
         return get_storage().get_signal_type_index(signal_type)

hasher-matcher-actioner/src/OpenMediaMatch/tests/test_e2e_workflow.py

@@ -13,7 +13,7 @@
 from threatexchange.signal_type.pdq.signal import PdqSignal
 
 
-def test_raw_hash_add_to_match(app: Flask, client: FlaskClient):
+def test_raw_hash_add_to_match_no_distance(app: Flask, client: FlaskClient):
     bank_name = "TEST_BANK"
     create_bank(client, bank_name)
 
@@ -40,3 +40,34 @@ def test_raw_hash_add_to_match(app: Flask, client: FlaskClient):
     resp = client.get(f"/m/raw_lookup?signal_type=pdq&signal={hashes[-1]}")
     assert resp.status_code == 200
     assert resp.json == {"matches": [16]}
+
+
+def test_raw_hash_add_to_match_with_distance(app: Flask, client: FlaskClient):
+    bank_name = "TEST_BANK"
+    create_bank(client, bank_name)
+
+    # PDQ hashes
+    hashes = PdqSignal.get_examples()
+    for pdq in hashes:
+        resp = client.post(f"/c/bank/{bank_name}/signal", json={"pdq": pdq})
+        assert resp.status_code == 200
+
+    # No background tasks in the test, so let's trigger it manually
+    storage = get_storage()
+    build_all_indices(storage, storage, storage)
+
+    # Sanity check that the index is build
+    resp = client.get(f"/m/index/status?signal_type=pdq")
+    assert resp.status_code == 200
+    all_build_info = resp.json
+    assert "pdq" in all_build_info  # type: ignore
+    build_info = all_build_info["pdq"]  # type: ignore
+    assert build_info["present"] == True
+    assert build_info["size"] == len(hashes)
+
+    # Now match
+    resp = client.get(
+        f"/m/raw_lookup?signal_type=pdq&include_distance=true&signal={hashes[-1]}"
+    )
+    assert resp.status_code == 200
+    assert resp.json == {"matches": [{"content_id": 16, "distance": "0"}]}