Notebooks
T
Together
Knowledge Graphs With Structured Outputs

Generating Knowledge Graphs with LLMs and Structured Outputs

Open In Colab

Introduction

In this notebook we show how to generate and visualize knowledge graphs using Large Language Models (LLMs), JSON mode for structured output generation. We will use Together AI JSON mode to access models that will generate JSON structures that represent a knowledge graph (including nodes, edges, labels) for us.

We will then pass these knowledge graph components to GraphViz that will help us generate and visualize these graphs.

Install relevant libraries

[1]
Collecting together
  Downloading together-1.3.3-py3-none-any.whl.metadata (11 kB)
Requirement already satisfied: aiohttp<4.0.0,>=3.9.3 in /usr/local/lib/python3.10/dist-packages (from together) (3.10.10)
Requirement already satisfied: click<9.0.0,>=8.1.7 in /usr/local/lib/python3.10/dist-packages (from together) (8.1.7)
Requirement already satisfied: eval-type-backport<0.3.0,>=0.1.3 in /usr/local/lib/python3.10/dist-packages (from together) (0.2.0)
Requirement already satisfied: filelock<4.0.0,>=3.13.1 in /usr/local/lib/python3.10/dist-packages (from together) (3.16.1)
Requirement already satisfied: numpy>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from together) (1.26.4)
Requirement already satisfied: pillow<11.0.0,>=10.3.0 in /usr/local/lib/python3.10/dist-packages (from together) (10.4.0)
Requirement already satisfied: pyarrow>=10.0.1 in /usr/local/lib/python3.10/dist-packages (from together) (16.1.0)
Requirement already satisfied: pydantic<3.0.0,>=2.6.3 in /usr/local/lib/python3.10/dist-packages (from together) (2.9.2)
Requirement already satisfied: requests<3.0.0,>=2.31.0 in /usr/local/lib/python3.10/dist-packages (from together) (2.32.3)
Requirement already satisfied: rich<14.0.0,>=13.8.1 in /usr/local/lib/python3.10/dist-packages (from together) (13.9.2)
Requirement already satisfied: tabulate<0.10.0,>=0.9.0 in /usr/local/lib/python3.10/dist-packages (from together) (0.9.0)
Requirement already satisfied: tqdm<5.0.0,>=4.66.2 in /usr/local/lib/python3.10/dist-packages (from together) (4.66.5)
Requirement already satisfied: typer<0.13,>=0.9 in /usr/local/lib/python3.10/dist-packages (from together) (0.12.5)
Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (2.4.3)
Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (1.3.1)
Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (24.2.0)
Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (1.4.1)
Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (6.1.0)
Requirement already satisfied: yarl<2.0,>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (1.15.2)
Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp<4.0.0,>=3.9.3->together) (4.0.3)
Requirement already satisfied: annotated-types>=0.6.0 in /usr/local/lib/python3.10/dist-packages (from pydantic<3.0.0,>=2.6.3->together) (0.7.0)
Requirement already satisfied: pydantic-core==2.23.4 in /usr/local/lib/python3.10/dist-packages (from pydantic<3.0.0,>=2.6.3->together) (2.23.4)
Requirement already satisfied: typing-extensions>=4.6.1 in /usr/local/lib/python3.10/dist-packages (from pydantic<3.0.0,>=2.6.3->together) (4.12.2)
Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (3.4.0)
Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (3.10)
Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (2.2.3)
Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.31.0->together) (2024.8.30)
Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich<14.0.0,>=13.8.1->together) (3.0.0)
Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich<14.0.0,>=13.8.1->together) (2.18.0)
Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from typer<0.13,>=0.9->together) (1.5.4)
Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich<14.0.0,>=13.8.1->together) (0.1.2)
Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from yarl<2.0,>=1.12.0->aiohttp<4.0.0,>=3.9.3->together) (0.2.0)
Downloading together-1.3.3-py3-none-any.whl (68 kB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 68.1/68.1 kB 2.4 MB/s eta 0:00:00
Installing collected packages: together
Successfully installed together-1.3.3
Requirement already satisfied: graphviz in /usr/local/lib/python3.10/dist-packages (0.20.3)
[11]

Define Knowledge Graph Schema in Pydantic

We need a way of telling the LLM what structure a knowledge graph has - including what information the nodes, edges must have and how they come together to create the overall graph structure. We will do this using pydantic models.

Below we define the required classes.

  • Each node in our graph needs to have an id and a label. The Node class specifies this.
  • Each edge in our graph needs to have a label and must also connect two nodes from the source to destination direction. The Edge class specified this.
  • Each Knowledge graph is a combination of multiple Node and Edge elements. The KnowledgeGraph class specifies this.
[38]
[39]
{'$defs': {'Edge': {'properties': {'source': {'title': 'Source',
,     'type': 'integer'},
,    'target': {'title': 'Target', 'type': 'integer'},
,    'label': {'title': 'Label', 'type': 'string'}},
,   'required': ['source', 'target', 'label'],
,   'title': 'Edge',
,   'type': 'object'},
,  'Node': {'properties': {'id': {'title': 'Id', 'type': 'integer'},
,    'label': {'title': 'Label', 'type': 'string'}},
,   'required': ['id', 'label'],
,   'title': 'Node',
,   'type': 'object'}},
, 'properties': {'nodes': {'items': {'$ref': '#/$defs/Node'},
,   'title': 'Nodes',
,   'type': 'array'},
,  'edges': {'items': {'$ref': '#/$defs/Edge'},
,   'title': 'Edges',
,   'type': 'array'}},
, 'title': 'KnowledgeGraph',
, 'type': 'object'}

Function to Generate a Knowledge Graph using LLMs

Here we use Together AI to access Meta-Llama-3.1-70B in strict JSON mode. We pass in instructions to get the LLM to generate a knowledge that captures the information about a passed in input string.

We pass in the above KnowledgeGraph JSON schema to guide the structured output generation.

[64]
[ ]
[ ]
{'nodes': [{'id': 1, 'label': 'NBA 2019 playoffs'},
,  {'id': 2, 'label': 'Eastern Conference'},
,  {'id': 3, 'label': 'Western Conference'},
,  {'id': 4, 'label': 'Milwaukee Bucks'},
,  {'id': 5, 'label': 'Toronto Raptors'},
,  {'id': 6, 'label': 'Golden State Warriors'},
,  {'id': 7, 'label': 'Portland Trail Blazers'},
,  {'id': 8, 'label': 'NBA Finals'},
,  {'id': 9, 'label': 'Conference Finals'},
,  {'id': 10, 'label': 'Semifinals'},
,  {'id': 11, 'label': 'First Round'}],
, 'edges': [{'source': 1, 'target': 2, 'label': 'conference'},
,  {'source': 1, 'target': 3, 'label': 'conference'},
,  {'source': 2, 'target': 4, 'label': 'participant'},
,  {'source': 2, 'target': 5, 'label': 'participant'},
,  {'source': 3, 'target': 6, 'label': 'participant'},
,  {'source': 3, 'target': 7, 'label': 'participant'},
,  {'source': 4, 'target': 9, 'label': 'qualified for'},
,  {'source': 5, 'target': 9, 'label': 'qualified for'},
,  {'source': 6, 'target': 8, 'label': 'qualified for'},
,  {'source': 7, 'target': 9, 'label': 'qualified for'},
,  {'source': 5, 'target': 8, 'label': 'qualified for'},
,  {'source': 6, 'target': 9, 'label': 'qualified for'},
,  {'source': 9, 'target': 8, 'label': 'preceded by'},
,  {'source': 10, 'target': 9, 'label': 'preceded by'},
,  {'source': 11, 'target': 10, 'label': 'preceded by'}]}

Use GraphViz to Create and Visualize the Graph

Below we write a function that takes the knowledge graph JSON structure above and creates a GraphViz Directed Graph(Digraph) and generates a picture of the graph.

[65]
[68]
Graph rendered and saved to knowledge_graph.png
[69]
Output

We were able to get a LLM to generate a knowledge graph by telling is the structure that a knowledge graph library, GraphViz, expects the output to be in. We then used GraphViz to create and visualize the knowledge graph above!

Learn more about how to use JSON mode in the docs here!