Update app.py

app.py

@@ -8,37 +8,92 @@ def generate_knowledge_graph(api_key, user_input):
 
     # Chamar a API da OpenAI
     print("Chamando a API da OpenAI...")
-    completion = openai.ChatCompletion.create(
+    completion = openai.Completion.create(
         model="gpt-3.5-turbo-16k",
         messages=[
+            {"role": "system", "content": "You are a helpful assistant."},
             {
                 "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. Use the colors to help differentiate between different node or edge types/categories. Always provide light pastel colors that work well with black font.",
+                "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"},  # Added color property
+                                    "properties": {
+                                        "type": "object",
+                                        "description": "Additional attributes for the node",
+                                    },
+                                },
+                                "required": [
+                                    "id",
+                                    "label",
+                                    "type",
+                                    "color",
+                                ],  # Added color to required
+                            },
+                        },
+                        "edges": {
+                            "type": "array",
+                            "items": {
+                                "type": "object",
+                                "properties": {
+                                    "from": {"type": "string"},
+                                    "to": {"type": "string"},
+                                    "relationship": {"type": "string"},
+                                    "direction": {"type": "string"},
+                                    "color": {"type": "string"},  # Added color property
+                                    "properties": {
+                                        "type": "object",
+                                        "description": "Additional attributes for the edge",
+                                    },
+                                },
+                                "required": [
+                                    "from",
+                                    "to",
+                                    "relationship",
+                                    "color",
+                                ],  # Added color to required
+                            },
+                        },
+                    },
+                    "required": ["nodes", "edges"],
+                },
             }
-        ]
+        ],
+        response_format="json",
     )
-    raw_response = completion.choices[0].message.to_dict()
-    print("Resposta bruta da API:")
-    print(raw_response)
-
-    # Verificar se a resposta contém conteúdo
-    if 'content' in raw_response and raw_response['content']:
-        try:
-            response_data = json.loads(raw_response['content'])
-        except json.JSONDecodeError:
-            print("Erro ao decodificar o JSON.")
-            return "Erro ao decodificar o JSON."
-    else:
-        print("Resposta da API vazia ou inválida.")
-        return "Resposta da API vazia ou inválida."
+    response_data = completion.choices[0]["message"]["function_results"]["knowledge_graph"]
+    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']})")
+        dot.node(node["id"], f"{node['label']} ({node['type']})", color=node.get("color", "lightblue"))
     for edge in response_data.get("edges", []):
-        dot.edge(edge["from"], edge["to"], label=edge["relationship"])
+        dot.edge(edge["from"], edge["to"], label=edge["relationship"], color=edge.get("color", "black"))
 
     # Renderizar para o formato PNG
     print("Renderizando o gráfico para o formato PNG...")
@@ -50,13 +105,13 @@ def generate_knowledge_graph(api_key, user_input):
     return "knowledge_graph.png"
 
 iface = gr.Interface(
-    fn=generate_knowledge_graph, 
+    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
+        gr.inputs.Textbox(label="OpenAI API Key", type="password"),
+        gr.inputs.Textbox(label="User Input for Graph"),
+    ],
+    outputs=gr.outputs.Image(type="file", label="Generated Knowledge Graph"),
+    live=False,
 )
 
 print("Iniciando a interface Gradio...")