Entity Extraction: Domain Schemas

from neo4j_agent_memory.schema import EntitySchemaConfig, EntityTypeConfig

insurance_schema = EntitySchemaConfig(
    name="insurance",
    version="1.0",
    description="Schema for insurance industry context graphs",
    entity_types=[
        EntityTypeConfig(
            name="POLICYHOLDER",
            description="Insurance policy holder or applicant",
            examples=["John Smith", "Acme Corporation"],
        ),
        EntityTypeConfig(
            name="POLICY",
            description="Insurance policy with number",
            examples=["Policy #INS-2024-001", "Auto Policy 12345"],
        ),
        EntityTypeConfig(
            name="CLAIM",
            description="Insurance claim reference",
            examples=["Claim #CLM-98765", "accident claim"],
        ),
        EntityTypeConfig(
            name="COVERAGE",
            description="Type of insurance coverage",
            examples=["liability coverage", "comprehensive", "collision"],
        ),
        EntityTypeConfig(
            name="PREMIUM",
            description="Insurance premium amount",
            examples=["$500/month", "annual premium of $6,000"],
        ),
        EntityTypeConfig(
            name="VEHICLE",
            description="Insured vehicle",
            examples=["2024 Toyota Camry", "Honda Accord"],
        ),
    ],
)

extractor = GLiNEREntityExtractor.for_schema(insurance_schema)

from neo4j_agent_memory.schema import get_schema, EntityTypeConfig

base_schema = get_schema("ecommerce")

custom_types = [
    EntityTypeConfig(
        name="LOYALTY_TIER",
        description="Customer loyalty program tier",
        examples=["Gold member", "Platinum status", "VIP"],
    ),
    EntityTypeConfig(
        name="SUBSCRIPTION",
        description="Subscription service",
        examples=["Prime membership", "monthly box subscription"],
    ),
]

extended_schema = base_schema.extend(
    name="ecommerce_extended",
    additional_types=custom_types,
)

extractor = GLiNEREntityExtractor.for_schema(extended_schema)

from neo4j_agent_memory.schema import SchemaManager

manager = SchemaManager(client)

stored = await manager.save_schema(
    insurance_schema,
    created_by="admin",
    set_active=True,
)

print(f"Schema saved with ID: {stored.id}")

# Load schema in another session
loaded = await manager.load_schema("insurance")
extractor = GLiNEREntityExtractor.for_schema(loaded.config)

from neo4j_agent_memory.extraction import GLiNERWithRelationsExtractor

extractor = GLiNERWithRelationsExtractor.for_schema("ecommerce")

text = """
    Jane Doe purchased Nike Air Max 90 from our Manhattan store.
    The product was manufactured by Nike and shipped via FedEx.
"""

result = await extractor.extract(text)

for rel in result.relations:
    print(f"  ({rel.source}) -[:{rel.type}]-> ({rel.target})")
    # (Jane Doe) -[:PURCHASED]-> (Nike Air Max 90)
    # (Nike Air Max 90) -[:SOLD_AT]-> (Manhattan store)
    # (Nike Air Max 90) -[:MANUFACTURED_BY]-> (Nike)
    # (Nike Air Max 90) -[:SHIPPED_VIA]-> (FedEx)

financial_relations = [
    {"name": "ADVISES",      "description": "Financial advisor advises client"},
    {"name": "HOLDS",        "description": "Account holds security position"},
    {"name": "TRADED",       "description": "Executed trade in security"},
    {"name": "SUBSIDIARY_OF","description": "Company is subsidiary of parent"},
    {"name": "CUSTODIED_AT", "description": "Assets custodied at institution"},
]

extractor = GLiNERWithRelationsExtractor(
    entity_schema="financial",
    relation_types=financial_relations,
)