2204 lines
71 KiB
Python
2204 lines
71 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
|
|
|
|
import json
|
|
from typing import Tuple, Any
|
|
from unittest.mock import AsyncMock, patch, Mock
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from neo4j_graphrag.exceptions import SchemaValidationError, SchemaExtractionError
|
|
from neo4j_graphrag.experimental.components.schema import (
|
|
SchemaBuilder,
|
|
NodeType,
|
|
PropertyType,
|
|
RelationshipType,
|
|
ConstraintType,
|
|
GraphConstraintType,
|
|
SchemaFromTextExtractor,
|
|
GraphSchema,
|
|
SchemaFromExistingGraphExtractor,
|
|
Pattern,
|
|
)
|
|
import os
|
|
import tempfile
|
|
import yaml
|
|
|
|
from neo4j_graphrag.generation import PromptTemplate
|
|
from neo4j_graphrag.llm.types import LLMResponse
|
|
from neo4j_graphrag.utils.file_handler import FileFormat
|
|
|
|
# Most tests below read ``PropertyType.required`` to assert migration; silence only that
|
|
# deprecation. Explicit warning behavior is covered in ``test_property_type_deprecation.py``.
|
|
pytestmark = pytest.mark.filterwarnings(
|
|
"ignore:Use GraphSchema.constraints with type EXISTENCE:DeprecationWarning"
|
|
)
|
|
|
|
|
|
def test_node_type_initialization_from_string() -> None:
|
|
node_type = NodeType.model_validate("Label")
|
|
assert isinstance(node_type, NodeType)
|
|
assert node_type.label == "Label"
|
|
assert len(node_type.properties) == 1
|
|
assert node_type.properties[0].name == "name"
|
|
assert node_type.properties[0].type == "STRING"
|
|
|
|
|
|
def test_node_type_additional_properties_default() -> None:
|
|
# default behavior:
|
|
node_type = NodeType.model_validate({"label": "Label"})
|
|
assert node_type.additional_properties is True
|
|
# Empty properties list now violates min_length=1 constraint
|
|
with pytest.raises(ValidationError):
|
|
NodeType.model_validate({"label": "Label", "properties": []})
|
|
node_type = NodeType.model_validate(
|
|
{"label": "Label", "properties": [{"name": "name", "type": "STRING"}]}
|
|
)
|
|
assert node_type.additional_properties is False
|
|
|
|
# manually changing the default value
|
|
# impossible cases: no properties and no additional
|
|
with pytest.raises(ValidationError):
|
|
NodeType.model_validate({"label": "Label", "additional_properties": False})
|
|
with pytest.raises(ValidationError):
|
|
NodeType.model_validate(
|
|
{"label": "Label", "properties": [], "additional_properties": False}
|
|
)
|
|
|
|
# working case: properties and additional allowed
|
|
node_type = NodeType.model_validate(
|
|
{
|
|
"label": "Label",
|
|
"properties": [{"name": "name", "type": "STRING"}],
|
|
"additional_properties": True,
|
|
}
|
|
)
|
|
assert node_type.additional_properties is True
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"label",
|
|
["__Person", "Person__", "__Person__", "__KNOWS", "KNOWS__", "__KNOWS__"],
|
|
)
|
|
def test_node_type_rejects_dunder_label(label: str) -> None:
|
|
with pytest.raises(ValidationError, match="reserved '__'"):
|
|
NodeType.model_validate(
|
|
{"label": label, "properties": [{"name": "name", "type": "STRING"}]}
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"label",
|
|
["__KNOWS", "KNOWS__", "__KNOWS__"],
|
|
)
|
|
def test_relationship_type_rejects_dunder_label(label: str) -> None:
|
|
with pytest.raises(ValidationError, match="reserved '__'"):
|
|
RelationshipType.model_validate({"label": label})
|
|
|
|
|
|
def test_property_type_initalization() -> None:
|
|
prop = PropertyType(name="email", type="STRING")
|
|
assert prop.name == "email"
|
|
assert prop.type == "STRING"
|
|
assert prop.required is False
|
|
|
|
|
|
def test_property_type_with_required_true() -> None:
|
|
prop = PropertyType(name="id", type="INTEGER", required=True)
|
|
assert prop.required is True
|
|
|
|
|
|
def test_property_type_is_frozen() -> None:
|
|
prop = PropertyType(name="email", type="STRING", required=False)
|
|
|
|
with pytest.raises(ValidationError):
|
|
prop.name = "other"
|
|
|
|
with pytest.raises(ValidationError):
|
|
prop.required = True
|
|
|
|
|
|
def test_relationship_type_initialization_from_string() -> None:
|
|
relationship_type = RelationshipType.model_validate("REL")
|
|
assert isinstance(relationship_type, RelationshipType)
|
|
assert relationship_type.label == "REL"
|
|
assert relationship_type.properties == []
|
|
|
|
|
|
def test_relationship_type_additional_properties_default() -> None:
|
|
relationship_type = RelationshipType.model_validate({"label": "REL"})
|
|
assert relationship_type.additional_properties is True
|
|
relationship_type = RelationshipType.model_validate(
|
|
{"label": "REL", "properties": []}
|
|
)
|
|
assert relationship_type.additional_properties is True
|
|
relationship_type = RelationshipType.model_validate(
|
|
{"label": "REL", "properties": [{"name": "name", "type": "STRING"}]}
|
|
)
|
|
assert relationship_type.additional_properties is False
|
|
|
|
# manually changing the default value
|
|
# auto-correction: no properties and additional_properties=False -> auto-corrected to True
|
|
relationship_type = RelationshipType.model_validate(
|
|
{"label": "REL", "additional_properties": False}
|
|
)
|
|
assert relationship_type.additional_properties is True # auto-corrected
|
|
relationship_type = RelationshipType.model_validate(
|
|
{"label": "REL", "properties": [], "additional_properties": False}
|
|
)
|
|
assert relationship_type.additional_properties is True # auto-corrected
|
|
|
|
# working case: properties and additional allowed
|
|
relationship_type = RelationshipType.model_validate(
|
|
{
|
|
"label": "REL",
|
|
"properties": [{"name": "name", "type": "STRING"}],
|
|
"additional_properties": True,
|
|
}
|
|
)
|
|
assert relationship_type.additional_properties is True
|
|
|
|
|
|
def test_constraint_type_initialization() -> None:
|
|
constraint = ConstraintType(
|
|
type=GraphConstraintType.UNIQUENESS,
|
|
node_type="Person",
|
|
property_names=("name",),
|
|
)
|
|
assert constraint.type == "UNIQUENESS"
|
|
assert constraint.node_type == "Person"
|
|
assert constraint.property_names == ("name",)
|
|
|
|
|
|
def test_constraint_type_is_frozen() -> None:
|
|
constraint = ConstraintType(
|
|
type=GraphConstraintType.UNIQUENESS,
|
|
node_type="Person",
|
|
property_names=("name",),
|
|
)
|
|
|
|
with pytest.raises(ValidationError):
|
|
constraint.type = GraphConstraintType.UNIQUENESS
|
|
|
|
with pytest.raises(ValidationError):
|
|
constraint.node_type = "Organization"
|
|
|
|
with pytest.raises(ValidationError):
|
|
constraint.property_names = ("id",)
|
|
|
|
|
|
def test_constraint_type_rejects_empty_property_names() -> None:
|
|
"""``property_names`` has ``min_length=1``; empty tuple/list fails validation."""
|
|
with pytest.raises(ValidationError) as exc_info:
|
|
ConstraintType(
|
|
type=GraphConstraintType.UNIQUENESS,
|
|
node_type="Person",
|
|
property_names=(),
|
|
)
|
|
assert "too_short" in str(exc_info.value).lower()
|
|
|
|
with pytest.raises(ValidationError):
|
|
ConstraintType.model_validate(
|
|
{
|
|
"type": "UNIQUENESS",
|
|
"node_type": "Person",
|
|
"property_names": [],
|
|
}
|
|
)
|
|
|
|
|
|
def test_schema_additional_node_types_default() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.additional_node_types is True
|
|
|
|
schema_dict = {
|
|
"node_types": ["Person"],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.additional_node_types is False
|
|
|
|
|
|
def test_schema_additional_relationship_types_default() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.additional_relationship_types is True
|
|
|
|
schema_dict = {
|
|
"node_types": [],
|
|
"relationship_types": ["REL"],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.additional_relationship_types is False
|
|
|
|
|
|
def test_schema_additional_patterns_default() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.additional_patterns is True
|
|
|
|
schema_dict = {
|
|
"node_types": ["Person"],
|
|
"relationship_types": ["REL"],
|
|
"patterns": [("Person", "REL", "Person")],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.additional_patterns is False
|
|
|
|
|
|
def test_schema_additional_parameter_validation() -> None:
|
|
"""Additional relationship types not allowed, but additional patterns allowed
|
|
|
|
=> raise Exception
|
|
"""
|
|
schema_dict = {
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{
|
|
"name": "name",
|
|
"type": "STRING",
|
|
},
|
|
{"name": "height", "type": "INTEGER"},
|
|
],
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "KNOWS",
|
|
}
|
|
],
|
|
"patterns": [
|
|
("Person", "KNOWS", "Person"),
|
|
],
|
|
"additional_relationship_types": True,
|
|
"additional_patterns": False,
|
|
}
|
|
with pytest.raises(
|
|
ValidationError,
|
|
match="`additional_relationship_types` must be set to False when using `additional_patterns=False`",
|
|
):
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
|
|
def test_schema_constraint_validation_property_not_in_node_type() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [{"name": "name", "type": "STRING"}],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "email"}
|
|
],
|
|
}
|
|
|
|
with pytest.raises(SchemaValidationError) as exc_info:
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
assert "Constraint references undefined property" in str(exc_info.value)
|
|
assert "on node type 'Person'" in str(exc_info.value)
|
|
|
|
|
|
def test_schema_constraint_with_additional_properties_with_allows_unknown_property() -> (
|
|
None
|
|
):
|
|
# if additional_properties is True, we can define constraints that are not in the node_type
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [{"name": "name", "type": "STRING"}],
|
|
"additional_properties": True,
|
|
}
|
|
],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "email"}
|
|
],
|
|
}
|
|
|
|
# Should raise - email is not allowed because the property is not defined in the node
|
|
with pytest.raises(SchemaValidationError) as exc_info:
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
assert "Constraint references undefined property 'email'" in str(exc_info.value)
|
|
|
|
|
|
def test_schema_with_valid_constraints() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
|
|
],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
|
|
],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].type == "UNIQUENESS"
|
|
assert schema.constraints[0].node_type == "Person"
|
|
assert schema.constraints[0].property_names == ("name",)
|
|
|
|
|
|
def test_schema_constraint_validation_invalid_node_type() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "UNIQUENESS",
|
|
"node_type": "NonExistentNode",
|
|
"property_name": "id",
|
|
}
|
|
],
|
|
}
|
|
|
|
with pytest.raises(SchemaValidationError) as exc_info:
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
assert "Constraint references undefined node type: NonExistentNode" in str(
|
|
exc_info.value
|
|
)
|
|
|
|
|
|
def test_schema_constraint_validation_missing_property_name() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
|
|
],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": ""}
|
|
],
|
|
}
|
|
|
|
with pytest.raises((SchemaValidationError, ValidationError)):
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
|
|
def test_schema_key_constraint_node_valid() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{"label": "Person", "properties": [{"name": "email", "type": "STRING"}]}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "KEY",
|
|
"node_type": "Person",
|
|
"property_name": "email",
|
|
"relationship_type": None,
|
|
}
|
|
],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.key_property_names_for_node("Person") == {"email"}
|
|
assert schema.mandatory_property_names_for_node("Person") == {"email"}
|
|
assert schema.uniqueness_property_names_for_node("Person") == set()
|
|
|
|
|
|
def test_schema_key_constraint_relationship_valid() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "KNOWS",
|
|
"properties": [{"name": "since", "type": "INTEGER"}],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "KEY",
|
|
"node_type": "",
|
|
"property_name": "since",
|
|
"relationship_type": "KNOWS",
|
|
}
|
|
],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.key_property_names_for_relationship("KNOWS") == {"since"}
|
|
assert schema.mandatory_property_names_for_relationship("KNOWS") == {"since"}
|
|
|
|
|
|
def test_schema_uniqueness_and_key_same_property_rejected() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{"label": "Person", "properties": [{"name": "id", "type": "STRING"}]}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "UNIQUENESS",
|
|
"node_type": "Person",
|
|
"property_name": "id",
|
|
"relationship_type": None,
|
|
},
|
|
{
|
|
"type": "KEY",
|
|
"node_type": "Person",
|
|
"property_name": "id",
|
|
"relationship_type": None,
|
|
},
|
|
],
|
|
}
|
|
with pytest.raises(SchemaValidationError, match="UNIQUENESS and KEY"):
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
|
|
# --- Composite (multi-property) constraint tests ---
|
|
|
|
|
|
def test_composite_key_constraint_node_valid() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Actor",
|
|
"properties": [
|
|
{"name": "firstname", "type": "STRING"},
|
|
{"name": "surname", "type": "STRING"},
|
|
],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "KEY",
|
|
"node_type": "Actor",
|
|
"property_names": ["firstname", "surname"],
|
|
"relationship_type": None,
|
|
}
|
|
],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].property_names == ("firstname", "surname")
|
|
assert schema.key_property_names_for_node("Actor") == {"firstname", "surname"}
|
|
assert schema.mandatory_property_names_for_node("Actor") == {
|
|
"firstname",
|
|
"surname",
|
|
}
|
|
|
|
|
|
def test_composite_uniqueness_constraint_valid() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Book",
|
|
"properties": [
|
|
{"name": "title", "type": "STRING"},
|
|
{"name": "year", "type": "INTEGER"},
|
|
],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "UNIQUENESS",
|
|
"node_type": "Book",
|
|
"property_names": ["title", "year"],
|
|
}
|
|
],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].property_names == ("title", "year")
|
|
assert schema.uniqueness_property_names_for_node("Book") == {"title", "year"}
|
|
|
|
|
|
def test_existence_rejects_composite() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"},
|
|
{"name": "email", "type": "STRING"},
|
|
],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "EXISTENCE",
|
|
"node_type": "Person",
|
|
"property_names": ["name", "email"],
|
|
}
|
|
],
|
|
}
|
|
with pytest.raises(ValueError, match="EXISTENCE constraint must have exactly one"):
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
|
|
def test_backward_compat_property_name_migrates_to_property_names() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{"label": "Person", "properties": [{"name": "name", "type": "STRING"}]}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "UNIQUENESS",
|
|
"node_type": "Person",
|
|
"property_name": "name",
|
|
}
|
|
],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert schema.constraints[0].property_names == ("name",)
|
|
with pytest.warns(DeprecationWarning, match="property_names"):
|
|
assert schema.constraints[0].property_name == "name"
|
|
|
|
|
|
def test_composite_constraint_validates_all_properties_exist() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Actor",
|
|
"properties": [
|
|
{"name": "firstname", "type": "STRING"},
|
|
],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "KEY",
|
|
"node_type": "Actor",
|
|
"property_names": ["firstname", "nonexistent"],
|
|
"relationship_type": None,
|
|
}
|
|
],
|
|
}
|
|
with pytest.raises(SchemaValidationError, match="nonexistent"):
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
|
|
def test_composite_uniqueness_and_key_same_properties_rejected() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Actor",
|
|
"properties": [
|
|
{"name": "firstname", "type": "STRING"},
|
|
{"name": "surname", "type": "STRING"},
|
|
],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "UNIQUENESS",
|
|
"node_type": "Actor",
|
|
"property_names": ["firstname", "surname"],
|
|
},
|
|
{
|
|
"type": "KEY",
|
|
"node_type": "Actor",
|
|
"property_names": ["firstname", "surname"],
|
|
"relationship_type": None,
|
|
},
|
|
],
|
|
}
|
|
with pytest.raises(SchemaValidationError, match="UNIQUENESS and KEY"):
|
|
GraphSchema.model_validate(schema_dict)
|
|
|
|
|
|
def test_composite_uniqueness_and_key_different_properties_allowed() -> None:
|
|
schema_dict: dict[str, Any] = {
|
|
"node_types": [
|
|
{
|
|
"label": "Actor",
|
|
"properties": [
|
|
{"name": "firstname", "type": "STRING"},
|
|
{"name": "surname", "type": "STRING"},
|
|
{"name": "email", "type": "STRING"},
|
|
],
|
|
}
|
|
],
|
|
"constraints": [
|
|
{
|
|
"type": "UNIQUENESS",
|
|
"node_type": "Actor",
|
|
"property_names": ["email"],
|
|
},
|
|
{
|
|
"type": "KEY",
|
|
"node_type": "Actor",
|
|
"property_names": ["firstname", "surname"],
|
|
"relationship_type": None,
|
|
},
|
|
],
|
|
}
|
|
schema = GraphSchema.model_validate(schema_dict)
|
|
assert len(schema.constraints) == 2
|
|
assert schema.key_property_names_for_node("Actor") == {"firstname", "surname"}
|
|
assert schema.uniqueness_property_names_for_node("Actor") == {"email"}
|
|
|
|
|
|
def test_extract_graph_constraints_from_metadata_composite_key() -> None:
|
|
from neo4j_graphrag.experimental.components.schema import (
|
|
SchemaFromExistingGraphExtractor,
|
|
)
|
|
|
|
structured_schema: dict[str, Any] = {
|
|
"metadata": {
|
|
"constraint": [
|
|
{
|
|
"type": "NODE_KEY",
|
|
"properties": ["firstname", "surname"],
|
|
"labelsOrTypes": ["Actor"],
|
|
}
|
|
]
|
|
}
|
|
}
|
|
result = SchemaFromExistingGraphExtractor._extract_graph_constraints_from_metadata(
|
|
structured_schema
|
|
)
|
|
assert len(result) == 1
|
|
assert result[0]["type"] == "KEY"
|
|
assert result[0]["node_type"] == "Actor"
|
|
assert result[0]["property_names"] == ("firstname", "surname")
|
|
assert result[0]["property_name"] == "firstname"
|
|
|
|
|
|
@pytest.fixture
|
|
def valid_node_types() -> tuple[NodeType, ...]:
|
|
# required=False so tuples match GraphSchema after validation (legacy required=True
|
|
# is migrated to EXISTENCE and cleared).
|
|
return (
|
|
NodeType(
|
|
label="PERSON",
|
|
description="An individual human being.",
|
|
properties=[
|
|
PropertyType(name="birth date", type="ZONED_DATETIME"),
|
|
PropertyType(name="name", type="STRING", required=False),
|
|
],
|
|
additional_properties=False,
|
|
),
|
|
NodeType(
|
|
label="ORGANIZATION",
|
|
description="A structured group of people with a common purpose.",
|
|
properties=[PropertyType(name="name", type="STRING")],
|
|
additional_properties=True,
|
|
),
|
|
NodeType(
|
|
label="AGE",
|
|
description="Age of a person in years.",
|
|
properties=[PropertyType(name="value", type="INTEGER")],
|
|
additional_properties=True,
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def valid_relationship_types() -> tuple[RelationshipType, ...]:
|
|
return (
|
|
RelationshipType(
|
|
label="EMPLOYED_BY",
|
|
description="Indicates employment relationship.",
|
|
properties=[
|
|
PropertyType(name="start_time", type="LOCAL_DATETIME", required=False),
|
|
PropertyType(name="end_time", type="LOCAL_DATETIME"),
|
|
],
|
|
additional_properties=False,
|
|
),
|
|
RelationshipType(
|
|
label="ORGANIZED_BY",
|
|
description="Indicates organization responsible for an event.",
|
|
),
|
|
RelationshipType(
|
|
label="ATTENDED_BY", description="Indicates attendance at an event."
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def valid_patterns() -> tuple[Pattern, ...]:
|
|
return (
|
|
Pattern(source="PERSON", relationship="EMPLOYED_BY", target="ORGANIZATION"),
|
|
Pattern(source="ORGANIZATION", relationship="ATTENDED_BY", target="PERSON"),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def patterns_with_invalid_entity() -> tuple[Pattern, ...]:
|
|
return (
|
|
Pattern(source="PERSON", relationship="EMPLOYED_BY", target="ORGANIZATION"),
|
|
Pattern(
|
|
source="NON_EXISTENT_ENTITY", relationship="ATTENDED_BY", target="PERSON"
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def valid_constraints() -> tuple[ConstraintType, ...]:
|
|
return (
|
|
ConstraintType(
|
|
type=GraphConstraintType.UNIQUENESS,
|
|
node_type="PERSON",
|
|
property_names=("name",),
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def patterns_with_invalid_relation() -> tuple[Pattern, ...]:
|
|
return (
|
|
Pattern(
|
|
source="PERSON", relationship="NON_EXISTENT_RELATION", target="ORGANIZATION"
|
|
),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_builder() -> SchemaBuilder:
|
|
return SchemaBuilder()
|
|
|
|
|
|
@pytest.fixture
|
|
def graph_schema(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_relationship_types: Tuple[RelationshipType, ...],
|
|
valid_patterns: tuple[Pattern, ...],
|
|
) -> GraphSchema:
|
|
return schema_builder.create_schema_model(
|
|
list(valid_node_types), list(valid_relationship_types), list(valid_patterns)
|
|
)
|
|
|
|
|
|
def test_create_schema_model_valid_data(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_relationship_types: Tuple[RelationshipType, ...],
|
|
valid_patterns: tuple[Pattern, ...],
|
|
) -> None:
|
|
schema = schema_builder.create_schema_model(
|
|
list(valid_node_types), list(valid_relationship_types), list(valid_patterns)
|
|
)
|
|
|
|
assert schema.node_types == valid_node_types
|
|
assert schema.relationship_types == valid_relationship_types
|
|
assert schema.patterns == valid_patterns
|
|
assert schema.additional_node_types is False
|
|
assert schema.additional_relationship_types is False
|
|
assert schema.additional_patterns is False
|
|
|
|
|
|
def test_create_schema_model_with_constraints(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_constraints: Tuple[ConstraintType, ...],
|
|
) -> None:
|
|
schema = schema_builder.create_schema_model(
|
|
list(valid_node_types),
|
|
constraints=list(valid_constraints),
|
|
)
|
|
|
|
assert schema.node_types == valid_node_types
|
|
assert schema.constraints == valid_constraints
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].type == "UNIQUENESS"
|
|
assert schema.constraints[0].node_type == "PERSON"
|
|
assert schema.constraints[0].property_names == ("name",)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_method(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_relationship_types: Tuple[RelationshipType, ...],
|
|
valid_patterns: tuple[Pattern, ...],
|
|
) -> None:
|
|
with patch.object(
|
|
schema_builder,
|
|
"create_schema_model",
|
|
return_value=GraphSchema(
|
|
node_types=valid_node_types,
|
|
relationship_types=valid_relationship_types,
|
|
patterns=valid_patterns,
|
|
),
|
|
):
|
|
schema = await schema_builder.run(
|
|
list(valid_node_types), list(valid_relationship_types), list(valid_patterns)
|
|
)
|
|
|
|
assert schema.node_types == valid_node_types
|
|
assert schema.relationship_types == valid_relationship_types
|
|
assert schema.patterns == valid_patterns
|
|
assert schema.additional_node_types is False
|
|
assert schema.additional_relationship_types is False
|
|
assert schema.additional_patterns is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_method_with_constraints(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_constraints: Tuple[ConstraintType, ...],
|
|
) -> None:
|
|
schema = await schema_builder.run(
|
|
list(valid_node_types),
|
|
constraints=list(valid_constraints),
|
|
)
|
|
|
|
assert schema.node_types == valid_node_types
|
|
assert schema.constraints == valid_constraints
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].type == "UNIQUENESS"
|
|
assert schema.constraints[0].node_type == "PERSON"
|
|
assert schema.constraints[0].property_names == ("name",)
|
|
|
|
|
|
def test_create_schema_model_invalid_entity(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_relationship_types: Tuple[RelationshipType, ...],
|
|
patterns_with_invalid_entity: tuple[Pattern, ...],
|
|
) -> None:
|
|
with pytest.raises(SchemaValidationError) as exc_info:
|
|
schema_builder.create_schema_model(
|
|
list(valid_node_types),
|
|
list(valid_relationship_types),
|
|
list(patterns_with_invalid_entity),
|
|
)
|
|
assert "Node type 'NON_EXISTENT_ENTITY' is not defined" in str(
|
|
exc_info.value
|
|
), "Should fail due to non-existent entity"
|
|
|
|
|
|
def test_create_schema_model_invalid_relation(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_relationship_types: Tuple[RelationshipType, ...],
|
|
patterns_with_invalid_relation: tuple[Pattern, ...],
|
|
) -> None:
|
|
with pytest.raises(SchemaValidationError) as exc_info:
|
|
schema_builder.create_schema_model(
|
|
list(valid_node_types),
|
|
list(valid_relationship_types),
|
|
list(patterns_with_invalid_relation),
|
|
)
|
|
assert "Relationship type 'NON_EXISTENT_RELATION' is not defined" in str(
|
|
exc_info.value
|
|
), "Should fail due to non-existent relation"
|
|
|
|
|
|
def test_create_schema_model_no_potential_schema(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_relationship_types: Tuple[RelationshipType, ...],
|
|
) -> None:
|
|
schema_instance = schema_builder.create_schema_model(
|
|
list(valid_node_types), list(valid_relationship_types)
|
|
)
|
|
assert schema_instance.node_types == valid_node_types
|
|
assert schema_instance.relationship_types == valid_relationship_types
|
|
assert schema_instance.patterns == ()
|
|
|
|
|
|
def test_create_schema_model_no_relations_or_potential_schema(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
) -> None:
|
|
schema_instance = schema_builder.create_schema_model(list(valid_node_types))
|
|
|
|
assert len(schema_instance.node_types) == 3
|
|
person = schema_instance.node_type_from_label("PERSON")
|
|
|
|
assert person is not None
|
|
assert person.description == "An individual human being."
|
|
assert len(person.properties) == 2
|
|
assert person.additional_properties is False
|
|
|
|
org = schema_instance.node_type_from_label("ORGANIZATION")
|
|
assert org is not None
|
|
assert org.description == "A structured group of people with a common purpose."
|
|
assert org.additional_properties is True
|
|
|
|
age = schema_instance.node_type_from_label("AGE")
|
|
assert age is not None
|
|
assert age.description == "Age of a person in years."
|
|
assert age.additional_properties is True
|
|
|
|
|
|
def test_create_schema_model_missing_relations(
|
|
schema_builder: SchemaBuilder,
|
|
valid_node_types: Tuple[NodeType, ...],
|
|
valid_patterns: tuple[Pattern, ...],
|
|
) -> None:
|
|
with pytest.raises(SchemaValidationError) as exc_info:
|
|
schema_builder.create_schema_model(
|
|
node_types=valid_node_types, patterns=valid_patterns
|
|
)
|
|
assert "Relationship types must also be provided when using patterns." in str(
|
|
exc_info.value
|
|
), "Should fail due to missing relations"
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_llm() -> AsyncMock:
|
|
mock = AsyncMock()
|
|
mock.ainvoke = AsyncMock()
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def valid_schema_json() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"]
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_valid_constraints() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"},
|
|
{"name": "email", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"]
|
|
],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_invalid_constraints() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"]
|
|
],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"},
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "email"},
|
|
{"type": "UNIQUENESS", "node_type": "NonExistentNode", "property_name": "id"},
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": ""}
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_null_constraints() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"]
|
|
],
|
|
"constraints": null
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_required_properties() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING", "required": true},
|
|
{"name": "email", "type": "STRING", "required": false},
|
|
{"name": "phone", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{"label": "KNOWS"}
|
|
],
|
|
"patterns": [
|
|
["Person", "KNOWS", "Person"]
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_string_required_values() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING", "required": "true"},
|
|
{"name": "email", "type": "STRING", "required": "yes"},
|
|
{"name": "phone", "type": "STRING", "required": "false"},
|
|
{"name": "address", "type": "STRING", "required": "no"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{"label": "KNOWS"}
|
|
],
|
|
"patterns": [
|
|
["Person", "KNOWS", "Person"]
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def invalid_schema_json() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
},
|
|
],
|
|
invalid json content
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_from_text(mock_llm: AsyncMock) -> SchemaFromTextExtractor:
|
|
return SchemaFromTextExtractor(llm=mock_llm)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_run_valid_response(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
valid_schema_json: str,
|
|
) -> None:
|
|
# configure the mock LLM to return a valid schema JSON
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)
|
|
|
|
# run the schema extraction
|
|
schema_config = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify the LLM was called with a prompt
|
|
mock_llm.ainvoke.assert_called_once()
|
|
prompt_arg = mock_llm.ainvoke.call_args[0][0]
|
|
assert isinstance(prompt_arg, str)
|
|
assert "Sample text for extraction" in prompt_arg
|
|
|
|
# verify the schema was correctly extracted
|
|
assert len(schema_config.node_types) == 2
|
|
assert schema_config.node_type_from_label("Person") is not None
|
|
assert schema_config.node_type_from_label("Organization") is not None
|
|
|
|
assert schema_config.relationship_types is not None
|
|
assert schema_config.relationship_type_from_label("WORKS_FOR") is not None
|
|
|
|
assert schema_config.patterns is not None
|
|
assert len(schema_config.patterns) == 1
|
|
assert schema_config.patterns[0] == ("Person", "WORKS_FOR", "Organization")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_run_invalid_json(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
invalid_schema_json: str,
|
|
) -> None:
|
|
# configure the mock LLM to return invalid JSON
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=invalid_schema_json)
|
|
|
|
# verify that running with invalid JSON raises a ValueError
|
|
with pytest.raises(SchemaExtractionError) as exc_info:
|
|
await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
assert "not valid JSON" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_custom_template(
|
|
mock_llm: AsyncMock, valid_schema_json: str
|
|
) -> None:
|
|
# create a custom template
|
|
custom_prompt = "This is a custom prompt with text: {text}"
|
|
custom_template = PromptTemplate(template=custom_prompt, expected_inputs=["text"])
|
|
|
|
# create SchemaFromTextExtractor with the custom template
|
|
schema_from_text = SchemaFromTextExtractor(
|
|
llm=mock_llm, prompt_template=custom_template
|
|
)
|
|
|
|
# configure mock LLM to return valid JSON and capture the prompt that was sent to it
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)
|
|
|
|
# run the schema extraction
|
|
await schema_from_text.run(text="Sample text")
|
|
|
|
# verify the custom prompt was passed to the LLM
|
|
prompt_sent_to_llm = mock_llm.ainvoke.call_args[0][0]
|
|
assert "This is a custom prompt with text" in prompt_sent_to_llm
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_llm_params(
|
|
mock_llm: AsyncMock, valid_schema_json: str
|
|
) -> None:
|
|
# configure custom LLM parameters
|
|
llm_params = {"temperature": 0.1, "max_tokens": 500}
|
|
|
|
# create SchemaFromTextExtractor with custom LLM parameters
|
|
schema_from_text = SchemaFromTextExtractor(llm=mock_llm, llm_params=llm_params)
|
|
|
|
# configure the mock LLM to return a valid schema JSON
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)
|
|
|
|
# run the schema extraction
|
|
await schema_from_text.run(text="Sample text")
|
|
|
|
# verify the LLM was called with the custom parameters
|
|
mock_llm.ainvoke.assert_called_once()
|
|
call_kwargs = mock_llm.ainvoke.call_args[1]
|
|
assert call_kwargs["temperature"] == 0.1
|
|
assert call_kwargs["max_tokens"] == 500
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_config_save_json(graph_schema: GraphSchema) -> None:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
# create file path
|
|
json_path = os.path.join(temp_dir, "schema.json")
|
|
|
|
# store the schema config
|
|
graph_schema.save(json_path)
|
|
|
|
# verify the file exists and has content
|
|
assert os.path.exists(json_path)
|
|
assert os.path.getsize(json_path) > 0
|
|
|
|
# verify the content is valid JSON and contains expected data
|
|
with open(json_path, "r") as f:
|
|
data = json.load(f)
|
|
assert "node_types" in data
|
|
assert len(data["node_types"]) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_config_save_yaml(graph_schema: GraphSchema) -> None:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
# Create file path
|
|
yaml_path = os.path.join(temp_dir, "schema.yaml")
|
|
|
|
# Store the schema config
|
|
graph_schema.save(yaml_path)
|
|
|
|
# Verify the file exists and has content
|
|
assert os.path.exists(yaml_path)
|
|
assert os.path.getsize(yaml_path) > 0
|
|
|
|
# Verify the content is valid YAML and contains expected data
|
|
with open(yaml_path, "r") as f:
|
|
data = yaml.safe_load(f)
|
|
assert "node_types" in data
|
|
assert len(data["node_types"]) == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_config_from_file(graph_schema: GraphSchema) -> None:
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
# create file paths with different extensions
|
|
json_path = os.path.join(temp_dir, "schema.json")
|
|
yaml_path = os.path.join(temp_dir, "schema.yaml")
|
|
yml_path = os.path.join(temp_dir, "schema.yml")
|
|
|
|
# store the schema config in the different formats
|
|
graph_schema.save(json_path)
|
|
graph_schema.save(yaml_path)
|
|
graph_schema.save(yml_path)
|
|
|
|
# load using from_file which should detect the format based on extension
|
|
json_schema = GraphSchema.from_file(json_path)
|
|
yaml_schema = GraphSchema.from_file(yaml_path)
|
|
yml_schema = GraphSchema.from_file(yml_path)
|
|
|
|
# simple verification that the objects were loaded correctly
|
|
assert isinstance(json_schema, GraphSchema)
|
|
assert isinstance(yaml_schema, GraphSchema)
|
|
assert isinstance(yml_schema, GraphSchema)
|
|
|
|
# verify basic structure is intact
|
|
assert "node_types" in json_schema.model_dump()
|
|
assert "node_types" in yaml_schema.model_dump()
|
|
assert "node_types" in yml_schema.model_dump()
|
|
|
|
# verify an unsupported extension raises the correct error
|
|
txt_path = os.path.join(temp_dir, "schema.txt")
|
|
graph_schema.save(
|
|
txt_path, format=FileFormat.JSON
|
|
) # Store as JSON but with .txt extension
|
|
|
|
with pytest.raises(ValueError, match="Unsupported file format: None"):
|
|
GraphSchema.from_file(txt_path)
|
|
|
|
|
|
@pytest.fixture
|
|
def valid_schema_json_array() -> str:
|
|
return """
|
|
[
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"]
|
|
]
|
|
}
|
|
]
|
|
"""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_run_valid_json_array(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
valid_schema_json_array: str,
|
|
) -> None:
|
|
# configure the mock LLM to return a valid JSON array
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json_array)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify the schema was correctly extracted from the array
|
|
assert len(schema.node_types) == 2
|
|
assert schema.node_type_from_label("Person") is not None
|
|
assert schema.node_type_from_label("Organization") is not None
|
|
|
|
assert schema.relationship_types is not None
|
|
assert schema.relationship_type_from_label("WORKS_FOR") is not None
|
|
|
|
assert schema.patterns is not None
|
|
assert len(schema.patterns) == 1
|
|
assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_invalid_node_patterns() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"],
|
|
["Person", "WORKS_FOR", "UndefinedNode"],
|
|
["UndefinedNode", "WORKS_FOR", "Organization"]
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_invalid_relationship_patterns() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"],
|
|
["Person", "UNDEFINED_RELATION", "Organization"],
|
|
["Organization", "ANOTHER_UNDEFINED_RELATION", "Person"]
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_nodes_without_labels() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
"Organization",
|
|
"",
|
|
"Company",
|
|
"Invalid description with spaces",
|
|
"{\\"invalid\\": \\"json object\\"}"
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
}
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"]
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_relationships_without_labels() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
},
|
|
{
|
|
"label": "Organization",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [
|
|
{
|
|
"label": "WORKS_FOR",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
},
|
|
{
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
},
|
|
{
|
|
"label": "",
|
|
"properties": [
|
|
{"name": "since", "type": "DATE"}
|
|
]
|
|
},
|
|
"MANAGES",
|
|
"",
|
|
"SUPERVISES",
|
|
"invalid relationship description",
|
|
"{\\"invalid\\": \\"json\\"}"
|
|
],
|
|
"patterns": [
|
|
["Person", "WORKS_FOR", "Organization"],
|
|
["Person", "MANAGES", "Organization"]
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def schema_json_with_nonexistent_property_constraint() -> str:
|
|
return """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING"}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [],
|
|
"patterns": [],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"},
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "nonexistent_property"}
|
|
]
|
|
}
|
|
"""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_filters_invalid_node_patterns(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_invalid_node_patterns: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with invalid node patterns
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_invalid_node_patterns
|
|
)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify that invalid node patterns were filtered out (2 out of 3 patterns should be removed)
|
|
assert schema.patterns is not None
|
|
assert len(schema.patterns) == 1
|
|
assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_filters_invalid_relationship_patterns(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_invalid_relationship_patterns: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with invalid relationship patterns
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_invalid_relationship_patterns
|
|
)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify that invalid relationship patterns were filtered out (2 out of 3 patterns should be removed)
|
|
assert schema.patterns is not None
|
|
assert len(schema.patterns) == 1
|
|
assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_filters_nodes_without_labels(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_nodes_without_labels: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with nodes without labels
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_nodes_without_labels
|
|
)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify that nodes without labels were filtered out (5 out of 8 nodes should be removed)
|
|
assert len(schema.node_types) == 3
|
|
assert schema.node_type_from_label("Person") is not None
|
|
assert schema.node_type_from_label("Organization") is not None
|
|
assert schema.node_type_from_label("Company") is not None
|
|
|
|
# verify that the pattern is still valid with the remaining nodes
|
|
assert schema.patterns is not None
|
|
assert len(schema.patterns) == 1
|
|
assert schema.patterns[0] == ("Person", "WORKS_FOR", "Organization")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_filters_relationships_without_labels(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_relationships_without_labels: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with relationships without labels
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_relationships_without_labels
|
|
)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify that relationships without labels were filtered out (5 out of 8 relationships should be removed)
|
|
assert schema.relationship_types is not None
|
|
assert len(schema.relationship_types) == 3
|
|
assert schema.relationship_type_from_label("WORKS_FOR") is not None
|
|
assert schema.relationship_type_from_label("MANAGES") is not None
|
|
assert schema.relationship_type_from_label("SUPERVISES") is not None
|
|
|
|
# verify that the patterns are still valid with the remaining relationships
|
|
assert schema.patterns is not None
|
|
assert len(schema.patterns) == 2
|
|
assert ("Person", "WORKS_FOR", "Organization") in schema.patterns
|
|
assert ("Person", "MANAGES", "Organization") in schema.patterns
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_with_valid_constraints(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_valid_constraints: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with valid constraints
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_valid_constraints
|
|
)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].type == "UNIQUENESS"
|
|
assert schema.constraints[0].node_type == "Person"
|
|
assert schema.constraints[0].property_names == ("name",)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_filters_invalid_constraints(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_invalid_constraints: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with invalid constraints
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_invalid_constraints
|
|
)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify that invalid constraints were filtered out:
|
|
# constraints with NonExistentNode should be removed
|
|
# constraint with empty property_name should be removed
|
|
# only the valid constraint should remain
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].node_type == "Person"
|
|
assert schema.constraints[0].property_names == ("name",)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_filters_constraint_with_nonexistent_property(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_nonexistent_property_constraint: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with constraint on nonexistent property
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_nonexistent_property_constraint
|
|
)
|
|
|
|
# run the schema extraction
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify that only the valid constraint (with "name" property) remains
|
|
# the constraint with "nonexistent_property" should be filtered out
|
|
assert len(schema.constraints) == 1
|
|
assert schema.constraints[0].property_names == ("name",)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_handles_null_constraints(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_null_constraints: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema with null constraints
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_null_constraints
|
|
)
|
|
|
|
# run the schema extraction - should not crash
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify schema was created with empty constraints
|
|
assert len(schema.constraints) == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_handles_missing_constraints(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
valid_schema_json: str,
|
|
) -> None:
|
|
# configure the mock LLM to return schema without constraints field
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)
|
|
|
|
# run the schema extraction - should not crash
|
|
schema = await schema_from_text.run(text="Sample text for extraction")
|
|
|
|
# verify schema was created with empty constraints
|
|
assert len(schema.constraints) == 0
|
|
|
|
|
|
def test_clean_json_content_markdown_with_json_language(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
) -> None:
|
|
content = """```json
|
|
{"node_types": [{"label": "Person"}]}
|
|
```"""
|
|
|
|
cleaned = schema_from_text._clean_json_content(content)
|
|
assert cleaned == '{"node_types": [{"label": "Person"}]}'
|
|
|
|
|
|
def test_clean_json_content_markdown_without_language(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
) -> None:
|
|
content = """```
|
|
{"node_types": [{"label": "Person"}]}
|
|
```"""
|
|
|
|
cleaned = schema_from_text._clean_json_content(content)
|
|
assert cleaned == '{"node_types": [{"label": "Person"}]}'
|
|
|
|
|
|
def test_clean_json_content_plain_json(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
) -> None:
|
|
content = '{"node_types": [{"label": "Person"}]}'
|
|
|
|
cleaned = schema_from_text._clean_json_content(content)
|
|
assert cleaned == '{"node_types": [{"label": "Person"}]}'
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_with_required_properties(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_required_properties: str,
|
|
) -> None:
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_required_properties
|
|
)
|
|
|
|
schema = await schema_from_text.run(text="Sample text for test")
|
|
|
|
person = schema.node_type_from_label("Person")
|
|
assert person is not None
|
|
|
|
# Legacy required:true migrates to EXISTENCE; flags are cleared on PropertyType
|
|
assert schema.existence_property_names_for_node("Person") == {"name"}
|
|
name_prop = next((p for p in person.properties if p.name == "name"), None)
|
|
email_prop = next((p for p in person.properties if p.name == "email"), None)
|
|
phone_prop = next((p for p in person.properties if p.name == "phone"), None)
|
|
|
|
assert name_prop is not None and name_prop.required is False
|
|
assert email_prop is not None and email_prop.required is False
|
|
assert phone_prop is not None and phone_prop.required is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_string_required_coerced_without_existence_migration(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
schema_json_with_string_required_values: str,
|
|
) -> None:
|
|
"""LLMs may emit string truthiness for ``required``; Pydantic coerces it on ``PropertyType``.
|
|
|
|
Migration of legacy ``required`` to ``EXISTENCE`` constraints only runs when the raw
|
|
property dict has ``required is True`` (JSON boolean). String values such as ``\"true\"``
|
|
are not migrated; add ``ConstraintType`` rows with type ``EXISTENCE`` on ``GraphSchema``
|
|
if you need existence semantics for that case.
|
|
"""
|
|
mock_llm.ainvoke.return_value = LLMResponse(
|
|
content=schema_json_with_string_required_values
|
|
)
|
|
|
|
schema = await schema_from_text.run(text="Sample text for test")
|
|
|
|
person = schema.node_type_from_label("Person")
|
|
assert person is not None
|
|
|
|
assert schema.existence_property_names_for_node("Person") == set()
|
|
name_prop = next((p for p in person.properties if p.name == "name"), None)
|
|
email_prop = next((p for p in person.properties if p.name == "email"), None)
|
|
assert name_prop is not None
|
|
assert email_prop is not None
|
|
assert name_prop.model_dump().get("required") is True
|
|
assert email_prop.model_dump().get("required") is True
|
|
|
|
phone_prop = next((p for p in person.properties if p.name == "phone"), None)
|
|
address_prop = next((p for p in person.properties if p.name == "address"), None)
|
|
assert phone_prop is not None
|
|
assert address_prop is not None
|
|
assert phone_prop.model_dump().get("required") is False
|
|
assert address_prop.model_dump().get("required") is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_handles_missing_required_field(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
valid_schema_json: str,
|
|
) -> None:
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=valid_schema_json)
|
|
|
|
schema = await schema_from_text.run(text="Sample text")
|
|
|
|
person = schema.node_type_from_label("Person")
|
|
assert person is not None
|
|
|
|
# All properties should have required=False (default)
|
|
for prop in person.properties:
|
|
assert prop.required is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schema_from_text_uniqueness_does_not_force_required_property(
|
|
schema_from_text: SchemaFromTextExtractor,
|
|
mock_llm: AsyncMock,
|
|
) -> None:
|
|
schema_json = """
|
|
{
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [
|
|
{"name": "name", "type": "STRING", "required": false},
|
|
{"name": "email", "type": "STRING", "required": false}
|
|
]
|
|
}
|
|
],
|
|
"relationship_types": [],
|
|
"patterns": [],
|
|
"constraints": [
|
|
{"type": "UNIQUENESS", "node_type": "Person", "property_name": "name"}
|
|
]
|
|
}
|
|
"""
|
|
mock_llm.ainvoke.return_value = LLMResponse(content=schema_json)
|
|
|
|
schema = await schema_from_text.run(text="Sample text")
|
|
|
|
person = schema.node_type_from_label("Person")
|
|
assert person is not None
|
|
|
|
name_prop = next((p for p in person.properties if p.name == "name"), None)
|
|
email_prop = next((p for p in person.properties if p.name == "email"), None)
|
|
|
|
assert name_prop is not None and name_prop.required is False
|
|
assert email_prop is not None and email_prop.required is False
|
|
assert schema.existence_property_names_for_node("Person") == set()
|
|
assert schema.constraints[0].type == "UNIQUENESS"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.experimental.components.schema.get_structured_schema")
|
|
async def test_schema_from_existing_graph(mock_get_structured_schema: Mock) -> None:
|
|
mock_get_structured_schema.return_value = {
|
|
"node_props": {
|
|
"Person": [
|
|
{"property": "id", "type": "INTEGER"},
|
|
{"property": "name", "type": "STRING"},
|
|
]
|
|
},
|
|
"rel_props": {"KNOWS": [{"property": "fromDate", "type": "DATE"}]},
|
|
"relationships": [
|
|
{"start": "Person", "type": "KNOWS", "end": "Person"},
|
|
{"start": "Person", "type": "LIVES_IN", "end": "City"},
|
|
],
|
|
"metadata": {
|
|
"constraint": [
|
|
{
|
|
"id": 7,
|
|
"name": "person_id",
|
|
"type": "NODE_PROPERTY_EXISTENCE",
|
|
"entityType": "NODE",
|
|
"labelsOrTypes": ["Person"],
|
|
"properties": ["id"],
|
|
"ownedIndex": "person_id",
|
|
"propertyType": None,
|
|
},
|
|
],
|
|
"index": [
|
|
{
|
|
"label": "Person",
|
|
"properties": ["name"],
|
|
"size": 2,
|
|
"type": "RANGE",
|
|
"valuesSelectivity": 1.0,
|
|
"distinctValues": 2.0,
|
|
},
|
|
],
|
|
},
|
|
}
|
|
driver = Mock()
|
|
schema_builder = SchemaFromExistingGraphExtractor(
|
|
driver=driver,
|
|
)
|
|
schema = await schema_builder.run()
|
|
assert isinstance(schema, GraphSchema)
|
|
assert len(schema.node_types) == 2
|
|
person_node_type = schema.node_type_from_label("Person")
|
|
assert person_node_type is not None
|
|
id_person_property = [p for p in person_node_type.properties if p.name == "id"][0]
|
|
assert id_person_property.required is False
|
|
assert schema.existence_property_names_for_node("Person") == {"id"}
|
|
assert any(
|
|
c.type == "EXISTENCE"
|
|
and c.node_type == "Person"
|
|
and c.property_names == ("id",)
|
|
for c in schema.constraints
|
|
)
|
|
assert person_node_type.additional_properties is False
|
|
city_node_type = schema.node_type_from_label("City")
|
|
assert city_node_type is not None
|
|
assert city_node_type.additional_properties is True
|
|
|
|
assert len(schema.relationship_types) == 2
|
|
knows_rel = schema.relationship_type_from_label("KNOWS")
|
|
assert knows_rel is not None
|
|
assert knows_rel.additional_properties is False
|
|
lives_rel = schema.relationship_type_from_label("LIVES_IN")
|
|
assert lives_rel is not None
|
|
assert lives_rel.additional_properties is True
|
|
|
|
assert schema.additional_node_types is False
|
|
assert schema.additional_relationship_types is False
|
|
assert schema.additional_patterns is False
|
|
|
|
assert schema.patterns == (
|
|
Pattern(source="Person", relationship="KNOWS", target="Person"),
|
|
Pattern(source="Person", relationship="LIVES_IN", target="City"),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@patch("neo4j_graphrag.experimental.components.schema.get_structured_schema")
|
|
async def test_schema_from_existing_graph_additional_params(
|
|
mock_get_structured_schema: Mock,
|
|
) -> None:
|
|
mock_get_structured_schema.return_value = {
|
|
"node_props": {
|
|
"Person": [
|
|
{"property": "id", "type": "INTEGER"},
|
|
{"property": "name", "type": "STRING"},
|
|
]
|
|
},
|
|
"rel_props": {"KNOWS": [{"property": "fromDate", "type": "DATE"}]},
|
|
"relationships": [
|
|
{"start": "Person", "type": "KNOWS", "end": "Person"},
|
|
{"start": "Person", "type": "LIVES_IN", "end": "City"},
|
|
],
|
|
}
|
|
driver = Mock()
|
|
schema_builder = SchemaFromExistingGraphExtractor(
|
|
driver=driver,
|
|
additional_node_types=True,
|
|
additional_relationship_types=True,
|
|
additional_patterns=True,
|
|
additional_properties=True,
|
|
)
|
|
schema = await schema_builder.run()
|
|
person_node_type = schema.node_type_from_label("Person")
|
|
assert person_node_type is not None
|
|
assert person_node_type.additional_properties is True
|
|
city_node_type = schema.node_type_from_label("City")
|
|
assert city_node_type is not None
|
|
assert city_node_type.additional_properties is True
|
|
|
|
knows_rel = schema.relationship_type_from_label("KNOWS")
|
|
assert knows_rel is not None
|
|
assert knows_rel.additional_properties is True
|
|
lives_rel = schema.relationship_type_from_label("LIVES_IN")
|
|
assert lives_rel is not None
|
|
assert lives_rel.additional_properties is True
|
|
|
|
assert schema.additional_node_types is True
|
|
assert schema.additional_relationship_types is True
|
|
assert schema.additional_patterns is True
|
|
|
|
|
|
def test_extract_graph_constraints_from_metadata_node_key_maps_to_key() -> None:
|
|
"""Neo4j ``NODE_KEY`` metadata must become ``GraphConstraintType.KEY``, not EXISTENCE."""
|
|
structured_schema: dict[str, Any] = {
|
|
"metadata": {
|
|
"constraint": [
|
|
{
|
|
"type": "NODE_KEY",
|
|
"labelsOrTypes": ["Person"],
|
|
"properties": ["email"],
|
|
}
|
|
]
|
|
}
|
|
}
|
|
out = SchemaFromExistingGraphExtractor._extract_graph_constraints_from_metadata(
|
|
structured_schema
|
|
)
|
|
assert out == [
|
|
{
|
|
"type": GraphConstraintType.KEY.value,
|
|
"node_type": "Person",
|
|
"property_name": "email",
|
|
"property_names": ("email",),
|
|
"relationship_type": None,
|
|
}
|
|
]
|
|
|
|
|
|
def test_extract_graph_constraints_from_metadata_relationship_key_maps_to_key() -> None:
|
|
"""Neo4j ``RELATIONSHIP_KEY`` metadata maps to relationship-scoped ``KEY`` constraints."""
|
|
structured_schema: dict[str, Any] = {
|
|
"metadata": {
|
|
"constraint": [
|
|
{
|
|
"type": "RELATIONSHIP_KEY",
|
|
"labelsOrTypes": ["WORKS_FOR"],
|
|
"properties": ["since"],
|
|
}
|
|
]
|
|
}
|
|
}
|
|
out = SchemaFromExistingGraphExtractor._extract_graph_constraints_from_metadata(
|
|
structured_schema
|
|
)
|
|
assert out == [
|
|
{
|
|
"type": GraphConstraintType.KEY.value,
|
|
"node_type": "",
|
|
"property_name": "since",
|
|
"property_names": ("since",),
|
|
"relationship_type": "WORKS_FOR",
|
|
}
|
|
]
|
|
|
|
|
|
def test_graph_schema_extraction_output_json_schema_lean_root() -> None:
|
|
"""Structured-output schema must not include pipeline-only GraphSchema flags."""
|
|
from neo4j_graphrag.experimental.components.graph_schema_extraction import (
|
|
GraphSchemaExtractionOutput,
|
|
)
|
|
|
|
raw = GraphSchemaExtractionOutput.model_json_schema()
|
|
dumped = json.dumps(raw)
|
|
assert "additional_node_types" not in dumped
|
|
assert "additional_relationship_types" not in dumped
|
|
assert "additional_patterns" not in dumped
|
|
|
|
|
|
def test_graph_schema_extraction_constraint_schema_avoids_null_type_for_vertex() -> (
|
|
None
|
|
):
|
|
"""Vertex AI maps JSON Schema to protobuf and rejects ``type: \"null\"`` (e.g. ``Optional``)."""
|
|
from neo4j_graphrag.experimental.components.graph_schema_extraction import (
|
|
GraphSchemaExtractionOutput,
|
|
)
|
|
|
|
raw = GraphSchemaExtractionOutput.model_json_schema()
|
|
ect = (raw.get("$defs") or {}).get("ExtractedConstraintType")
|
|
assert ect is not None
|
|
rel_schema = (ect.get("properties") or {}).get("relationship_type", {})
|
|
assert "anyOf" not in rel_schema
|
|
assert '"type": "null"' not in json.dumps(rel_schema)
|
|
|
|
|
|
def test_graph_schema_from_extraction_output() -> None:
|
|
from neo4j_graphrag.experimental.components.graph_schema_extraction import (
|
|
ExtractedConstraintType,
|
|
ExtractedNodeType,
|
|
ExtractedPropertyType,
|
|
GraphSchemaExtractionOutput,
|
|
)
|
|
|
|
dto = GraphSchemaExtractionOutput(
|
|
node_types=[
|
|
ExtractedNodeType(
|
|
label="Person",
|
|
properties=[
|
|
ExtractedPropertyType(name="name", type="STRING"),
|
|
],
|
|
)
|
|
],
|
|
relationship_types=[],
|
|
patterns=[],
|
|
constraints=[
|
|
ExtractedConstraintType(
|
|
type="UNIQUENESS",
|
|
node_type="Person",
|
|
property_names=["name"],
|
|
),
|
|
ExtractedConstraintType(
|
|
type="EXISTENCE",
|
|
node_type="Person",
|
|
property_names=["name"],
|
|
relationship_type="",
|
|
),
|
|
],
|
|
)
|
|
gs = GraphSchema.from_extraction_output(dto)
|
|
assert gs.node_types[0].label == "Person"
|
|
assert gs.node_types[0].properties[0].required is False
|
|
assert gs.existence_property_names_for_node("Person") == {"name"}
|
|
assert {c.property_names for c in gs.constraints} == {("name",)}
|
|
assert gs.additional_node_types is False
|
|
|
|
|
|
def test_validate_extraction_dict_to_graph_schema() -> None:
|
|
from neo4j_graphrag.experimental.components.schema import (
|
|
validate_extraction_dict_to_graph_schema,
|
|
)
|
|
|
|
d = {
|
|
"node_types": [
|
|
{
|
|
"label": "Person",
|
|
"properties": [{"name": "name", "type": "STRING", "required": False}],
|
|
}
|
|
],
|
|
"relationship_types": [],
|
|
"patterns": [],
|
|
"constraints": [],
|
|
}
|
|
gs = validate_extraction_dict_to_graph_schema(d)
|
|
assert len(gs.node_types) == 1
|
|
assert gs.node_types[0].label == "Person"
|