import gradio as gr
import os
import openai

# Defina sua chave da API OpenAI
api_key = None

def create_knowledge_graph(user_input):
    global api_key
    if not api_key:
        return "Por favor, insira sua chave da API OpenAI."

    # Configurar a chamada para a API OpenAI
    openai.api_key = api_key
    response = 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",
                            "properties": {
                                "createdDate": {"type": "string"},
                                "lastUpdated": {"type": "string"},
                                "description": {"type": "string"},
                            },
                        },
                        "nodes": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "id": {"type": "string"},
                                    "label": {"type": "string"},
                                    "type": {"type": "string"},
                                    "color": {"type": "string"},
                                    "properties": {
                                        "type": "object",
                                        "description": "Additional attributes for the node",
                                    },
                                },
                                "required": [
                                    "id",
                                    "label",
                                    "type",
                                    "color",
                                ],
                            },
                        },
                        "edges": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "from": {"type": "string"},
                                    "to": {"type": "string"},
                                    "relationship": {"type": "string"},
                                    "direction": {"type": "string"},
                                    "color": {"type": "string"},
                                    "properties": {
                                        "type": "object",
                                        "description": "Additional attributes for the edge",
                                    },
                                },
                                "required": [
                                    "from",
                                    "to",
                                    "relationship",
                                    "color",
                                ],
                            },
                        },
                    },
                    "required": ["nodes", "edges"],
                },
            }
        ],
        function_call={"name": "knowledge_graph"},
    )

    response_data = response.choices[0]["message"]["function_call"]["arguments"]
    return response_data

# Defina a interface Gradio
iface = gr.Interface(
    fn=create_knowledge_graph,
    inputs=gr.Textbox("Texto para criar o gráfico de conhecimento:"),
    outputs=gr.Image(type="pil"),  # Imagem de saída para o gráfico
    live=True,
)

if __name__ == "__main__":
    iface.launch()