from gradio_client import Client
from trulens_eval.schema import Select
from trulens_eval.tru import Tru
from trulens_eval.feedback import Feedback
from trulens_eval.feedback import OpenAI as Feedback_OpenAI
from langchain.llms import HuggingFacePipeline
from langchain.prompts import PromptTemplate
from langchain.chains import ConversationChain
from langchain.chains.conversation.memory import ConversationBufferWindowMemory
import os

# Access environment variables
openai_api_key = os.environ.get("OPENAI_API_KEY")
huggingface_api_token = os.environ.get("HUGGINGFACE_API_TOKEN")


# Define a feedback function for query-statement relevance using OpenAI
feedback_openai = Feedback_OpenAI()
qa_relevance = Feedback(feedback_openai.relevance, name="Answer Relevance").on_input_output()

# Create a Tru object
tru = Tru()

# Initialize the HuggingFacePipeline for local LLM
local_llm = HuggingFacePipeline.from_model_id(
    model_id="Tonic/stablemed",
    task="text-generation",
    model_kwargs={"temperature": 0.2, "top_p": 0.95, "max_length": 256}
)

# Set the window memory to go back 4 turns
window_memory = ConversationBufferWindowMemory(k=4)

# Create the ConversationChain with the given window memory
conversation = ConversationChain(
    llm=local_llm,
    verbose=True,
    memory=window_memory
)

# Update the conversation prompt template to prime it as a gardening expert
conversation.prompt.template = '''The following is a consult between a clinical consultant and a public health and medical expert. The AI is an expert on medicine and public health and gives recommendations specific to location and conditions. If the AI does not know the answer to a question, it truthfully says it does not know.

Current conversation:
{history}
Human: {input}
AI:'''

# Wrap the conversation with TruChain to instrument it
tc_conversation = tru.Chain(conversation, app_id='Trulens-StableMed', feedbacks=[qa_relevance])

# Initialize Gradio Client
client = Client("https://tonic-stablemed-chat.hf.space/")

# Make a prediction using the wrapped conversation
result = client.predict(
    "Howdy!",  # str in 'user_input' Textbox component
    "Howdy!",  # str in 'system_prompt' Textbox component
    api_name="/predict"
)

# Print the result
print(result)

tru.run_dashboard()