{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating Knowledge Graphs with LLMs and Structured Outputs\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/meta-llama/llama-cookbook/blob/main/3p-integrations/togetherai/knowledge_graphs_with_structured_outputs.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "In this notebook we show how to generate and visualize knowledge graphs using Large Language Models (LLMs), and how to use JSON mode for structured output generation. We will use [Together AI JSON mode](https://docs.together.ai/docs/json-mode) to access models that will generate JSON structures that represent a knowledge graph (including nodes, edges, labels) for us. \n",
    "\n",
    "We will then pass these knowledge graph components to [GraphViz](https://graphviz.org/) that will help us generate and visualize these graphs.\n",
    "\n",
    "<img src=\"images/cluster.png\" height=\"500\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install relevant libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DFAjay1FZVrn",
    "outputId": "d4b17b31-c125-4de5-ad54-6d4d08d81eaa"
   },
   "outputs": [],
   "source": [
    "!pip install together\n",
    "!pip install graphviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "5L-bdWL8Zz2B"
   },
   "outputs": [],
   "source": [
    "from together import Together\n",
    "import together\n",
    "import json\n",
    "\n",
    "# Paste in your Together AI API Key\n",
    "TOGETHER_API_KEY = \"YOUR_TOGETHER_API_KEY\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Knowledge Graph Schema in Pydantic\n",
    "\n",
    "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.\n",
    "\n",
    "Below we define the required classes. \n",
    "\n",
    "- Each node in our graph needs to have an `id` and a `label`. The `Node` class specifies this.\n",
    "- 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.\n",
    "- Each Knowledge graph is a combination of multiple `Node` and `Edge` elements. The `KnowledgeGraph` class specifies this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "ycvtkJdhZ5Zo"
   },
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from typing import List\n",
    "\n",
    "class Node(BaseModel, frozen=True):\n",
    "    id: int\n",
    "    label: str\n",
    "\n",
    "class Edge(BaseModel, frozen=True):\n",
    "    source: int\n",
    "    target: int\n",
    "    label: str\n",
    "\n",
    "class KnowledgeGraph(BaseModel):\n",
    "    nodes: List[Node] = Field(..., default_factory=list)\n",
    "    edges: List[Edge] = Field(..., default_factory=list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5auwyteuadki",
    "outputId": "38fecb4a-5c11-4755-e311-a376bc308bee"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'$defs': {'Edge': {'properties': {'source': {'title': 'Source',\n",
       "     'type': 'integer'},\n",
       "    'target': {'title': 'Target', 'type': 'integer'},\n",
       "    'label': {'title': 'Label', 'type': 'string'}},\n",
       "   'required': ['source', 'target', 'label'],\n",
       "   'title': 'Edge',\n",
       "   'type': 'object'},\n",
       "  'Node': {'properties': {'id': {'title': 'Id', 'type': 'integer'},\n",
       "    'label': {'title': 'Label', 'type': 'string'}},\n",
       "   'required': ['id', 'label'],\n",
       "   'title': 'Node',\n",
       "   'type': 'object'}},\n",
       " 'properties': {'nodes': {'items': {'$ref': '#/$defs/Node'},\n",
       "   'title': 'Nodes',\n",
       "   'type': 'array'},\n",
       "  'edges': {'items': {'$ref': '#/$defs/Edge'},\n",
       "   'title': 'Edges',\n",
       "   'type': 'array'}},\n",
       " 'title': 'KnowledgeGraph',\n",
       " 'type': 'object'}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Print out the JSON Schema for our knowledge graph\n",
    "KnowledgeGraph.model_json_schema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function to Generate a Knowledge Graph using LLMs\n",
    "\n",
    "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. \n",
    "\n",
    "We pass in the above `KnowledgeGraph` JSON schema to guide the structured output generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "GHwlq8vuZ9QM"
   },
   "outputs": [],
   "source": [
    "# Call the LLM with the JSON schema\n",
    "def generate_graph(input) -> KnowledgeGraph:\n",
    "    together = Together(api_key = TOGETHER_API_KEY)\n",
    "\n",
    "    extract = together.chat.completions.create(\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": f\"Build a knowledge graph to explain: {input}. Only answer in JSON.\",\n",
    "            }\n",
    "        ],\n",
    "        model=\"meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo\",\n",
    "        response_format={\n",
    "            \"type\": \"json_object\",\n",
    "            \"schema\": KnowledgeGraph.model_json_schema(),\n",
    "        },\n",
    "    )\n",
    "\n",
    "    output = json.loads(extract.choices[0].message.content)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's use the function above to generate a knowledge graph about the paper Attention Is All You Need,\n",
    "# a landmark in AI which introduced the Transformer architecture, enabling breakthroughs in LLMs.\n",
    "\n",
    "graph = generate_graph(\"The paper Attention Is All You Need\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'nodes': [{'id': 1, 'label': 'Transformer'},\n",
       "  {'id': 2, 'label': 'Attention Mechanism'},\n",
       "  {'id': 3, 'label': 'Encoder-Decoder Architecture'},\n",
       "  {'id': 4, 'label': 'Self-Attention'},\n",
       "  {'id': 5, 'label': 'Vaswani et al.'},\n",
       "  {'id': 6, 'label': 'Sequence-to-Sequence Models'},\n",
       "  {'id': 7, 'label': 'Recurrent Neural Networks (RNNs)'},\n",
       "  {'id': 8, 'label': 'Convolutional Neural Networks (CNNs)'},\n",
       "  {'id': 9, 'label': 'Parallelization'}],\n",
       " 'edges': [{'source': 1, 'target': 2, 'label': 'utilizes'},\n",
       "  {'source': 1, 'target': 3, 'label': 'employs'},\n",
       "  {'source': 2, 'target': 4, 'label': 'includes'},\n",
       "  {'source': 3, 'target': 6, 'label': 'enhances'},\n",
       "  {'source': 5, 'target': 1, 'label': 'proposes'},\n",
       "  {'source': 6, 'target': 7, 'label': 'replaces'},\n",
       "  {'source': 6, 'target': 8, 'label': 'outperforms'},\n",
       "  {'source': 1, 'target': 9, 'label': 'enables'}]}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Lets see the knowledge graph components generated!\n",
    "graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use GraphViz to Create and Visualize the Graph\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "dso_caYbapzO"
   },
   "outputs": [],
   "source": [
    "from graphviz import Digraph\n",
    "\n",
    "def visualize_knowledge_graph(kg):\n",
    "    dot = Digraph(comment=\"Knowledge Graph\", format='png')\n",
    "\n",
    "    # Add nodes\n",
    "    for node in kg['nodes']:\n",
    "        dot.node(str(node['id']), node['label'])\n",
    "\n",
    "    # Add edges\n",
    "    for edge in kg['edges']:\n",
    "        dot.edge(str(edge['source']), str(edge['target']), label=edge['label'])\n",
    "\n",
    "    # Render the graph to a file and open it\n",
    "    output_path = dot.render(\"knowledge_graph\", view=True)\n",
    "    print(f\"Graph rendered and saved to {output_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q9v0fnk9dgLY",
    "outputId": "668c6e7a-4460-4248-e7af-ed6c5eb54d2b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Graph rendered and saved to knowledge_graph.png\n"
     ]
    }
   ],
   "source": [
    "visualize_knowledge_graph(graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 540
    },
    "id": "w8JnsjGaifAx",
    "outputId": "d08dabeb-5777-4aa0-a5a3-9f32917b1dc1"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 95216 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "# Display the image of the knowledge graph\n",
    "display(Image(filename='knowledge_graph.png'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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!\n",
    "\n",
    "Learn more about how to use JSON mode in the [docs](https://docs.together.ai/docs/json-mode) here!"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
