from models.base import RTEAgent_base
import threading
import requests
import json
from tqdm import tqdm
import copy

class RTEAgent(RTEAgent_base):

    def __init__(self, args):

        super(RTEAgent, self).__init__(args.url, args.headers, args.model, args.temperature, args.tokens, args.threads, args.dataset)
        self.url = self.url
        self.headers = self.headers
        self.model = self.model
        self.temperature = self.temperature
        self.tokens = self.tokens
        self.threads = self.threads
        self.dataset = self.dataset
        self.memories = []


class RTEAgent_TE_Vanilla(RTEAgent):

    def prompt_construction(self, input, index):

        text = input[index]['text description']

        prompt_completion = {
                "model": self.model,
                "messages": [
                    {"role": "user",
                     "content": "Suppose you are an urban triplet extraction model and I will provide you some text data. Please extract triplets in the sentence." + '\n'},
                {"role": "user", "content": "Given the urban text: " + text + '\n'},
                    {"role": "user", "content": "Return the result with the following format without any extra explanation: {\"Triplet 1\": {\"Head entity\": \"\", \"Relation\": \"\", \"Tail entity\": \"\"}, \"Triplet 2\": {\"Head entity\": \"\", \"Relation\": \"\", \"Tail entity\": \"\"}}."
                     }
                ],
                "temperature": self.temperature,
                "tokens": self.tokens,
            }
        if "llama" in self.model:
            prompt_completion = self.prompt_gpt_to_llama(prompt_completion)
        if "mistral" in self.model:
            prompt_completion = self.prompt_gpt_to_mistral(prompt_completion)
        return prompt_completion