graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/search/search_filters.py

@@ -19,7 +19,8 @@ from enum import Enum
 from typing import Any
 
 from pydantic import BaseModel, Field
-from typing_extensions import LiteralString
+
+from graphiti_core.driver.driver import GraphProvider
 
 
 class ComparisonOperator(Enum):
@@ -29,143 +30,233 @@ class ComparisonOperator(Enum):
     less_than = '<'
     greater_than_equal = '>='
     less_than_equal = '<='
+    is_null = 'IS NULL'
+    is_not_null = 'IS NOT NULL'
 
 
 class DateFilter(BaseModel):
-    date: datetime = Field(description='A datetime to filter on')
+    date: datetime | None = Field(description='A datetime to filter on')
     comparison_operator: ComparisonOperator = Field(
         description='Comparison operator for date filter'
     )
 
 
+class PropertyFilter(BaseModel):
+    property_name: str = Field(description='Property name')
+    property_value: str | int | float | None = Field(
+        description='Value you want to match on for the property'
+    )
+    comparison_operator: ComparisonOperator = Field(
+        description='Comparison operator for the property'
+    )
+
+
 class SearchFilters(BaseModel):
     node_labels: list[str] | None = Field(
         default=None, description='List of node labels to filter on'
     )
+    edge_types: list[str] | None = Field(
+        default=None, description='List of edge types to filter on'
+    )
     valid_at: list[list[DateFilter]] | None = Field(default=None)
     invalid_at: list[list[DateFilter]] | None = Field(default=None)
     created_at: list[list[DateFilter]] | None = Field(default=None)
     expired_at: list[list[DateFilter]] | None = Field(default=None)
+    edge_uuids: list[str] | None = Field(default=None)
+    property_filters: list[PropertyFilter] | None = Field(default=None)
+
+
+def cypher_to_opensearch_operator(op: ComparisonOperator) -> str:
+    mapping = {
+        ComparisonOperator.greater_than: 'gt',
+        ComparisonOperator.less_than: 'lt',
+        ComparisonOperator.greater_than_equal: 'gte',
+        ComparisonOperator.less_than_equal: 'lte',
+    }
+    return mapping.get(op, op.value)
 
 
 def node_search_filter_query_constructor(
     filters: SearchFilters,
-) -> tuple[LiteralString, dict[str, Any]]:
-    filter_query: LiteralString = ''
+    provider: GraphProvider,
+) -> tuple[list[str], dict[str, Any]]:
+    filter_queries: list[str] = []
     filter_params: dict[str, Any] = {}
 
     if filters.node_labels is not None:
-        node_labels = '|'.join(filters.node_labels)
-        node_label_filter = ' AND n:' + node_labels
-        filter_query += node_label_filter
+        if provider == GraphProvider.KUZU:
+            node_label_filter = 'list_has_all(n.labels, $labels)'
+            filter_params['labels'] = filters.node_labels
+        else:
+            node_labels = '|'.join(filters.node_labels)
+            node_label_filter = 'n:' + node_labels
+        filter_queries.append(node_label_filter)
 
-    return filter_query, filter_params
+    return filter_queries, filter_params
+
+
+def date_filter_query_constructor(
+    value_name: str, param_name: str, operator: ComparisonOperator
+) -> str:
+    query = '(' + value_name + ' '
+
+    if operator == ComparisonOperator.is_null or operator == ComparisonOperator.is_not_null:
+        query += operator.value + ')'
+    else:
+        query += operator.value + ' ' + param_name + ')'
+
+    return query
 
 
 def edge_search_filter_query_constructor(
     filters: SearchFilters,
-) -> tuple[LiteralString, dict[str, Any]]:
-    filter_query: LiteralString = ''
+    provider: GraphProvider,
+) -> tuple[list[str], dict[str, Any]]:
+    filter_queries: list[str] = []
     filter_params: dict[str, Any] = {}
 
+    if filters.edge_types is not None:
+        edge_types = filters.edge_types
+        filter_queries.append('e.name in $edge_types')
+        filter_params['edge_types'] = edge_types
+
+    if filters.edge_uuids is not None:
+        filter_queries.append('e.uuid in $edge_uuids')
+        filter_params['edge_uuids'] = filters.edge_uuids
+
+    if filters.node_labels is not None:
+        if provider == GraphProvider.KUZU:
+            node_label_filter = (
+                'list_has_all(n.labels, $labels) AND list_has_all(m.labels, $labels)'
+            )
+            filter_params['labels'] = filters.node_labels
+        else:
+            node_labels = '|'.join(filters.node_labels)
+            node_label_filter = 'n:' + node_labels + ' AND m:' + node_labels
+        filter_queries.append(node_label_filter)
+
     if filters.valid_at is not None:
-        valid_at_filter = ' AND ('
+        valid_at_filter = '('
         for i, or_list in enumerate(filters.valid_at):
             for j, date_filter in enumerate(or_list):
