instructor

1
from openai import OpenAI
2
import instructor
3

4
from graphviz import Digraph
5
from typing import Optional
6

7
from pydantic import BaseModel, Field
8

9
client = instructor.from_openai(OpenAI())
10

11

12
class Node(BaseModel):
13
    id: int
14
    label: str
15
    color: str
16

17
    def __hash__(self) -> int:
18
        return hash((id, self.label))
19

20

21
class Edge(BaseModel):
22
    source: int
23
    target: int
24
    label: str
25
    color: str = "black"
26

27
    def __hash__(self) -> int:
28
        return hash((self.source, self.target, self.label))
29

30

31
class KnowledgeGraph(BaseModel):
32
    nodes: Optional[list[Node]] = Field(..., default_factory=list)
33
    edges: Optional[list[Edge]] = Field(..., default_factory=list)
34

35
    def update(self, other: "KnowledgeGraph") -> "KnowledgeGraph":
36
        """Updates the current graph with the other graph, deduplicating nodes and edges."""
37
        return KnowledgeGraph(
38
            nodes=list(set(self.nodes + other.nodes)),
39
            edges=list(set(self.edges + other.edges)),
40
        )
41

42
    def draw(self, prefix: str = None):
43
        dot = Digraph(comment="Knowledge Graph")
44

45
        # Add nodes
46
        for node in self.nodes:
47
            dot.node(str(node.id), node.label, color=node.color)
48

49
        # Add edges
50
        for edge in self.edges:
51
            dot.edge(
52
                str(edge.source), str(edge.target), label=edge.label, color=edge.color
53
            )
54
        dot.render(prefix, format="png", view=True)
55

56

57
def generate_graph(input: list[str]) -> KnowledgeGraph:
58
    cur_state = KnowledgeGraph()
59
    num_iterations = len(input)
60
    for i, inp in enumerate(input):
61
        new_updates = client.chat.completions.create(
62
            model="gpt-3.5-turbo-16k",
63
            messages=[
64
                {
65
                    "role": "system",
66
                    "content": """You are an iterative knowledge graph builder.
67
                    You are given the current state of the graph, and you must append the nodes and edges 
68
                    to it Do not procide any duplcates and try to reuse nodes as much as possible.""",
69
                },
70
                {
71
                    "role": "user",
72
                    "content": f"""Extract any new nodes and edges from the following:
73
                    # Part {i}/{num_iterations} of the input:
74

75
                    {inp}""",
76
                },
77
                {
78
                    "role": "user",
79
                    "content": f"""Here is the current state of the graph:
80
                    {cur_state.model_dump_json(indent=2)}""",
81
                },
82
            ],
83
            response_model=KnowledgeGraph,
84
        )  # type: ignore
85

86
        # Update the current state
87
        cur_state = cur_state.update(new_updates)
88
        cur_state.draw(prefix=f"iteration_{i}")
89
    return cur_state
90

91

92
# here we assume that we have to process the text in chunks
93
# one at a time since they may not fit in the prompt otherwise
94
text_chunks = [
95
    "Jason knows a lot about quantum mechanics. He is a physicist. He is a professor",
96
    "Professors are smart.",
97
    "Sarah knows Jason and is a student of his.",
98
    "Sarah is a student at the University of Toronto. and UofT is in Canada.",
99
]
100

101
graph: KnowledgeGraph = generate_graph(text_chunks)
102

103
graph.draw(prefix="final")
104