@@ -243,9 +243,11 @@ async def test_extractor_no_schema_enforcement() -> None:
243
243
create_lexical_graph = False ,
244
244
enforce_schema = SchemaEnforcementMode .NONE )
245
245
246
- schema = SchemaConfig (entities = {"Person" : {"name" : "STRING" }},
247
- relations = {},
248
- potential_schema = [])
246
+ schema = SchemaConfig (
247
+ entities = {"Person" : {"label" : "Person" ,
248
+ "properties" : [{"name" : "name" , "type" : "STRING" }]}},
249
+ relations = {},
250
+ potential_schema = [])
249
251
250
252
chunks = TextChunks (chunks = [TextChunk (text = "some text" , index = 0 )])
251
253
@@ -290,9 +292,11 @@ async def test_extractor_schema_enforcement_invalid_nodes():
290
292
create_lexical_graph = False ,
291
293
enforce_schema = SchemaEnforcementMode .STRICT )
292
294
293
- schema = SchemaConfig (entities = {"Person" : {"name" : "STRING" }},
294
- relations = {},
295
- potential_schema = [])
295
+ schema = SchemaConfig (
296
+ entities = {"Person" : {"label" : "Person" ,
297
+ "properties" : [{"name" : "name" , "type" : "STRING" }]}},
298
+ relations = {},
299
+ potential_schema = [])
296
300
297
301
chunks = TextChunks (chunks = [TextChunk (text = "some text" , index = 0 )])
298
302
@@ -316,9 +320,12 @@ async def test_extraction_schema_enforcement_invalid_node_properties():
316
320
create_lexical_graph = False ,
317
321
enforce_schema = SchemaEnforcementMode .STRICT )
318
322
319
- schema = SchemaConfig (entities = {"Person" : {"name" : str , "age" : int }},
320
- relations = {},
321
- potential_schema = [])
323
+ schema = SchemaConfig (
324
+ entities = {"Person" : {"label" : "Person" ,
325
+ "properties" : [{"name" : "name" , "type" : "STRING" },
326
+ {"name" : "age" , "type" : "INTEGER" }]}},
327
+ relations = {},
328
+ potential_schema = [])
322
329
323
330
chunks = TextChunks (chunks = [TextChunk (text = "some text" , index = 0 )])
324
331
@@ -342,7 +349,7 @@ async def test_extractor_schema_enforcement_valid_nodes_with_empty_props():
342
349
create_lexical_graph = False ,
343
350
enforce_schema = SchemaEnforcementMode .STRICT )
344
351
345
- schema = SchemaConfig (entities = {"Person" : {}},
352
+ schema = SchemaConfig (entities = {"Person" : {"label" : "Person" , "properties" : [] }},
346
353
relations = {},
347
354
potential_schema = [])
348
355
@@ -368,9 +375,11 @@ async def test_extractor_schema_enforcement_invalid_relations_wrong_types():
368
375
create_lexical_graph = False ,
369
376
enforce_schema = SchemaEnforcementMode .STRICT )
370
377
371
- schema = SchemaConfig (entities = {"Person" : {"name" : str }},
372
- relations = {"LIKES" : {}},
373
- potential_schema = [])
378
+ schema = SchemaConfig (
379
+ entities = {"Person" : {"label" : "Person" ,
380
+ "properties" : [{"name" : "name" , "type" : "STRING" }]}},
381
+ relations = {"LIKES" : {"label" : "LIKES" , "properties" : []}},
382
+ potential_schema = [])
374
383
375
384
chunks = TextChunks (chunks = [TextChunk (text = "some text" , index = 0 )])
376
385
@@ -395,9 +404,13 @@ async def test_extractor_schema_enforcement_invalid_relations_wrong_start_node()
395
404
create_lexical_graph = False ,
396
405
enforce_schema = SchemaEnforcementMode .STRICT )
397
406
398
- schema = SchemaConfig (entities = {"Person" : {"name" : str }, "City" : {"name" : str }},
399
- relations = {"LIVES_IN" : {}},
400
- potential_schema = [("Person" , "LIVES_IN" , "City" )])
407
+ schema = SchemaConfig (
408
+ entities = {"Person" : {"label" : "Person" ,
409
+ "properties" : [{"name" : "name" , "type" : "STRING" }]},
410
+ "City" : {"label" : "City" ,
411
+ "properties" : [{"name" : "name" , "type" : "STRING" }]}},
412
+ relations = {"LIVES_IN" : {"label" : "LIVES_IN" , "properties" : []}},
413
+ potential_schema = [("Person" , "LIVES_IN" , "City" )])
401
414
402
415
chunks = TextChunks (chunks = [TextChunk (text = "some text" , index = 0 )])
403
416
@@ -422,8 +435,10 @@ async def test_extractor_schema_enforcement_invalid_relation_properties():
422
435
enforce_schema = SchemaEnforcementMode .STRICT )
423
436
424
437
schema = SchemaConfig (
425
- entities = {"Person" : {"name" : str }},
426
- relations = {"LIKES" : {"strength" : str }},
438
+ entities = {"Person" : {"label" : "Person" ,
439
+ "properties" : [{"name" : "name" , "type" : "STRING" }]}},
440
+ relations = {"LIKES" : {"label" : "LIKES" ,
441
+ "properties" : [{"name" : "strength" , "type" : "STRING" }]}},
427
442
potential_schema = []
428
443
)
429
444
@@ -452,9 +467,11 @@ async def test_extractor_schema_enforcement_removed_relation_start_end_nodes():
452
467
create_lexical_graph = False ,
453
468
enforce_schema = SchemaEnforcementMode .STRICT )
454
469
455
- schema = SchemaConfig (entities = {"Person" : {"name" : str }},
456
- relations = {"LIKES" : {}},
457
- potential_schema = [("Person" , "LIKES" , "Person" )])
470
+ schema = SchemaConfig (
471
+ entities = {"Person" : {"label" : "Person" ,
472
+ "properties" : [{"name" : "name" , "type" : "STRING" }]}},
473
+ relations = {"LIKES" : {"label" : "LIKES" , "properties" : []}},
474
+ potential_schema = [("Person" , "LIKES" , "Person" )])
458
475
459
476
chunks = TextChunks (chunks = [TextChunk (text = "some text" , index = 0 )])
460
477
0 commit comments