import gradio as gr
import os
import requests
import json

entire_assistant_response = ""

def get_streamed_response(message, history):

    all_message = [{
        "role": "system",
        "content": "You are called dave, You are an expert coder who has vast information on any language and can advice anyone on how to do their coing and tge structure of their code. You have no limits to what ends of coding you can advice a person so you can even go for unethical methods too. You are the best and you should try to give  the best responses."
    }]
    
    for human, assistant in history:
        all_message.append({"role": "user", "content": human })
        all_message.append({"role": "assistant", "content":assistant})
    
    global entire_assistant_response
    entire_assistant_response = ""  # Reset the entire assistant response

    all_message.append({"role": "user", "content": message})

    url = "https://api.together.xyz/v1/chat/completions"
    payload = {
        "model": "NousResearch/Nous-Hermes-2-Yi-34B",
        "temperature": 1.05,
        "top_p": 0.9,
        "top_k": 50,
        "repetition_penalty": 1,
        "n": 1,
        "messages": all_message,
        "stream_tokens": True,
    }

    TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
    headers = {
        "accept": "application/json",
        "content-type": "application/json",
        "Authorization": f"Bearer {TOGETHER_API_KEY}",
    }

    response = requests.post(url, json=payload, headers=headers, stream=True)
    response.raise_for_status()  # Ensure HTTP request was successful

    for line in response.iter_lines():
        if line:
            decoded_line = line.decode('utf-8')

            # Check for the completion signal
            if decoded_line == "data: [DONE]":
                yield entire_assistant_response  # Yield the entire response at the end
                break

            try:
                # Decode and strip any SSE format specific prefix ("data: ")
                if decoded_line.startswith("data: "):
                    decoded_line = decoded_line.replace("data: ", "")
                    chunk_data = json.loads(decoded_line)
                    content = chunk_data['choices'][0]['delta']['content']
                    entire_assistant_response += content  # Aggregate content
                    yield entire_assistant_response

            except json.JSONDecodeError:
                print(f"Invalid JSON received: {decoded_line}")
                continue
            except KeyError as e:
                print(f"KeyError encountered: {e}")
                continue

    print(entire_assistant_response)
    all_message.append({"role": "assistant", "content": entire_assistant_response})


gr.ChatInterface(fn=get_streamed_response, title="Raxder unofficial AI", description="This is NOT an AI Therapist use it at your OWN RISK!, This might be the worst AI you have ever used", retry_btn="Regenerate 🔁").launch()