Free form Knowledge Graph Extraction

import asyncio
from typing import List

from dotenv import load_dotenv

import synalinks

# =============================================================================
# Generic (free-form) schema: open labels, no Literal discriminators.
# =============================================================================


class Document(synalinks.DataModel):
    """A piece of unstructured text to extract a graph from."""

    text: str = synalinks.Field(description="The raw document text")


class Node(synalinks.Entity):
    # `label` is inherited from Entity as a plain `str`. The model fills
    # it with the entity TYPE it discovers ("Person", "City", ...).
    name: str = synalinks.Field(description="The entity's name / identifier.")
    description: str = synalinks.Field(description="A short description of the entity.")


class Edge(synalinks.Relation):
    # `label` inherited as `str` (the relation type); generic endpoints.
    subj: Node = synalinks.Field(description="The source entity.")
    obj: Node = synalinks.Field(description="The target entity.")


class Graph(synalinks.KnowledgeGraph):
    entities: List[Node] = synalinks.Field(
        description="Every entity mentioned in the text, with a discovered label.",
    )
    relations: List[Edge] = synalinks.Field(
        description="Every relation between the entities, with a discovered label.",
    )


# =============================================================================
# Main Demonstration
# =============================================================================


async def main():
    load_dotenv()
    synalinks.clear_session()

    language_model = synalinks.LanguageModel(model="ollama/llama3.2:latest")
    embedding_model = synalinks.EmbeddingModel(model="ollama/mxbai-embed-large")

    document = Document(
        text=(
            "Marie Curie was a physicist and chemist born in Warsaw. She "
            "conducted pioneering research on radioactivity at the University "
            "of Paris, and won the Nobel Prize in Physics in 1903 together "
            "with Pierre Curie and Henri Becquerel."
        ),
    )

    # The knowledge base declares NO entity_models / relation_models — the
    # labels aren't known until the model discovers them, and the store
    # creates the tables on demand.
    knowledge_base = synalinks.KnowledgeBase(
        graph_uri="ladybug://:memory:",
        embedding_model=embedding_model,
        metric="cosine",
        wipe_on_start=True,
    )

    # -------------------------------------------------------------------------
    # Extract a free-form graph, embed it, and store it
    # -------------------------------------------------------------------------
    print("=" * 60)
    print("Free-form extraction")
    print("=" * 60)

    inputs = synalinks.Input(data_model=Document)
    graph = await synalinks.Generator(
        data_model=Graph,
        language_model=language_model,
        instructions=(
            "Extract every entity and every relation between them from the "
            "text. Pick a PascalCase label for each entity type (e.g. "
            "'Person', 'City', 'Award') and a concise PascalCase label "
            "for each relation (e.g. 'BornIn', 'Won')."
        ),
    )(inputs)
    embedded = await synalinks.EmbedKnowledge(
        embedding_model=embedding_model,
        in_mask=["name", "description"],
    )(graph)
    stored = await synalinks.UpdateKnowledge(
        knowledge_base=knowledge_base,
    )(embedded)

    program = synalinks.Program(
        inputs=inputs,
        outputs=stored,
        name="free_form_kg_extraction",
        description="Extract an open-label knowledge graph and store it.",
    )

    await program(document)

    # -------------------------------------------------------------------------
    # Introspect: which labels did the model actually invent?
    # -------------------------------------------------------------------------
    print("\n" + "=" * 60)
    print("Discovered entity types")
    print("=" * 60)

    types = await knowledge_base.cypher(
        "MATCH (n) RETURN DISTINCT label(n) AS type ORDER BY type"
    )
    for row in types:
        print(f"  - {row['type']}")

    print("\n" + "=" * 60)
    print("Discovered relations")
    print("=" * 60)

    edges = await knowledge_base.cypher(
        "MATCH (a)-[r]->(b) RETURN a.name AS subj, label(r) AS rel, b.name AS obj"
    )
    for row in edges:
        print(f"  - ({row['subj']}) -[{row['rel']}]-> ({row['obj']})")

    # -------------------------------------------------------------------------
    # Query: the neighbourhood around an entity (properties preserved)
    # -------------------------------------------------------------------------
    print("\n" + "=" * 60)
    print("Local graph search around 'Marie Curie'")
    print("=" * 60)

    subgraph = await knowledge_base.local_graph_search(
        "Marie Curie",
        label="Person",
        max_hops=2,
        k=1,
    )
    for entity in subgraph.entities:
        print(f"  - {entity.label}: {entity.name}")

    print("\nDone!")


if __name__ == "__main__":
    asyncio.run(main())