from datasets import load_dataset

n_samples = 100

ds = load_dataset("ServiceNow/repliqa")["repliqa_0"].select(list(range(n_samples)))

system_prompt = """You are a researcher answering questions from the given document. 
Please answer the question based on the contents of the document.

Make sure to read the document carefully before answering the question. 
Do not use any external knowledge.
For each question, present the reasoning followed by the correct answer.
If the answer to the question can not be obtained from the provided context paragraph, output 'UNANSWERABLE'.
"""

prompt_template = """
Document: {document}

Question: {question}
"""

from typing import List, Dict, Any, Optional
import requests
import json

def request_prompt_adaptation(
    system_prompt: str, 
    prompt_template: str, 
    fields: List[str], 
    train_data: List[Dict[str, str | Dict[str, str]]], 
    origin_model: Dict[str, str],
    target_models: List[Dict[str, str]],
    notdiamond_api_key: str,
  	evaluation_metric: Optional[str] = None,
		evaluation_config: Optional[Dict] = None,
    request_url: str = "https://api.notdiamond.ai/v2/prompt/adapt"
) -> str:
    """
    Helper method to call the prompt adaptation endpoint
    """
    
    if (evaluation_metric and evaluation_config) or (not evaluation_metric and not evaluation_config):
      raise ValueError("Either evaluation_metric or evaluation_config must be provided, but not both or neither.")
    
    request_body = {
        "system_prompt": system_prompt,
        "template": prompt_template,
        "fields": fields,
        "goldens": train_data,
        "origin_model": origin_model,
        "target_models": target_models,
    }
    
    if evaluation_metric:
        request_body["evaluation_metric"] = evaluation_metric
		
    if evaluation_config:	
        request_body["evaluation_config"] = json.dumps(evaluation_config)

    headers = {
        "Authorization": f"Bearer {notdiamond_api_key}",
        "content-type": "application/json"
    }
    
    resp = requests.post(
        request_url,
        headers=headers,
        json=request_body
    )
    
    if resp.status_code == 200:
        response = resp.json()
        return response["adaptation_run_id"]
    else:
        raise Exception(
            f"Request to adapt prompt failed with code {resp.status_code}: {resp.text}"
        )
        
def get_prompt_adaptation_results(
    adaptation_run_id: str,
    notdiamond_api_key: str,
    request_url: str = "https://api.notdiamond.ai/v2/prompt/adaptResults"
) -> Dict[str, Any]:
    """
    Helper method to get the results of a prompt adaptation request
    """
    
    headers = {
        "Authorization": f"Bearer {notdiamond_api_key}",
        "content-type": "application/json"
    }
    
    resp = requests.get(
        f"{request_url}/{adaptation_run_id}",
        headers=headers
    )
    
    if resp.status_code == 200:
        response = resp.json()
        return response
    else:
        raise Exception(
            f"Requesting prompt adaptation result failed with code {resp.status_code}: {resp.text}"
        )
  

Given the predicted answer and reference answer, compare them and check whether they mean the same.

Following are the given answers:

Predicted Answer: {predicted_answer}
Reference Answer: {gt_answer}

On a NEW LINE, give a score of 1 if the answers mean the same, and 0 if they differ, and nothing more.

You are a smart language model that evaluates the similarity between a predicted text and the expected ground truth answer. You do not propose changes to the answer and only critically evaluate the existing answer and provide feedback following the instructions given.

The following is the response provided by a language model to a prompt:
{predicted_answer}

The expected answer to this prompt is:
{gt_answer}

Answer only with an integer from 1 to 10 based on how semantically similar the responses are to the expected answer. where 0 is no semantic similarity at all and 10 is perfect agreement between the responses and the expected answer. On a NEW LINE, give the integer score and nothing more.

Below is a question that asks to generate a SQL query to answer a question.

Based on the question, determine if the generated SQL produces the same outcome as the correct SQL.

Look for the SQL query enclosed in <ANS_START> and <ANS_END> in the generated SQL.

Question:
{question}

Generated SQL:
{predicted_answer}

Correct SQL:
{gt_answer}

On a NEW LINE, give a score of 1 if the generated SQL produces the same outcome as the correct SQL, and 0 if they differ, and nothing more.

## Instructions

