import gradio as gr
from uuid import uuid4 as uuid
from huggingface_hub import HfApi
from typing import Any
from griptape.structures import Agent, Workflow
from griptape.tasks import PromptTask
from griptape.drivers import LocalConversationMemoryDriver
from griptape.memory.structure import ConversationMemory
from griptape.tools import Calculator
from griptape.rules import Rule
import time

repo_id = "kateforsberg/gradio-test"
api = HfApi()

def user(user_message, history):
    history.append([user_message, None])
    return ("", history)

def bot(history):
    response = agent.send_message(history[-1][0])
    history[-1][1] = ""
    for character in response:
        history[-1][1] += character

        time.sleep(0.005)

        yield history

agent = Agent(
        conversation_memory=ConversationMemory(
            driver=LocalConversationMemoryDriver(
                file_path="conversation_memory.json"
            )),
        tools=[Calculator()],
        rules=[
            Rule(
                value = "You are an intelligent agent tasked with answering questions."
            ),
            Rule(
                value = "All of your responses are less than 5 sentences."
            ) ]
)

workflow = Workflow()
start_task = PromptTask("I will provide you two dog breeds to compare", id="START")
end_task = PromptTask("How are the dog breeds different", id="END")
workflow.add_task(start_task)
workflow.add_task(end_task)
dog_task = PromptTask("Poodle and a Labrador", id="DOG")
workflow.insert_tasks(start_task,[dog_task],end_task)

def use_workflow(message:str, history) -> Any:
    response = workflow.run()
    return response.output.value

def send_message(message:str, history,) -> Any:
    response = agent.run(message)
    return response.output.value

demo = gr.ChatInterface(
    fn=send_message,
    )
demo.launch(
    auth=("griptape","griptaper"))