File size: 2,667 Bytes
8ea927a
 
410031a
 
 
 
8ea927a
410031a
b79f8d4
8ea927a
410031a
b79f8d4
b31a1e4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8ea927a
 
b79f8d4
b31a1e4
b79f8d4
410031a
b79f8d4
b31a1e4
b79f8d4
 
410031a
 
 
 
 
 
 
 
 
 
 
b79f8d4
410031a
 
 
 
b79f8d4
410031a
 
 
b79f8d4
410031a
8ea927a
 
410031a
 
 
 
 
 
 
8ea927a
 
b79f8d4
410031a
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import gradio as gr
import openai
import json
from graphviz import Digraph
import base64
from PIL import Image

def generate_knowledge_graph(api_key, user_input):
    print("Setting OpenAI API key...")
    openai.api_key = api_key

    print("Making API call to OpenAI...")
    completion = openai.ChatCompletion.create(
        model="gpt-3.5-turbo-16k",
        messages=[
            {
                "role": "user",
                "content": f"Help me understand following by describing as a detailed knowledge graph: {user_input}",
            }
        ],
        functions=[
            {
                "name": "knowledge_graph",
                "description": "Generate a knowledge graph with entities and relationships.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "metadata": {"type": "object"},
                        "nodes": {"type": "array"},
                        "edges": {"type": "array"}
                    },
                    "required": ["nodes", "edges"]
                }
            }
        ],
        function_call={"name": "knowledge_graph"}
    )

    print("Received response from OpenAI.")
    response_data = completion.choices[0]["message"]["function_call"]["arguments"]
    print(f"Response data: {response_data}")

    print("Converting response to JSON...")
    response_dict = json.loads(response_data)

    print("Generating knowledge graph using Graphviz...")
    dot = Digraph(comment="Knowledge Graph")

    # Add nodes to the graph
    for node in response_dict.get("nodes", []):
        dot.node(node["id"], f"{node['label']} ({node['type']})")

    # Add edges to the graph
    for edge in response_dict.get("edges", []):
        dot.edge(edge["from"], edge["to"], label=edge["relationship"])

    # Render to PNG format
    print("Rendering graph to PNG format...")
    dot.format = "png"
    dot.render(filename="knowledge_graph", cleanup=True)

    # Convert PNG to base64 to display in Gradio
    print("Converting PNG to base64...")
    with open("knowledge_graph.png", "rb") as img_file:
        img_base64 = base64.b64encode(img_file.read()).decode()

    print("Returning base64 image to Gradio interface.")
    return f"data:image/png;base64,{img_base64}"

iface = gr.Interface(
    fn=generate_knowledge_graph,
    inputs=[
        gr.inputs.Textbox(label="OpenAI API Key", type="password"),
        gr.inputs.Textbox(label="Text to Generate Knowledge Graph")
    ],
    outputs=gr.outputs.Image(type="pil", label="Generated Knowledge Graph"),
    live=False
)

print("Launching Gradio interface...")
iface.launch()