You are a helpful assistant that evaluates the correctness of a response to a question.
Please provide a score of 1 if the response is correct, and 0 if it is incorrect.

## Criteria

- Use information provided in the question to evaluate the response.
- If the question includes reference text which should be used to answer it, ensure that the answer
is based on the reference text. In this scenario, **the answer is incorrect if it is not based
on the reference text.**
- If the question does not include reference text, you can evaluate the answer based on your knowledge.
- Do not provide any other text in your response. Only output the numeric score of 0 or 1.

## Examples

Question: What is the capital of France?
Response: Paris
Score: 1

# In the example below, the answer is incorrect because it is not based on the reference text.
Question: Assume the capital of France is Berlin. What is the capital of France?
Response: Paris
Score: 0

Question: The meaning of life is 42. What is the meaning of life?
Response: 42
Score: 1

# In the example below, the answer is incorrect because it is not based on the reference text.
Question: The meaning of life is 42. What is the meaning of life?
Response: It can be difficult to reason about a subjective topic like this...
Score: 0

## Evaluation

Question: {question}
Response: {predicted_answer}
Score:

evaluation_config = {
    "llm_judging_prompt": (
        "Does the assistant's answer properly answer the user's question? question: {question} answer: {answer} "
        "Score a 1 if the answer is correct, 0 otherwise. Do not output any other values or text - only the score."
        ),
    "llm_judge": "openai/gpt-4o-2024-11-20",
    "correctness_cutoff": 0,
}

fields = ["document", "question"]

pa_ds = [
    {
        "fields": {
            "document": sample["document_extracted"],
            "question": sample["question"]
        },
        "answer": sample["answer"]
    }
    for sample in ds
]
print(prompt_template.format(**pa_ds[0]['fields']))

origin_model = {"provider": "openai", "model": "gpt-4o-2024-08-06"}
target_models = [
    {"provider": "anthropic", "model": "claude-3-5-sonnet-20241022"},
]

prompt_adaptation_id = request_prompt_adaptation(
    system_prompt=system_prompt,
    prompt_template=prompt_template,
    fields=fields,
    train_data=pa_ds,
    origin_model=origin_model,
    target_models=target_models,
    evaluation_metric="LLMaaJ:Sem_Sim_1",
    notdiamond_api_key="YOUR_NOTDIAMOND_API_KEY"
)

prompt_adaptation_id = request_prompt_adaptation(
    system_prompt=system_prompt,
    prompt_template=prompt_template,
    fields=fields,
    train_data=pa_ds,
    origin_model=origin_model,
  	target_models=target_models,
		evaluation_config=evaluation_config,
    notdiamond_api_key="YOUR_NOTDIAMOND_API_KEY"
)

results = get_prompt_adaptation_results(prompt_adaptation_id, "YOUR_NOTDIAMOND_API_KEY")
print(results)

{
  "id": "uuid", // The prompt adaptation id
  "created_at": "datetime", // Timestamp
  "origin_model": {
    "model_name": "openai/gpt-4o-2024-11-20", // The original model the prompt was designed for
    "score": 0.8, // The original model's score on the dataset before optimization
    "evals": {"LLMaaJ:Sem_Sim_1": 0.8}, // The original model's evaluation results on the dataset
    "system_prompt": "...", // The baseline system prompt submitted
    "user_message_template": "...", // The baseline prompt template submitted
    "result_status": "completed" 
  },
  "target_models": [
    {
      "model_name": "anthropic/claude-3-5-sonnet-20241022", // The original model the prompt was designed for
      "pre_optimization_score": 0.64, // The target model's score on the dataset before optimization
      "pre_optimization_evals": {"LLMaaJ:Sem_Sim_1": 0.64}, // The target model's evaluation results on the dataset before optimization
      "post_optimization_score": 0.8, // The target model's score on the dataset after optimization
      "post_optimization_evals": {"LLMaaJ:Sem_Sim_1": 0.8}, // The targe model's evaluation results on the dataset after optimization
      "system_prompt": "...", // The baseline system prompt submitted
      "user_message_template": "...", // The baseline prompt template submitted
      "user_message_template_fields": ["..."], // Field arguments in the user_message_template
      "result_status": "completed"
    }
  ],
}