953 lines
30 KiB
Python
953 lines
30 KiB
Python
# Copyright (c) "Neo4j"
|
|
# Neo4j Sweden AB [https://neo4j.com]
|
|
# #
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
# #
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
# #
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from neo4j import Driver, Query
|
|
from neo4j_graphrag.schema import (
|
|
BASE_ENTITY_LABEL,
|
|
BASE_KG_BUILDER_LABEL,
|
|
EXCLUDED_LABELS,
|
|
EXCLUDED_RELS,
|
|
INDEX_QUERY,
|
|
LIST_LIMIT,
|
|
NODE_PROPERTIES_QUERY,
|
|
REL_PROPERTIES_QUERY,
|
|
REL_QUERY,
|
|
_value_sanitize,
|
|
format_schema,
|
|
get_enhanced_schema_cypher,
|
|
get_schema,
|
|
get_structured_schema,
|
|
)
|
|
|
|
|
|
def _query_return_value(*args: Any, **kwargs: Any) -> list[Any]:
|
|
query = kwargs.get("query", args[1] if len(args) > 1 else "")
|
|
if NODE_PROPERTIES_QUERY in query:
|
|
return [
|
|
{
|
|
"output": {
|
|
"properties": [{"property": "property_a", "type": "STRING"}],
|
|
"label": "LabelA",
|
|
}
|
|
}
|
|
]
|
|
if REL_PROPERTIES_QUERY in query:
|
|
return [
|
|
{
|
|
"output": {
|
|
"type": "REL_TYPE",
|
|
"properties": [{"property": "rel_prop", "type": "STRING"}],
|
|
}
|
|
}
|
|
]
|
|
if REL_QUERY in query:
|
|
return [
|
|
{"output": {"start": "LabelA", "type": "REL_TYPE", "end": "LabelB"}},
|
|
{"output": {"start": "LabelA", "type": "REL_TYPE", "end": "LabelC"}},
|
|
]
|
|
if "SHOW CONSTRAINTS" == query:
|
|
return ["fake constraints"]
|
|
if INDEX_QUERY == query:
|
|
return ["fake indexes"]
|
|
|
|
raise AssertionError("Unexpected query")
|
|
|
|
|
|
@patch("neo4j_graphrag.schema.query_database", side_effect=_query_return_value)
|
|
def test_get_schema_ensure_formatted_response(driver: Driver) -> None:
|
|
result = get_schema(driver)
|
|
assert (
|
|
result
|
|
== """Node properties:
|
|
LabelA {property_a: STRING}
|
|
Relationship properties:
|
|
REL_TYPE {rel_prop: STRING}
|
|
The relationships:
|
|
(:LabelA)-[:REL_TYPE]->(:LabelB)
|
|
(:LabelA)-[:REL_TYPE]->(:LabelC)"""
|
|
)
|
|
|
|
|
|
def test_get_structured_schema_happy_path(driver: MagicMock) -> None:
|
|
get_structured_schema(driver)
|
|
assert 5 == driver.execute_query.call_count
|
|
calls = driver.execute_query.call_args_list
|
|
|
|
args, kwargs = calls[0]
|
|
query_obj = args[0]
|
|
assert isinstance(query_obj, Query)
|
|
assert query_obj.text == NODE_PROPERTIES_QUERY
|
|
assert query_obj.timeout is None
|
|
assert kwargs["database_"] is None
|
|
assert kwargs["parameters_"] == {
|
|
"EXCLUDED_LABELS": EXCLUDED_LABELS + [BASE_ENTITY_LABEL, BASE_KG_BUILDER_LABEL],
|
|
"SAMPLE": 1000,
|
|
}
|
|
|
|
args, kwargs = calls[1]
|
|
query_obj = args[0]
|
|
assert isinstance(query_obj, Query)
|
|
assert query_obj.text == REL_PROPERTIES_QUERY
|
|
assert query_obj.timeout is None
|
|
assert kwargs["database_"] is None
|
|
assert kwargs["parameters_"] == {"EXCLUDED_LABELS": EXCLUDED_RELS, "SAMPLE": 1000}
|
|
|
|
args, kwargs = calls[2]
|
|
query_obj = args[0]
|
|
assert isinstance(query_obj, Query)
|
|
assert query_obj.text == REL_QUERY
|
|
assert query_obj.timeout is None
|
|
assert kwargs["database_"] is None
|
|
assert kwargs["parameters_"] == {
|
|
"EXCLUDED_LABELS": EXCLUDED_LABELS + [BASE_ENTITY_LABEL, BASE_KG_BUILDER_LABEL],
|
|
"SAMPLE": 1000,
|
|
}
|
|
|
|
args, kwargs = calls[3]
|
|
query_obj = args[0]
|
|
assert isinstance(query_obj, Query)
|
|
assert query_obj.text == "SHOW CONSTRAINTS"
|
|
assert query_obj.timeout is None
|
|
assert kwargs["database_"] is None
|
|
assert kwargs["parameters_"] == {}
|
|
|
|
args, kwargs = calls[4]
|
|
query_obj = args[0]
|
|
assert isinstance(query_obj, Query)
|
|
assert query_obj.text == INDEX_QUERY
|
|
assert query_obj.timeout is None
|
|
assert kwargs["database_"] is None
|
|
assert kwargs["parameters_"] == {}
|
|
|
|
|
|
@patch("neo4j_graphrag.schema.query_database", side_effect=_query_return_value)
|
|
def test_get_schema_ensure_structured_response(driver: MagicMock) -> None:
|
|
result = get_structured_schema(driver)
|
|
|
|
assert result["node_props"]["LabelA"] == [
|
|
{"property": "property_a", "type": "STRING"}
|
|
]
|
|
assert result["rel_props"]["REL_TYPE"] == [
|
|
{"property": "rel_prop", "type": "STRING"}
|
|
]
|
|
assert result["relationships"] == [
|
|
{"end": "LabelB", "start": "LabelA", "type": "REL_TYPE"},
|
|
{"end": "LabelC", "start": "LabelA", "type": "REL_TYPE"},
|
|
]
|
|
assert result["metadata"]["constraint"] == ["fake constraints"]
|
|
assert result["metadata"]["index"] == ["fake indexes"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"description, input_value, expected_output",
|
|
[
|
|
(
|
|
"Small list",
|
|
{"key1": "value1", "small_list": list(range(15))},
|
|
{"key1": "value1", "small_list": list(range(15))},
|
|
),
|
|
(
|
|
"Oversized list",
|
|
{"key1": "value1", "oversized_list": list(range(LIST_LIMIT + 1))},
|
|
{"key1": "value1"},
|
|
),
|
|
(
|
|
"Nested oversized list",
|
|
{"key1": "value1", "oversized_list": {"key": list(range(150))}},
|
|
{"key1": "value1", "oversized_list": {}},
|
|
),
|
|
(
|
|
"Dict in list",
|
|
{
|
|
"key1": "value1",
|
|
"oversized_list": [1, 2, {"key": list(range(LIST_LIMIT + 1))}],
|
|
},
|
|
{"key1": "value1", "oversized_list": [1, 2, {}]},
|
|
),
|
|
(
|
|
"Dict in nested list",
|
|
{
|
|
"key1": "value1",
|
|
"deeply_nested_lists": [
|
|
[[[{"final_nested_key": list(range(LIST_LIMIT + 1))}]]]
|
|
],
|
|
},
|
|
{"key1": "value1", "deeply_nested_lists": [[[[{}]]]]},
|
|
),
|
|
(
|
|
"Bare oversized list",
|
|
list(range(LIST_LIMIT + 1)),
|
|
None,
|
|
),
|
|
(
|
|
"None value",
|
|
None,
|
|
None,
|
|
),
|
|
],
|
|
)
|
|
def test__value_sanitize(
|
|
description: str, input_value: Dict[str, Any], expected_output: Any
|
|
) -> None:
|
|
"""Test the _value_sanitize function."""
|
|
assert (
|
|
_value_sanitize(input_value) == expected_output
|
|
), f"Failed test case: {description}"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"description, schema, is_enhanced, expected_output",
|
|
[
|
|
(
|
|
"Enhanced, string property with high distinct count",
|
|
{
|
|
"node_props": {
|
|
"Person": [
|
|
{
|
|
"property": "name",
|
|
"type": "STRING",
|
|
"values": ["Alice", "Bob", "Charlie"],
|
|
"distinct_count": 11,
|
|
}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Person**\n"
|
|
' - `name`: STRING Example: "Alice"\n'
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, string property with low distinct count",
|
|
{
|
|
"node_props": {
|
|
"Animal": [
|
|
{
|
|
"property": "species",
|
|
"type": "STRING",
|
|
"values": ["Cat", "Dog"],
|
|
"distinct_count": 2,
|
|
}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Animal**\n"
|
|
" - `species`: STRING Available options: ['Cat', 'Dog']\n"
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, numeric property with min and max",
|
|
{
|
|
"node_props": {
|
|
"Person": [
|
|
{"property": "age", "type": "INTEGER", "min": 20, "max": 70}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Person**\n"
|
|
" - `age`: INTEGER Min: 20, Max: 70\n"
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, numeric property with values",
|
|
{
|
|
"node_props": {
|
|
"Event": [
|
|
{
|
|
"property": "date",
|
|
"type": "DATE",
|
|
"values": ["2021-01-01", "2021-01-02"],
|
|
}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Event**\n"
|
|
' - `date`: DATE Example: "2021-01-01"\n'
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, list property that should be skipped",
|
|
{
|
|
"node_props": {
|
|
"Document": [
|
|
{
|
|
"property": "embedding",
|
|
"type": "LIST",
|
|
"min_size": 150,
|
|
"max_size": 200,
|
|
}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Document**\n"
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, list property that should be included",
|
|
{
|
|
"node_props": {
|
|
"Document": [
|
|
{
|
|
"property": "keywords",
|
|
"type": "LIST",
|
|
"min_size": 2,
|
|
"max_size": 5,
|
|
}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Document**\n"
|
|
" - `keywords`: LIST Min Size: 2, Max Size: 5\n"
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, relationship string property with high distinct count",
|
|
{
|
|
"node_props": {},
|
|
"rel_props": {
|
|
"KNOWS": [
|
|
{
|
|
"property": "since",
|
|
"type": "STRING",
|
|
"values": ["2000", "2001", "2002"],
|
|
"distinct_count": 15,
|
|
}
|
|
]
|
|
},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"\n"
|
|
"Relationship properties:\n"
|
|
"- **KNOWS**\n"
|
|
' - `since`: STRING Example: "2000"\n'
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, relationship string property with low distinct count",
|
|
{
|
|
"node_props": {},
|
|
"rel_props": {
|
|
"LIKES": [
|
|
{
|
|
"property": "intensity",
|
|
"type": "STRING",
|
|
"values": ["High", "Medium", "Low"],
|
|
"distinct_count": 3,
|
|
}
|
|
]
|
|
},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"\n"
|
|
"Relationship properties:\n"
|
|
"- **LIKES**\n"
|
|
" - `intensity`: STRING Available options: ['High', 'Medium', 'Low']\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, relationship numeric property with min and max",
|
|
{
|
|
"node_props": {},
|
|
"rel_props": {
|
|
"WORKS_WITH": [
|
|
{
|
|
"property": "since",
|
|
"type": "INTEGER",
|
|
"min": 1995,
|
|
"max": 2020,
|
|
}
|
|
]
|
|
},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"\n"
|
|
"Relationship properties:\n"
|
|
"- **WORKS_WITH**\n"
|
|
" - `since`: INTEGER Min: 1995, Max: 2020\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, relationship list property that should be skipped",
|
|
{
|
|
"node_props": {},
|
|
"rel_props": {
|
|
"KNOWS": [
|
|
{
|
|
"property": "embedding",
|
|
"type": "LIST",
|
|
"min_size": 150,
|
|
"max_size": 200,
|
|
}
|
|
]
|
|
},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"\n"
|
|
"Relationship properties:\n"
|
|
"- **KNOWS**\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, relationship list property that should be included",
|
|
{
|
|
"node_props": {},
|
|
"rel_props": {
|
|
"KNOWS": [
|
|
{
|
|
"property": "messages",
|
|
"type": "LIST",
|
|
"min_size": 2,
|
|
"max_size": 5,
|
|
}
|
|
]
|
|
},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"\n"
|
|
"Relationship properties:\n"
|
|
"- **KNOWS**\n"
|
|
" - `messages`: LIST Min Size: 2, Max Size: 5\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, relationship numeric property without min and max",
|
|
{
|
|
"node_props": {},
|
|
"rel_props": {
|
|
"OWES": [
|
|
{
|
|
"property": "amount",
|
|
"type": "FLOAT",
|
|
"values": [3.14, 2.71],
|
|
}
|
|
]
|
|
},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"\n"
|
|
"Relationship properties:\n"
|
|
"- **OWES**\n"
|
|
' - `amount`: FLOAT Example: "3.14"\n'
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, property with empty values list",
|
|
{
|
|
"node_props": {
|
|
"Person": [
|
|
{
|
|
"property": "name",
|
|
"type": "STRING",
|
|
"values": [],
|
|
"distinct_count": 15,
|
|
}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Person**\n"
|
|
" - `name`: STRING \n"
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
(
|
|
"Enhanced, property with missing values",
|
|
{
|
|
"node_props": {
|
|
"Person": [
|
|
{
|
|
"property": "name",
|
|
"type": "STRING",
|
|
"distinct_count": 15,
|
|
}
|
|
]
|
|
},
|
|
"rel_props": {},
|
|
"relationships": [],
|
|
},
|
|
True,
|
|
(
|
|
"Node properties:\n"
|
|
"- **Person**\n"
|
|
" - `name`: STRING \n"
|
|
"Relationship properties:\n"
|
|
"\n"
|
|
"The relationships:\n"
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_format_schema(
|
|
description: str, schema: Dict[str, Any], is_enhanced: bool, expected_output: str
|
|
) -> None:
|
|
result = format_schema(schema, is_enhanced)
|
|
assert result == expected_output, f"Failed test case: {description}"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"description, structured_schema, label_or_type, properties, exhaustive, sample_size, is_relationship, expected_query",
|
|
[
|
|
(
|
|
"Exhaustive, string property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "name", "type": "STRING"}],
|
|
True,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`)\n"
|
|
"WITH collect(distinct substring(toString(n.`name`), 0, 50)) AS `name_values`\n"
|
|
"RETURN {`name`: {values: `name_values`[..10], distinct_count: size(`name_values`)}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Non-exhaustive, string property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "name", "type": "STRING"}],
|
|
False,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH collect(distinct substring(toString(n.`name`), 0, 50)) AS `name_values`\n"
|
|
"RETURN {`name`: {values: `name_values`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Exhaustive, integer property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "age", "type": "INTEGER"}],
|
|
True,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`)\n"
|
|
"WITH min(n.`age`) AS `age_min`,\n"
|
|
" max(n.`age`) AS `age_max`,\n"
|
|
" count(distinct n.`age`) AS `age_distinct`\n"
|
|
"RETURN {`age`: {min: toString(`age_min`), max: toString(`age_max`), distinct_count: `age_distinct`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Non-exhaustive, integer property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "age", "type": "INTEGER"}],
|
|
False,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH collect(distinct toString(n.`age`)) AS `age_values`\n"
|
|
"RETURN {`age`: {values: `age_values`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Exhaustive, float property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "salary", "type": "FLOAT"}],
|
|
True,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`)\n"
|
|
"WITH min(n.`salary`) AS `salary_min`,\n"
|
|
" max(n.`salary`) AS `salary_max`,\n"
|
|
" count(distinct n.`salary`) AS `salary_distinct`\n"
|
|
"RETURN {`salary`: {min: toString(`salary_min`), max: toString(`salary_max`), distinct_count: `salary_distinct`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Non-exhaustive, float property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "salary", "type": "FLOAT"}],
|
|
False,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH collect(distinct toString(n.`salary`)) AS `salary_values`\n"
|
|
"RETURN {`salary`: {values: `salary_values`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Exhaustive, date property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "birth_date", "type": "DATE"}],
|
|
True,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`)\n"
|
|
"WITH min(n.`birth_date`) AS `birth_date_min`,\n"
|
|
" max(n.`birth_date`) AS `birth_date_max`,\n"
|
|
" count(distinct n.`birth_date`) AS `birth_date_distinct`\n"
|
|
"RETURN {`birth_date`: {min: toString(`birth_date_min`), max: toString(`birth_date_max`), distinct_count: `birth_date_distinct`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Non-exhaustive, date property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "birth_date", "type": "DATE"}],
|
|
False,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH collect(distinct toString(n.`birth_date`)) AS `birth_date_values`\n"
|
|
"RETURN {`birth_date`: {values: `birth_date_values`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Exhaustive, date time property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "birth_date", "type": "DATE_TIME"}],
|
|
True,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`)\n"
|
|
"WITH min(n.`birth_date`) AS `birth_date_min`,\n"
|
|
" max(n.`birth_date`) AS `birth_date_max`,\n"
|
|
" count(distinct n.`birth_date`) AS `birth_date_distinct`\n"
|
|
"RETURN {`birth_date`: {min: toString(`birth_date_min`), max: toString(`birth_date_max`), distinct_count: `birth_date_distinct`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Non-exhaustive, date time property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "birth_date", "type": "DATE_TIME"}],
|
|
False,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH collect(distinct toString(n.`birth_date`)) AS `birth_date_values`\n"
|
|
"RETURN {`birth_date`: {values: `birth_date_values`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Exhaustive, local date time property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "birth_date", "type": "LOCAL_DATE_TIME"}],
|
|
True,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`)\n"
|
|
"WITH min(n.`birth_date`) AS `birth_date_min`,\n"
|
|
" max(n.`birth_date`) AS `birth_date_max`,\n"
|
|
" count(distinct n.`birth_date`) AS `birth_date_distinct`\n"
|
|
"RETURN {`birth_date`: {min: toString(`birth_date_min`), max: toString(`birth_date_max`), distinct_count: `birth_date_distinct`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Non-exhaustive, local date time property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "birth_date", "type": "LOCAL_DATE_TIME"}],
|
|
False,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH collect(distinct toString(n.`birth_date`)) AS `birth_date_values`\n"
|
|
"RETURN {`birth_date`: {values: `birth_date_values`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Exhaustive, list property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "tags", "type": "LIST"}],
|
|
True,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`)\n"
|
|
"WITH min(size(n.`tags`)) AS `tags_size_min`, max(size(n.`tags`)) AS `tags_size_max`\n"
|
|
"RETURN {`tags`: {min_size: `tags_size_min`, max_size: `tags_size_max`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Non-exhaustive, list property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "tags", "type": "LIST"}],
|
|
False,
|
|
5,
|
|
False,
|
|
(
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH min(size(n.`tags`)) AS `tags_size_min`, max(size(n.`tags`)) AS `tags_size_max`\n"
|
|
"RETURN {`tags`: {min_size: `tags_size_min`, max_size: `tags_size_max`}} AS output"
|
|
),
|
|
),
|
|
(
|
|
"Exhaustive, boolean property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "active", "type": "BOOLEAN"}],
|
|
True,
|
|
5,
|
|
False,
|
|
"MATCH (n:`Person`)\nRETURN {} AS output",
|
|
),
|
|
(
|
|
"Non-exhaustive, boolean property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "active", "type": "BOOLEAN"}],
|
|
False,
|
|
5,
|
|
False,
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\nRETURN {} AS output",
|
|
),
|
|
(
|
|
"Exhaustive, point property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "location", "type": "POINT"}],
|
|
True,
|
|
5,
|
|
False,
|
|
"MATCH (n:`Person`)\nRETURN {} AS output",
|
|
),
|
|
(
|
|
"Non-exhaustive, point property",
|
|
{"metadata": {"index": []}},
|
|
"Person",
|
|
[{"property": "location", "type": "POINT"}],
|
|
False,
|
|
5,
|
|
False,
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\nRETURN {} AS output",
|
|
),
|
|
(
|
|
"Exhaustive, duration property",
|
|
{"metadata": {"index": []}},
|
|
"Journey",
|
|
[{"property": "duration", "type": "DURATION"}],
|
|
True,
|
|
5,
|
|
False,
|
|
("MATCH (n:`Journey`)\nRETURN {} AS output"),
|
|
),
|
|
(
|
|
"Non-exhaustive, duration property",
|
|
{"metadata": {"index": []}},
|
|
"Journey",
|
|
[{"property": "duration", "type": "DURATION"}],
|
|
False,
|
|
5,
|
|
False,
|
|
"MATCH (n:`Journey`) WITH n LIMIT 5\nRETURN {} AS output",
|
|
),
|
|
(
|
|
"Relationship property",
|
|
{"metadata": {"index": []}},
|
|
"FRIENDS_WITH",
|
|
[{"property": "since", "type": "INTEGER"}],
|
|
True,
|
|
5,
|
|
True,
|
|
(
|
|
"MATCH ()-[n:`FRIENDS_WITH`]->()\n"
|
|
"WITH min(n.`since`) AS `since_min`,\n"
|
|
" max(n.`since`) AS `since_max`,\n"
|
|
" count(distinct n.`since`) AS `since_distinct`\n"
|
|
"RETURN {`since`: {min: toString(`since_min`), max: toString(`since_max`), distinct_count: `since_distinct`}} AS output"
|
|
),
|
|
),
|
|
],
|
|
)
|
|
def test_get_enhanced_schema_cypher(
|
|
driver: MagicMock,
|
|
description: str,
|
|
structured_schema: Dict[str, Any],
|
|
label_or_type: str,
|
|
properties: List[Dict[str, Any]],
|
|
exhaustive: bool,
|
|
sample_size: int,
|
|
is_relationship: bool,
|
|
expected_query: str,
|
|
) -> None:
|
|
query = get_enhanced_schema_cypher(
|
|
driver=driver,
|
|
structured_schema=structured_schema,
|
|
label_or_type=label_or_type,
|
|
properties=properties,
|
|
exhaustive=exhaustive,
|
|
sample_size=sample_size,
|
|
is_relationship=is_relationship,
|
|
)
|
|
assert query == expected_query, f"Failed test case: {description}"
|
|
|
|
|
|
def test_enhanced_schema_cypher_integer_exhaustive_false_with_index(
|
|
driver: MagicMock,
|
|
) -> None:
|
|
structured_schema = {
|
|
"metadata": {
|
|
"index": [
|
|
{
|
|
"label": "Person",
|
|
"properties": ["age"],
|
|
"type": "RANGE",
|
|
}
|
|
]
|
|
}
|
|
}
|
|
properties = [{"property": "age", "type": "INTEGER"}]
|
|
query = get_enhanced_schema_cypher(
|
|
driver=driver,
|
|
structured_schema=structured_schema,
|
|
label_or_type="Person",
|
|
properties=properties,
|
|
exhaustive=False,
|
|
)
|
|
assert query == (
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n"
|
|
"WITH min(n.`age`) AS `age_min`,\n"
|
|
" max(n.`age`) AS `age_max`,\n"
|
|
" count(distinct n.`age`) AS `age_distinct`\n"
|
|
"RETURN {`age`: {min: toString(`age_min`), max: toString(`age_max`), distinct_count: `age_distinct`}} AS output"
|
|
)
|
|
|
|
|
|
@patch("neo4j_graphrag.schema.query_database")
|
|
def test_enhanced_schema_cypher_string_exhaustive_false_with_index(
|
|
query_database_mock: MagicMock,
|
|
driver: MagicMock,
|
|
) -> None:
|
|
structured_schema = {
|
|
"metadata": {
|
|
"index": [
|
|
{
|
|
"label": "Person",
|
|
"properties": ["status"],
|
|
"type": "RANGE",
|
|
"size": 5,
|
|
"distinctValues": 5,
|
|
}
|
|
]
|
|
}
|
|
}
|
|
query_database_mock.return_value = [{"value": ["Single", "Married", "Divorced"]}]
|
|
properties = [{"property": "status", "type": "STRING"}]
|
|
query = get_enhanced_schema_cypher(
|
|
driver=driver,
|
|
structured_schema=structured_schema,
|
|
label_or_type="Person",
|
|
properties=properties,
|
|
exhaustive=False,
|
|
)
|
|
assert query == (
|
|
"MATCH (n:`Person`) WITH n LIMIT 5\n\n"
|
|
"RETURN {`status`: {values: ['Single', 'Married', 'Divorced'], distinct_count: 3}} AS output"
|
|
)
|