import gradio as gr
import openai
import json
from graphviz import Digraph

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

    # Chamar a API da OpenAI
    print("Chamando a API da 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}",
            }
        ]
    )
    response_data = completion.choices[0].message.to_dict()
    response_data = json.loads(response_data['content'])

    print("Dados da resposta:")
    print(response_data)

    # Visualizar o conhecimento usando Graphviz
    print("Gerando o conhecimento usando Graphviz...")
    dot = Digraph(comment="Knowledge Graph")
    for node in response_data.get("nodes", []):
        dot.node(node["id"], f"{node['label']} ({node['type']})")
    for edge in response_data.get("edges", []):
        dot.edge(edge["from"], edge["to"], label=edge["relationship"])

    # Renderizar para o formato PNG
    print("Renderizando o gráfico para o formato PNG...")
    dot.format = "png"
    dot.render(filename="knowledge_graph", cleanup=True)

    print("Gráfico gerado com sucesso!")

    return "knowledge_graph.png"

iface = gr.Interface(
    fn=generate_knowledge_graph, 
    inputs=[
        gr.components.Textbox(label="OpenAI API Key", type="password"),
        gr.components.Textbox(label="User Input for Graph")
    ], 
    outputs=gr.components.Image(type="filepath", label="Generated Knowledge Graph"),
    live=False
)

print("Iniciando a interface Gradio...")
iface.launch()