-                filter_params['valid_at_' + str(j)] = date_filter.date
+                if date_filter.comparison_operator not in [
+                    ComparisonOperator.is_null,
+                    ComparisonOperator.is_not_null,
+                ]:
+                    filter_params['valid_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.valid_at ' + date_filter.comparison_operator.value + f' $valid_at_{j})'
+                date_filter_query_constructor(
+                    'e.valid_at', f'$valid_at_{j}', date_filter.comparison_operator
+                )
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             valid_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.valid_at) - 1:
                 valid_at_filter += ')'
             else:
                 valid_at_filter += ' OR '
 
-        filter_query += valid_at_filter
+        filter_queries.append(valid_at_filter)
 
     if filters.invalid_at is not None:
-        invalid_at_filter = ' AND ('
+        invalid_at_filter = '('
         for i, or_list in enumerate(filters.invalid_at):
             for j, date_filter in enumerate(or_list):
-                filter_params['invalid_at_' + str(j)] = date_filter.date
+                if date_filter.comparison_operator not in [
+                    ComparisonOperator.is_null,
+                    ComparisonOperator.is_not_null,
+                ]:
+                    filter_params['invalid_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.invalid_at ' + date_filter.comparison_operator.value + f' $invalid_at_{j})'
+                date_filter_query_constructor(
+                    'e.invalid_at', f'$invalid_at_{j}', date_filter.comparison_operator
+                )
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             invalid_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.invalid_at) - 1:
                 invalid_at_filter += ')'
             else:
                 invalid_at_filter += ' OR '
 
-        filter_query += invalid_at_filter
+        filter_queries.append(invalid_at_filter)
 
     if filters.created_at is not None:
-        created_at_filter = ' AND ('
+        created_at_filter = '('
         for i, or_list in enumerate(filters.created_at):
             for j, date_filter in enumerate(or_list):
-                filter_params['created_at_' + str(j)] = date_filter.date
+                if date_filter.comparison_operator not in [
+                    ComparisonOperator.is_null,
+                    ComparisonOperator.is_not_null,
+                ]:
+                    filter_params['created_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.created_at ' + date_filter.comparison_operator.value + f' $created_at_{j})'
+                date_filter_query_constructor(
+                    'e.created_at', f'$created_at_{j}', date_filter.comparison_operator
+                )
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             created_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.created_at) - 1:
                 created_at_filter += ')'
             else:
                 created_at_filter += ' OR '
 
-        filter_query += created_at_filter
+        filter_queries.append(created_at_filter)
 
     if filters.expired_at is not None:
-        expired_at_filter = 'AND ('
+        expired_at_filter = '('
         for i, or_list in enumerate(filters.expired_at):
             for j, date_filter in enumerate(or_list):
-                filter_params['expired_at_' + str(j)] = date_filter.date
+                if date_filter.comparison_operator not in [
+                    ComparisonOperator.is_null,
+                    ComparisonOperator.is_not_null,
+                ]:
+                    filter_params['expired_at_' + str(j)] = date_filter.date
 
             and_filters = [
-                '(r.expired_at ' + date_filter.comparison_operator.value + f' $expired_at_{j})'
+                date_filter_query_constructor(
+                    'e.expired_at', f'$expired_at_{j}', date_filter.comparison_operator
+                )
                 for j, date_filter in enumerate(or_list)
             ]
             and_filter_query = ''
             for j, and_filter in enumerate(and_filters):
                 and_filter_query += and_filter
-                if j != len(and_filter_query) - 1:
+                if j != len(and_filters) - 1:
                     and_filter_query += ' AND '
 
             expired_at_filter += and_filter_query
 
-            if i == len(or_list) - 1:
+            if i == len(filters.expired_at) - 1:
                 expired_at_filter += ')'
             else:
                 expired_at_filter += ' OR '
 
-        filter_query += expired_at_filter
+        filter_queries.append(expired_at_filter)
 
-    return filter_query, filter_params
+    return filter_queries, filter_params

graphiti_core/search/search_helpers.py

@@ -14,9 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import json
-
 from graphiti_core.edges import EntityEdge
+from graphiti_core.prompts.prompt_helpers import to_prompt_json
 from graphiti_core.search.search_config import SearchResults
 
 
@@ -57,16 +56,16 @@ def search_results_to_context_string(search_results: SearchResults) -> str:
     These are the most relevant facts and their valid and invalid dates. Facts are considered valid
     between their valid_at and invalid_at dates. Facts with an invalid_at date of "Present" are considered valid.
     <FACTS>
-    {json.dumps(fact_json, indent=12)}
+            {to_prompt_json(fact_json)}
     </FACTS>
     <ENTITIES>
-    {json.dumps(entity_json, indent=12)}
+            {to_prompt_json(entity_json)}
     </ENTITIES>
     <EPISODES>
-    {json.dumps(episode_json, indent=12)}
+            {to_prompt_json(episode_json)}
     </EPISODES>
     <COMMUNITIES>
-    {json.dumps(community_json, indent=12)}
+            {to_prompt_json(community_json)}
     </COMMUNITIES>
 """