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

def generate_graph(input_text,api_key,graph):
  if graph[0] =="":
    kg = KnowledgeGraph(api_key)
    graph[0] = kg
  else:
    kg = graph[0]
  kg.learn(str(input_text))
  img = kg.display_graph()
  return img,graph

def answer_question(question,api_key,graph):
  if graph[0] =="":
    kg = KnowledgeGraph(api_key)
    graph[0] = kg
  else:
    kg = graph[0]
  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:
  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")
  graph = gr.State([""])
  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,graph], outputs=[output_image,graph])

  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,graph], outputs=[answer])

demo.launch()