import json
import requests

class llama3chat:

    def __init__(self, url, bearer_token, max_new_tokens,top_p,temperature):
        self.url = url
        self.bearer_token = bearer_token
        self.headers = {
          'Content-Type': 'application/json',
          'Authorization': f'Bearer {bearer_token}'
        }
        self.max_new_tokens=max_new_tokens
        self.top_p=top_p
        self.temperature=temperature

    def format_llama_messages(self, messages):
        
      prompt = []

      if messages[0]["role"] == "system":
          content = "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"+messages[0]["content"]+ "<|eot_id|>\n\n"
          prompt.append(content)
      for user, answer in zip(messages[1::2], messages[2::2]):
          prompt.append("<|start_header_id|>user<|end_header_id|>"+ (user["content"]).strip()+ "<|eot_id|><|start_header_id|>assistant<|end_header_id|>"+ (answer["content"]).strip()+ "<|eot_id|>")

      prompt.append("<|start_header_id|>user<|end_header_id|>"+ (messages[-1]["content"]).strip()+ "<|eot_id|>")

      return "".join(prompt)

    def query(self, all_messages):
      formatted_messages=self.format_llama_messages(all_messages)
      json_body = {
          "inputs": formatted_messages,
                  "parameters": {"max_new_tokens":self.max_new_tokens, "top_p":self.top_p, "temperature":self.temperature}
          }
      data = json.dumps(json_body)
      response = requests.request("POST", self.url, headers=self.headers, data=data)
      try:
        result=json.loads(response.content.decode("utf-8"))[0]['generated_text']
        target_substring = '|eot_id|>assistant'
        substring_length = len(target_substring)
        last_occurrence_index = result.rfind(target_substring)
        if last_occurrence_index != -1:
          return result[last_occurrence_index + substring_length:].strip()
      except:
          return response