Merge pull request #130 from BioPack-team/constraints-autogen

Disable attribute constraints generation during `transpiling` phase for now

Author: tokebe

src/retriever/data_tiers/tier_1/elasticsearch/transpiler.py

@@ -137,8 +137,53 @@ def process_qedge(self, qedge: QEdgeDict) -> list[ESFilterClause]:
             if (values := qedge.get(qfield))
         ]
 
-    def generate_query_for_merged_edges(
-        self, in_node: QNodeDict, edge: QEdgeDict, out_node: QNodeDict
+    def generate_attribute_constraints(
+        self,
+        in_node: QNodeDict,
+        edge: QEdgeDict,
+        out_node: QNodeDict,
+        query_kwargs: ESBooleanQuery,
+    ) -> ESBooleanQuery:
+        """Generate attribute constraints based on QNode/QEdge payload."""
+        constraint_origins: list[AttributeOrigin] = ["edge", "subject", "object"]
+
+        all_must: list[AttributeFilterQuery] = []
+        all_must_not: list[AttributeFilterQuery] = []
+
+        for origin in constraint_origins:
+            entity = (
+                edge
+                if origin == "edge"
+                else in_node
+                if origin == "subject"
+                else out_node
+            )
+
+            if origin == "edge":
+                constraints = entity.get("attribute_constraints", None)
+            else:
+                constraints = entity.get("constraints", None)
+
+            if constraints:
+                must, must_not = process_attribute_constraints(constraints, origin)
+                if must:
+                    all_must.extend(must)
+                if must_not:
+                    all_must_not.extend(must_not)
+
+        if all_must:
+            query_kwargs["must"] = all_must
+        if all_must_not:
+            query_kwargs["must_not"] = all_must_not
+
+        return query_kwargs
+
+    def generate_queries(
+        self,
+        in_node: QNodeDict,
+        edge: QEdgeDict,
+        out_node: QNodeDict,
+        gen_attribute_constraints: bool = False,  # disable attribute constraints for now
     ) -> ESPayload:
         """Generate query based on merged edges schema on Elasticsearch.
 
@@ -185,37 +230,12 @@ def generate_query_for_merged_edges(
                 query_kwargs["filter"].append(qualifier_terms)
 
         # generate constraint terms for edges and associated nodes
-        constraint_origins: list[AttributeOrigin] = ["edge", "subject", "object"]
-
-        all_must: list[AttributeFilterQuery] = []
-        all_must_not: list[AttributeFilterQuery] = []
-
-        for origin in constraint_origins:
-            entity = (
-                edge
-                if origin == "edge"
-                else in_node
-                if origin == "subject"
-                else out_node
+        # currently, this is DISABLED by default to favor post-processing
+        if gen_attribute_constraints:
+            query_kwargs = self.generate_attribute_constraints(
+                in_node, edge, out_node, query_kwargs
             )
 
-            if origin == "edge":
-                constraints = entity.get("attribute_constraints", None)
-            else:
-                constraints = entity.get("constraints", None)
-
-            if constraints:
-                must, must_not = process_attribute_constraints(constraints, origin)
-                if must:
-                    all_must.extend(must)
-                if must_not:
-                    all_must_not.extend(must_not)
-
-        if all_must:
-            query_kwargs["must"] = all_must
-        if all_must_not:
-            query_kwargs["must_not"] = all_must_not
-
         return ESPayload(query=ESQueryContext(bool=ESBooleanQuery(**query_kwargs)))
 
     @override
@@ -226,7 +246,7 @@ def convert_triple(self, qgraph: QueryGraphDict) -> ESPayload:
             raise ValueError("Query graph must contain exactly one edge.")
         in_node = qgraph["nodes"][edge["subject"]]
         out_node = qgraph["nodes"][edge["object"]]
-        return self.generate_query_for_merged_edges(in_node, edge, out_node)
+        return self.generate_queries(in_node, edge, out_node)
 
     @override
     def convert_batch_triple(self, qgraphs: list[QueryGraphDict]) -> list[ESPayload]:

tests/data_tiers/tier_1/elasticsearch_tests/test_tier1_driver.py

@@ -9,6 +9,7 @@
 from retriever.data_tiers.tier_1.elasticsearch.types import ESPayload, ESEdge
 from payload.trapi_qgraphs import DINGO_QGRAPH, VALID_REGEX_QGRAPHS, INVALID_REGEX_QGRAPHS, ID_BYPASS_PAYLOAD
 from retriever.utils.redis import REDIS_CLIENT
+from test_tier1_transpiler import _convert_triple, _convert_batch_triple
 
 
 def esp(d: dict[str, Any]) -> ESPayload:
@@ -148,15 +149,15 @@ def assert_single_result(res, expected_result_num: int):
 def test_invalid_regex_qgraph(qgraph):
     transpiler = ElasticsearchTranspiler()
     with pytest.raises(ValueError):
-        transpiler.convert_triple(qgraph)
+        _convert_triple(transpiler, qgraph)
 
 
 @pytest.mark.usefixtures("mock_elasticsearch_config")
 @pytest.mark.asyncio
 async def test_valid_regex_query():
     transpiler = ElasticsearchTranspiler()
 
-    qgraphs_with_valid_regex = transpiler.convert_batch_triple(VALID_REGEX_QGRAPHS)
+    qgraphs_with_valid_regex = _convert_batch_triple(transpiler, VALID_REGEX_QGRAPHS)
 
     driver: driver_mod.ElasticSearchDriver = driver_mod.ElasticSearchDriver()
 
@@ -211,7 +212,7 @@ async def test_metadata_retrieval():
 )
 async def test_end_to_end(qgraph, expected_hits):
     transpiler = ElasticsearchTranspiler()
-    payload = transpiler.convert_triple(qgraph)
+    payload = _convert_triple(transpiler, qgraph)
 
     driver: driver_mod.ElasticSearchDriver = driver_mod.ElasticSearchDriver()
 

tests/data_tiers/tier_1/elasticsearch_tests/test_tier1_transpiler.py

@@ -33,6 +33,26 @@ def es_transpiler() -> ElasticsearchTranspiler:
     return ElasticsearchTranspiler()
 
 
+def _convert_triple(
+    es_transpiler: ElasticsearchTranspiler,
+    qgraph: QueryGraphDict,
+) -> ESPayload:
+    """Wrapper for convert_triple so attribute constraints are generated by default."""
+    edge = next(iter(qgraph["edges"].values()), None)
+    if edge is None:
+        raise ValueError("Query graph must contain exactly one edge.")
+    in_node = qgraph["nodes"][edge["subject"]]
+    out_node = qgraph["nodes"][edge["object"]]
+    return es_transpiler.generate_queries(in_node, edge, out_node, gen_attribute_constraints=True)
+
+
+def _convert_batch_triple(
+    es_transpiler: ElasticsearchTranspiler,
+    qgraphs: list[QueryGraphDict]
+) -> list[ESPayload]:
+    return [_convert_triple(es_transpiler, qgraph) for qgraph in qgraphs]
+
+
 def check_list_fields(reference: list, against: list):
     for ref, ag in zip(reference, against):
         if ref.startswith("biolink:"):
@@ -205,7 +225,9 @@ def add_to_required_fields(_field: str, _side: side_type):
 def test_bypass_payload(
         es_transpiler: ElasticsearchTranspiler
 ) -> None:
-    generated_payload = es_transpiler.convert_triple(ID_BYPASS_PAYLOAD)
+    # generated_payload = es_transpiler.convert_triple(ID_BYPASS_PAYLOAD)
+
+    generated_payload = _convert_triple(es_transpiler, ID_BYPASS_PAYLOAD)
     check_single_query_payload(ID_BYPASS_PAYLOAD, generated_payload)
     must_clauses = generated_payload["query"]["bool"]["filter"]
 
@@ -225,7 +247,8 @@ def test_bypass_payload(
 def test_convert_triple(
     q_graph: QueryGraphDict, es_transpiler: ElasticsearchTranspiler
 ) -> None:
-    generated_payload = es_transpiler.convert_triple(q_graph)
+    # generated_payload = es_transpiler.convert_triple(q_graph)
+    generated_payload = _convert_triple(es_transpiler, q_graph)
     check_single_query_payload(q_graph, generated_payload)
 
 
@@ -235,7 +258,8 @@ def test_convert_batch_triple(
 ) -> None:
     batch_q_graphs = [q_graph for i in range(10)]
 
-    generated_payload_list = es_transpiler.convert_batch_triple(batch_q_graphs)
+    # generated_payload_list = es_transpiler.convert_batch_triple(batch_q_graphs)
+    generated_payload_list = _convert_batch_triple(es_transpiler, batch_q_graphs)
     for generated_payload in generated_payload_list:
         check_single_query_payload(q_graph, generated_payload)