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

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

def answer_question(question,api_key,graph):
  if graph == []:
    kg = KnowledgeGraph(api_key)
    graph.append(kg)
  else:
    kg = graph[0]
  return kg.chat_qa(question)

def clear_graph(api_key,graph):
  graph = []
  kg = KnowledgeGraph(api_key)
  graph.append(kg)
  return graph,None#Image.new('RGB', (400, 100),(255, 255, 255))

title = "Associative Memory with GPT4KG"
description = "Enter text to generate a semantically searchable knowledge graph:"

with gr.Blocks() as demo:
  gr.Markdown(f"<h1><center>{title}</center></h1>")
  gr.Markdown(f"<h3><center>{description}</center></h3>")
  
  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])

  clear_btn = gr.Button("Clear graph")
  clear_btn.click(fn=clear_graph, inputs=[api_key,graph], outputs=[graph,output_image],api_name="clear")

demo.launch()