import gradio as gr
import openai
import json
import requests
from bs4 import BeautifulSoup
from graphviz import Digraph
import base64
from io import BytesIO
from PIL import Image

def generate_knowledge_graph(api_key, user_input):
    openai.api_key = api_key

    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}",
            }
        ],
        function_call={"name": "knowledge_graph"},
    )

    response_data = completion.choices[0]["message"]["function_call"]["arguments"]
    response_dict = json.loads(response_data)

    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
    dot.format = "png"
    dot.render(filename="knowledge_graph", cleanup=True)

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

    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
)

iface.launch()