Fix Freebase entity prefix validation

src/server/tasks/knowledgegraph/api.py

@@ -57,7 +57,7 @@ def get_relations(self, variable: Union[Variable, str]):
         Returns (None, "Observation: [...]").
         """
         if not isinstance(variable, Variable):
-            if not re.match(r'^([mf])\.[\w_]+$', variable):
+            if not re.match(r'^[mg]\.[\w_]+$', variable):
                 raise ValueError("get_relations: variable must be a variable or an entity")
 
         cache_key = (self.task_id, variable if isinstance(variable, str) else hash(variable))
@@ -92,7 +92,7 @@ def get_neighbors(self, variable: Union[Variable, str], relation: str):
         Get neighbors via a relation. Returns (new_variable, "Observation: ...").
         """
         if not isinstance(variable, Variable):
-            if not re.match(r'^([mf])\.[\w_]+$', variable):
+            if not re.match(r'^[mg]\.[\w_]+$', variable):
                 raise ValueError("get_neighbors: variable must be a variable or an entity")
 
         cache_key = (self.task_id, variable if isinstance(variable, str) else hash(variable))

tests/test_knowledgegraph_entity_ids.py

@@ -0,0 +1,48 @@
+import unittest
+
+from src.server.tasks.knowledgegraph import api as knowledgegraph_api
+
+
+class StubSparqlExecutor:
+    def __init__(self):
+        self.entities = []
+
+    def get_out_relations(self, entity):
+        self.entities.append(entity)
+        return []
+
+
+class KnowledgeGraphEntityIdTest(unittest.TestCase):
+    def setUp(self):
+        knowledgegraph_api.relation_cache.clear()
+        knowledgegraph_api.variable_relations_cache.clear()
+        self.executor = StubSparqlExecutor()
+        self.api = knowledgegraph_api.API(self.executor)
+
+    def test_get_relations_accepts_g_prefixed_entity(self):
+        entity = "g.11b5lzm6b0"
+
+        self.api.get_relations(entity)
+
+        self.assertEqual(self.executor.entities, [entity])
+
+    def test_get_neighbors_accepts_g_prefixed_entity(self):
+        entity = "g.11b5lzm6b0"
+        relation = "test.relation"
+        knowledgegraph_api.variable_relations_cache[entity] = [relation]
+        knowledgegraph_api.range_info[relation] = "test.type"
+        self.addCleanup(knowledgegraph_api.range_info.pop, relation)
+
+        variable, _ = self.api.get_neighbors(entity, relation)
+
+        self.assertEqual(variable.program, f"(JOIN {relation}_inv {entity})")
+
+    def test_f_prefixed_value_is_rejected(self):
+        with self.assertRaises(ValueError):
+            self.api.get_relations("f.not_an_entity")
+        with self.assertRaises(ValueError):
+            self.api.get_neighbors("f.not_an_entity", "test.relation")
+
+
+if __name__ == "__main__":
+    unittest.main()