from morphik.models import EntityExtractionPromptOverride, EntityExtractionExample, GraphPromptOverrides
# Update with custom entity extraction examples
updated_graph = db.update_graph(
name="medical_graph",
additional_filters={"category": "new_medical_data"},
prompt_overrides=GraphPromptOverrides(
entity_extraction=EntityExtractionPromptOverride(
examples=[
EntityExtractionExample(label="Insulin", type="MEDICATION"),
EntityExtractionExample(label="Diabetes", type="CONDITION"),
EntityExtractionExample(label="Heart rate", type="VITAL_SIGN"),
EntityExtractionExample(label="Cardiology", type="SPECIALTY")
]
)
)
)
# With custom entity extraction template
updated_graph = db.update_graph(
name="legal_graph",
additional_documents=["contract1", "contract2"],
prompt_overrides=GraphPromptOverrides(
entity_extraction=EntityExtractionPromptOverride(
prompt_template="""Extract legal entities from the following document:
{content}
Focus on these types of entities:
{examples}
Return the extracted entities in JSON format with the following structure:
[
{"label": "entity name", "type": "ENTITY_TYPE", "properties": {"key": "value"}}
]
""",
examples=[
EntityExtractionExample(
label="John Smith",
type="PERSON",
properties={"role": "Plaintiff"}
),
EntityExtractionExample(
label="Acme Corporation",
type="ORGANIZATION",
properties={"type": "Corporation"}
),
EntityExtractionExample(
label="January 15, 2025",
type="DATE"
)
]
)
)
)