import pandas as pd
from deepchecks_llm_client.client import DeepchecksLLMClient
from deepchecks_llm_client.data_types import EnvType, LogInteractionType, ApplicationType

DEEPCHECKS_LLM_API_KEY = "YOUR API KEY"
APP_NAME = "DemoApp"

# download data and read as csv
df = pd.read_csv('https://figshare.com/ndownloader/files/49102981')

# Init SDK's client and create an application
dc_client = DeepchecksLLMClient(api_token=DEEPCHECKS_LLM_API_KEY)
dc_client.create_application(APP_NAME, ApplicationType.QA)

# Log a batch of 97 samples used as the evaluation set to the Deepchecks platform
interactions = [
    LogInteractionType(user_interaction_id=idx, input=row["input"], 
                       information_retrieval=row["information_retrieval"], output=row["output"],
                       custom_props={'Question Type': row['cp_Question Type']},
                       annotation="Good",  # If annotations are unknown, remove this argument.
                       ) for idx, row in df.iterrows()]
dc_client.log_batch_interactions(app_name=APP_NAME, version_name="Ground Truth", 
                                 env_type=EnvType.EVAL, interactions=interactions)

from datetime import datetime

# Retrieve the evaluation set inputs
evaluation_set = dc_client.get_data(
    app_name=APP_NAME, version_name="Ground Truth", env_type=EnvType.EVAL)

interactions = []
# Run v1 version pipeline on the evaluation set
for _, row in evaluation_set.iterrows():
    start_time = datetime.now().astimezone()
    rephrased_input = rephrase(row['input'])
    retrieved_documents = v1_retrieval_func(rephrased_input)
    output = v1_llm_pipeline(rephrased_input, retrieved_documents)
    
    interactions.append(LogInteractionType(
        user_interaction_id=str(row['user_interaction_id']), input=row['input'],
        information_retrieval=retrieved_documents, output=output,
        steps=[Step(name="Rephrase", input=row['input'], output=rephrased_input)],
        started_at=start_time, finished_at=datetime.now().astimezone()))

# Log the interactions into the Deepchecks' platform
dc_client.log_batch_interactions(
    app_name=APP_NAME, version_name="v1_gpt3.5",
    env_type=EnvType.EVAL, interactions=interactions)

from uuid import uuid4

from deepchecks_llm_client.client import DeepchecksLLMClient
from deepchecks_llm_client.data_types import EnvType, Step
from openai import OpenAI

OPENAI_CLIENT = OpenAI(api_key="YOUR_API_TOKEN")
DC_CLIENT = DeepchecksLLMClient(api_token="YOUR_API_TOKEN")
APP_NAME = "MY_Q&A_APP"
VERSION_NAME = "V1"
SESSION_ID = None


def _call_open_ai_with_prompt(user_msg):
    response = OPENAI_CLIENT.chat.completions.create(model="gpt-4o-mini", messages=[
        {"role": "user", "content": [{"text": user_msg, "type": "text"}]}])
    return response.choices[0].message.content


def rephrase_question(question):
    """Rephrase the user question and generate prompt for LLM call."""
    prompt = f"Rephrase the below text into a well structured question:\n\n{question}"
    rephrased_question = _call_open_ai_with_prompt(prompt)
    DC_CLIENT.update_interaction(app_name=APP_NAME, version_name=VERSION_NAME, user_interaction_id=SESSION_ID, steps=[
        Step(name="Question Rephrase", input=question, output=rephrased_question)], )
    return rephrased_question


def get_ai_response(rephrased_question):
    """Sends a request to the OpenAI API and returns the response."""
    answer = _call_open_ai_with_prompt(rephrased_question)
    DC_CLIENT.update_interaction(app_name=APP_NAME, version_name=VERSION_NAME, user_interaction_id=SESSION_ID,
                                 output=answer, is_completed=True)
    return answer


if __name__ == "__main__":
    while True:
        user_input = input("Please enter your question: ")
        SESSION_ID = uuid4().hex
        DC_CLIENT.log_interaction(app_name=APP_NAME, version_name=VERSION_NAME, env_type=EnvType.PROD,
                                  user_interaction_id=SESSION_ID, input=user_input, is_completed=False)
        rephrased_question = rephrase_question(user_input)
        answer = get_ai_response(rephrased_question)
        print("\nAnswer:", answer)