|
import gradio as gr |
|
import openai |
|
import json |
|
from graphviz import Digraph |
|
import base64 |
|
from io import BytesIO |
|
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.Completion.create( |
|
engine="text-davinci-002", |
|
prompt=f"Help me understand the following by describing it as a detailed knowledge graph: {user_input}", |
|
max_tokens=100 |
|
) |
|
|
|
print("Received response from OpenAI.") |
|
response_data = completion.choices[0].text |
|
print(f"Response data: {response_data}") |
|
|
|
|
|
|
|
print("Converting response to JSON...") |
|
try: |
|
response_dict = json.loads(response_data) |
|
except json.JSONDecodeError: |
|
print("Failed to decode JSON. Using empty dictionary as a fallback.") |
|
response_dict = {} |
|
|
|
print("Generating knowledge graph using Graphviz...") |
|
dot = Digraph(comment="Knowledge Graph") |
|
|
|
|
|
for node in response_dict.get("nodes", []): |
|
dot.node(node["id"], f"{node['label']} ({node['type']})") |
|
|
|
|
|
for edge in response_dict.get("edges", []): |
|
dot.edge(edge["from"], edge["to"], label=edge["relationship"]) |
|
|
|
|
|
print("Rendering graph to PNG format...") |
|
dot.format = "png" |
|
dot.render(filename="knowledge_graph", cleanup=True) |
|
|
|
|
|
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() |
|
|