Update app.py

app.py

@@ -1,18 +1,17 @@
 import gradio as gr
-import os
 import openai
+import json
+import requests
+from bs4 import BeautifulSoup
+from graphviz import Digraph
+import base64
+from io import BytesIO
+from PIL import Image
 
-# 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
+def generate_knowledge_graph(api_key, user_input):
     openai.api_key = api_key
-    response = openai.ChatCompletion.create(
+
+    completion = openai.ChatCompletion.create(
         model="gpt-3.5-turbo-16k",
         messages=[
             {
@@ -20,84 +19,40 @@ def create_knowledge_graph(user_input):
                 "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
+    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}"
 
-# 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,
+    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
 )
 
-if __name__ == "__main__":
-    iface.launch()
+iface.launch()