graph-g / app.py
artificialguybr's picture
Update app.py
b79f8d4
raw
history blame
No virus
2.34 kB
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}")
# For demonstration, let's assume the response_data is a JSON string that can be converted to a dictionary.
# You'll need to write code to interpret the text-based response to generate this dictionary.
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")
# 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()