from GPT4KG import KnowledgeGraph
import gradio as gr
from PIL import Image

def generate_graph(input_text,api_key):
  try:
    kg = KnowledgeGraph(api_key,"temp.kg")
  except:
    kg = KnowledgeGraph(api_key)
  kg.learn(str(input_text))
  kg.save_graph("temp.kg")
  kg.display_graph("temp.png")
  return Image.open("temp.png")

def answer_question(question,api_key):
  try:
    kg = KnowledgeGraph(api_key,"temp.kg")
  except:
    kg = KnowledgeGraph(api_key)
  return kg.chat_qa(question)

title = "GPT-4 Knowledge Graph Generator"
description = "Enter text to generate a knowledge graph using GPT4KG:"

with gr.Blocks() as demo:
  with open("temp.kg","w") as f:
    f.write("")
  gr.Markdown("""<h1><center>GPT-4 Knowledge Graph Generator</center></h1>""")
  output_image = gr.Image(label="Knowledge Graph", type="pil")
  api_key = gr.Textbox(lines=1, label="OpenAI API Key")
  input_text = gr.Textbox(lines=5, label="Information to be added to graph")
  submit_btn = gr.Button("Add info to graph")
  submit_btn.click(fn=generate_graph, inputs=[input_text,api_key], outputs=[output_image])

  question = gr.Textbox(lines=1, label="Question about the info in this graph")
  answer = gr.Textbox(lines=1, label="Answer")
  qa_btn = gr.Button("Ask question")
  qa_btn.click(fn=answer_question, inputs=[question,api_key], outputs=[answer])

demo.launch(share=True)