add support for JSONPath to entity type field extraction (Netflix#2994)

* add support for JSONPath to entity type field extraction

* delete old function

* add custom type defintion for entity type pari

* add custom type defintion for entity type pari

* improve custom type

* compile dependencies and remove unused import

* Update description of field on entitytype edit sheet component

* remove duplicate entities, parse valid jsonpath, and fix uuid bug

* Update src/dispatch/signal/service.py

Author: wssheldon

requirements-base.in

@@ -18,6 +18,7 @@ jinja2
 jira==2.0.0
 atlassian-python-api==3.32.0
 joblib
+jsonpath_ng
 numpy
 oauth2client
 pandas

requirements-base.txt

@@ -90,7 +90,9 @@ cymem==2.0.7
     #   spacy
     #   thinc
 decorator==5.1.1
-    # via validators
+    # via
+    #   jsonpath-ng
+    #   validators
 defusedxml==0.7.1
     # via jira
 deprecated==1.2.13
@@ -103,6 +105,10 @@ email-validator==1.3.1
     # via -r requirements-base.in
 emails==0.6
     # via -r requirements-base.in
+exceptiongroup==1.1.0
+    # via
+    #   hypothesis
+    #   pytest
 fastapi==0.90.1
     # via -r requirements-base.in
 frozenlist==1.3.3
@@ -169,6 +175,8 @@ jira==2.0.0
     # via -r requirements-base.in
 joblib==1.2.0
     # via -r requirements-base.in
+jsonpath-ng==1.5.3
+    # via -r requirements-base.in
 jsonschema==4.17.3
     # via
     #   hypothesis-jsonschema
@@ -233,6 +241,8 @@ pdpyras==4.5.2
     # via -r requirements-base.in
 pluggy==1.0.0
     # via pytest
+ply==3.11
+    # via jsonpath-ng
 premailer==3.10.0
     # via emails
 preshed==3.0.8
@@ -347,6 +357,7 @@ six==1.16.0
     #   google-auth
     #   google-auth-httplib2
     #   jira
+    #   jsonpath-ng
     #   junit-xml
     #   oauth2client
     #   patsy
@@ -410,7 +421,9 @@ text-unidecode==1.3
 thinc==8.1.5
     # via spacy
 tomli==2.0.1
-    # via schemathesis
+    # via
+    #   pytest
+    #   schemathesis
 tomli-w==1.0.0
     # via schemathesis
 tqdm==4.64.1
@@ -423,6 +436,7 @@ typing-extensions==4.4.0
     # via
     #   pydantic
     #   schemathesis
+    #   starlette
 uritemplate==4.1.1
     # via google-api-python-client
 urllib3==1.26.13

src/dispatch/entity/models.py

@@ -51,6 +51,9 @@ class EntityBase(DispatchBase):
 
 
 class EntityCreate(EntityBase):
+    def __hash__(self):
+        return hash((self.id, self.value))
+
     id: Optional[PrimaryKey]
     entity_type: EntityTypeCreate
     project: ProjectRead

src/dispatch/entity/service.py

@@ -1,7 +1,8 @@
 from datetime import datetime, timedelta
-from typing import Optional, Sequence
+from typing import Generator, Optional, Sequence, Union, NewType, NamedTuple
 import re
 
+import jsonpath_ng
 from pydantic.error_wrappers import ErrorWrapper, ValidationError
 from sqlalchemy.orm import Session, joinedload
 
@@ -163,114 +164,177 @@ def get_signal_instances_with_entity(
     return signal_instances
 
 
+EntityTypePair = NewType(
+    "EntityTypePair",
+    NamedTuple(
+        "EntityTypePairTuple",
+        [
+            ("entity_type", EntityType),
+            ("regex", Union[re.Pattern[str], None]),
+            ("json_path", Union[jsonpath_ng.JSONPath, None]),
+        ],
+    ),
+)
+
+
 def find_entities(
     db_session: Session, signal_instance: SignalInstance, entity_types: Sequence[EntityType]
 ) -> list[Entity]:
-    """Find entities of the given types in the raw data of a signal instance.
+    """
+    Find entities in a SignalInstance based on a list of EntityTypes.
 
     Args:
-        db_session (Session): SQLAlchemy database session.
-        signal_instance (SignalInstance): SignalInstance to search for entities in.
-        entity_types (list[EntityType]): List of EntityType objects to search for.
+        db_session (Session): The database session to use for entity creation.
+        signal_instance (SignalInstance): The SignalInstance to extract entities from.
+        entity_types (Sequence[EntityType]): A list of EntityTypes to search for in the SignalInstance.
 
     Returns:
-        list[Entity]: List of Entity objects found.
-
-    Example:
-        >>> signal_instance = SignalInstance(
-        ...     raw={
-        ...         "name": "John Smith",
-        ...         "email": "[email protected]",
-        ...         "phone": "555-555-1212",
-        ...         "address": {
-        ...             "street": "123 Main St",
-        ...             "city": "Anytown",
-        ...             "state": "CA",
-        ...             "zip": "12345"
-        ...         },
-        ...         "notes": "Customer is interested in buying a product."
-        ...     }
-        ... )
-        >>> entity_types = [
-        ...     EntityType(name="Name", field="name", regular_expression=r"\b[A-Z][a-z]+ [A-Z][a-z]+\b"),
-        ...     EntityType(name="Phone", field=None, regular_expression=r"\b\\d{3}[-.]?\\d{3}[-.]?\\d{4}\b"),
-        ...     EntityType(name="Street", field="address.street"),
-        ... ]
-        >>> entities = find_entities(db_session, signal_instance, entity_types)
-
-    Notes:
-        This function uses depth-first search to traverse the raw data of the signal instance. It searches for
-        the regular expressions specified in the EntityType objects in the values of the dictionary, list, and
-        string objects encountered during the traversal. The search can be limited to a specific key in the
-        dictionary objects by specifying a value for the 'field' attribute of the EntityType object.
+        list[Entity]: A list of entities found in the SignalInstance.
     """
 
-    def _search(key, val, entity_type_pairs):
-        # Create a list to hold any entities that are found in this value
-        entities = []
-
-        # If this value has been searched before, return the cached entities
-        if id(val) in cache:
-            return cache[id(val)]
-
+    def _find_entites_by_regex(
+        val: Union[dict, str, list],
+        signal_instance: SignalInstance,
+        entity_type_pairs: list[EntityTypePair],
+    ) -> Generator[EntityCreate, None, None]:
+        """
+        Find entities in a value using regular expressions.
+
+        Args:
+            val: The value to search for entities in.
+            signal_instance (SignalInstance): The SignalInstance being processed.
+            entity_type_pairs (list): A list of (entity_type, entity_regex, field) tuples to search for.
+
+        Yields:
+            EntityCreate: An entity found in the value.
+
+        Examples:
+            >>> entity_type_pairs = [
+            ...     (
+            ...         EntityType("PERSON", r"([A-Z][a-z]+)+"),
+            ...         re.compile(r"([A-Z][a-z]+)+"),
+            ...         None
+            ...     ),
+            ...     (
+            ...         EntityType("DATE", r"(\d{4}(-\d{2}){2}|\d{4}\/\d{2}\/\d{2})"), # noqa
+            ...         re.compile(r"(\d{4}(-\d{2}){2}|\\d{4}\/\d{2}\/\d{2})"), # noqa
+            ...         None
+            ...     )
+            ... ]
+
+            >>> signal_instance = SignalInstance(raw={"text": "John Doe was born on 1987-05-12."})
+
+            >>> entities = list(_find_entites_by_regex(signal_instance.raw, signal_instance, entity_type_pairs))
+
+            >>> entities[0].value
+            'John Doe'
+            >>> entities[0].entity_type.name
+            'PERSON'
+            >>> entities[1].value
+            '1987-05-12'
+            >>> entities[1].entity_type.name
+            'DATE'
+        """
         # If the value is a dictionary, search its key-value pairs recursively
         if isinstance(val, dict):
-            for subkey, subval in val.items():
-                entities.extend(_search(subkey, subval, entity_type_pairs))
+            for _, subval in val.items():
+                yield from _find_entites_by_regex(
+                    subval,
+                    signal_instance,
+                    entity_type_pairs,
+                )
 
         # If the value is a list, search its items recursively
         elif isinstance(val, list):
             for item in val:
-                entities.extend(_search(None, item, entity_type_pairs))
+                yield from _find_entites_by_regex(
+                    item,
+                    signal_instance,
+                    entity_type_pairs,
+                )
 
         # If the value is a string, search it for entity matches
         elif isinstance(val, str):
-            for entity_type, entity_regex, field in entity_type_pairs:
-                # If a field was specified for this entity type, only search that field
-                if not field or key == field:
-                    if entity_regex is None:
-                        # If no regular expression was specified, return the value of the field/key
-                        entity = EntityCreate(
-                            value=val,
-                            entity_type=entity_type,
-                            project=signal_instance.project,
-                        )
-                        entities.append(entity)
-                    else:
-                        # Search the string for matches to the entity type's regular expression
-                        if match := entity_regex.search(val):
-                            entity = EntityCreate(
-                                value=match.group(0),
-                                entity_type=entity_type,
-                                project=signal_instance.project,
-                            )
-                            entities.append(entity)
-
-        # Cache the entities found for this value
-        cache[id(val)] = entities
-        return entities
+            for entity_type, entity_regex, _ in entity_type_pairs:
+                # Search the string for matches to the entity type's regular expression
+                if match := entity_regex.search(val):
+                    yield EntityCreate(
+                        value=match.group(0),
+                        entity_type=entity_type,
+                        project=signal_instance.project,
+                    )
+
+    def _find_entities_by_regex_and_jsonpath_expression(
+        signal_instance: SignalInstance,
+        entity_type_pairs: list[EntityTypePair],
+    ) -> Generator[EntityCreate, None, None]:
+        """
+        Yield entities found in a SignalInstance by searching its fields using regular expressions and JSONPath expressions.
+
+        Args:
+            signal_instance: The SignalInstance to extract entities from.
+            entity_type_pairs: A list of (entity_type, entity_regex, field) tuples to search for.
+
+        Yields:
+            EntityCreate: An entity found in the SignalInstance.
+        """
+        for entity_type, entity_regex, field in entity_type_pairs:
+            if field:
+                try:
+                    matches = field.find(signal_instance.raw)
+                    for match in matches:
+                        if isinstance(match.value, str):
+                            if entity_regex is None:
+                                yield EntityCreate(
+                                    value=match.value,
+                                    entity_type=entity_type,
+                                    project=signal_instance.project,
+                                )
+                            else:
+                                if match := entity_regex.search(match.value):
+                                    yield EntityCreate(
+                                        value=match.group(0),
+                                        entity_type=entity_type,
+                                        project=signal_instance.project,
+                                    )
+                except jsonpath_ng.PathNotFound:
+                    # field not found in signal_instance.raw
+                    pass
 
     # Create a list of (entity type, regular expression, field) tuples
     entity_type_pairs = [
-        (type, re.compile(type.regular_expression) if type.regular_expression else None, type.field)
+        (
+            type,
+            re.compile(type.regular_expression) if type.regular_expression else None,
+            jsonpath_ng.parse(type.field) if type.field else None,
+        )
         for type in entity_types
         if isinstance(type.regular_expression, str) or type.field is not None
     ]
 
-    # Initialize a cache of previously searched values
-    cache = {}
+    # Filter the entity type pairs based on the field
+    filtered_entity_type_pairs = [
+        (entity_type, entity_regex, field)
+        for entity_type, entity_regex, field in entity_type_pairs
+        if not field
+    ]
 
-    # Traverse the signal data using depth-first search
+    # Use the recursive search function to find entities in the raw data
     entities = [
         entity
-        for key, val in signal_instance.raw.items()
-        for entity in _search(key, val, entity_type_pairs)
+        for _, val in signal_instance.raw.items()
+        for entity in _find_entites_by_regex(val, signal_instance, filtered_entity_type_pairs)
     ]
 
-    # Create the entities in the database and add them to the signal instance
+    entities.extend(
+        _find_entities_by_regex_and_jsonpath_expression(signal_instance, entity_type_pairs)
+    )
+
+    # Filter out duplicate entities
+    entities = list(set(entities))
+
     entities_out = [
         get_by_value_or_create(db_session=db_session, entity_in=entity_in) for entity_in in entities
     ]
 
-    # Return the list of entities found
     return entities_out

src/dispatch/signal/service.py

@@ -1,4 +1,5 @@
 import json
+import uuid
 import hashlib
 from typing import Optional