Add filter for topics to match API

opensanctions · Nov 24, 2023 · 4fd0a67 · 4fd0a67
1 parent 3bea378
commit 4fd0a67
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git a/tests/unit/test_match.py b/tests/unit/test_match.py
@@ -102,6 +102,16 @@ def test_match_exclude_dataset():
     assert len(res["results"]) == 0, res
 
 
+def test_filter_topic():
+    query = {"queries": {"vv": EXAMPLE}}
+    params = {"algorithm": "name-based", "topics": "crime.cyber"}
+    resp = client.post("/match/default", json=query, params=params)
+    assert resp.status_code == 200, resp.text
+    data = resp.json()
+    res = data["responses"]["vv"]
+    assert len(res["results"]) == 0, res
+
+
 def test_id_pass_through():
     body = dict(ERMAKOV)
     body["id"] = "ermakov"

diff --git a/yente/routers/match.py b/yente/routers/match.py
@@ -7,7 +7,7 @@
 from yente.data.common import ErrorResponse
 from yente.data.common import EntityMatchQuery, EntityMatchResponse, EntityExample
 from yente.data.common import EntityMatches
-from yente.search.queries import entity_query
+from yente.search.queries import entity_query, FilterDict
 from yente.search.search import search_entities, result_entities, result_total
 from yente.data.entity import Entity
 from yente.util import limit_window
@@ -53,6 +53,9 @@ async def match(
     exclude_dataset: List[str] = Query(
         [], title="Remove the given datasets from results"
     ),
+    topics: List[str] = Query(
+        [], title="Only return results that match the given topics"
+    ),
     fuzzy: bool = Query(
         settings.MATCH_FUZZY,
         title="Use slow matching for candidate generation, does not affect scores",
@@ -126,6 +129,7 @@ async def match(
         msg = "Too many queries in one batch (limit: %d)" % settings.MAX_BATCH
         raise HTTPException(400, detail=msg)
 
+    filters: FilterDict = {"topics": topics}
     queries = []
     entities = []
     responses: Dict[str, EntityMatches] = {}
@@ -138,6 +142,7 @@ async def match(
             query = entity_query(
                 ds,
                 entity,
+                filters=filters,
                 fuzzy=fuzzy,
                 exclude_schema=exclude_schema,
                 exclude_dataset=exclude_dataset,

diff --git a/yente/search/queries.py b/yente/search/queries.py
@@ -85,6 +85,7 @@ def names_query(entity: EntityProxy, fuzzy: bool = True) -> List[Clause]:
 def entity_query(
     dataset: Dataset,
     entity: EntityProxy,
+    filters: FilterDict = {},
     fuzzy: bool = True,
     exclude_schema: List[str] = [],
     exclude_dataset: List[str] = [],
@@ -104,6 +105,7 @@ def entity_query(
 
     return filter_query(
         shoulds,
+        filters=filters,
         dataset=dataset,
         schema=entity.schema,
         exclude_schema=exclude_schema,