참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
952
참고/neo4j-graphrag-python-main/tests/unit/test_schema.py
Normal file
952
참고/neo4j-graphrag-python-main/tests/unit/test_schema.py
Normal file
@@ -0,0 +1,952 @@
|
||||
# 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"
|
||||
)
|
||||
Reference in New Issue
Block a user