Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import openai | |
| import json | |
| from graphviz import Digraph | |
| from PIL import Image | |
| import io | |
| 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}", | |
| } | |
| ], | |
| 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"], | |
| }, | |
| } | |
| ], | |
| function_call={"name": "knowledge_graph"}, | |
| ) | |
| response_data = completion.choices[0]["message"]["function_call"]["arguments"] | |
| print(response_data) | |
| print("Type of response_data:", type(response_data)) | |
| print("Value of response_data:", response_data) | |
| # Convert to dictionary if it's a string | |
| if isinstance(response_data, str): | |
| response_data = json.loads(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']})", color=node.get("color", "lightblue")) | |
| for edge in response_data.get("edges", []): | |
| 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...") | |
| dot.format = "png" | |
| image_data = dot.pipe(format="png") | |
| image = Image.open(io.BytesIO(image_data)) | |
| print("Gráfico gerado com sucesso!") | |
| return image | |
| iface = gr.Interface( | |
| fn=generate_knowledge_graph, | |
| inputs=[ | |
| gr.inputs.Textbox(label="OpenAI API Key", type="password"), | |
| gr.inputs.Textbox(label="User Input for Graph"), | |
| ], | |
| outputs=gr.outputs.Image(type="pil", label="Generated Knowledge Graph"), | |
| live=False, | |
| ) | |
| print("Iniciando a interface Gradio...") | |
| iface.launch() |