import requests 
import os  
from openai import OpenAI 

# Set up API credentials and endpoints
openai_api_key=os.getenv("OPENAI_API_KEY")  # Get OpenAI API key from environment variables
api_url="https://api.tokenfactory.nebius.com"  # Base URL for the Nebius Token Factory API
base_url = api_url+"/v1"  # v1 API endpoint

# Initialize the OpenAI client with custom base URL and API key
client = OpenAI(
    base_url=base_url,
    api_key=os.getenv('NEBIUS_API_KEY', 'ANY'),  # Use Nebius Token Factory API key or fallback to 'ANY'
)

def upload_file(file_name):
    # Upload the file to the API server
    with open(file_name, "rb") as file_data:
        files = {"file": (os.path.basename(file_name), file_data)}
        upload_response = requests.post(
            f"{api_url}/v0/models/upload",
            files=files,
            headers={"Authorization": f"Bearer {openai_api_key}"}
        )
    
    # Check if the upload was successful
    if upload_response.status_code != 200:
        print(f"Error uploading file: {upload_response.text}")
        return upload_response.json()
    
    # Get and return the uploaded file ID
    file_info = upload_response.json()
    file_id = file_info["id"]
    print(f"File uploaded successfully with ID: {file_id}")
    return file_id

def create_lora_from_file(name, file_id, base_model):    
    # Create the LoRA model using the uploaded file
    lora_creation_request = {
        "source": file_id,  # ID of the uploaded training data
        "base_model": base_model,  # Base model to adapt
        "name": name,  # Name for the new model
        "description": "description"  # Description of the model
    }
    
    # Send request to create the model
    model_response = requests.post(
        f"{api_url}/v0/models", 
        json=lora_creation_request,
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {openai_api_key}"}
    ) 
    
    return(model_response.json())

def delete_lora(model):
    return requests.delete(f"{api_url}/v1/models/{model}",                         
      headers={"Content-Type": "application/json","Authorization": f"Bearer {os.getenv('NEBIUS_API_KEY')}"})  


def list_loras():
    return requests.get(f"{api_url}/v0/models",                         
      headers={"Content-Type": "application/json","Authorization": f"Bearer {os.getenv('NEBIUS_API_KEY')}"})  

def get_completion(model):
    completion = client.chat.completions.create(
        model=model,
        messages=[{"role":"user","content":'hi'}],  # Simple test prompt
    )
    return completion.choices[0].message.content


# Define path to the lora adapter archive file with adapter_config.json and adapter_model.safetensors
zip_file_name="/Users/sofrony/tmp/LORA.zip"
    # Upload the file and get the file ID
file_id = upload_file(zip_file_name)

# Create a new LoRA model using the file_id
create_lora_from_file("Sofrony-Test", file_id, "meta-llama/Meta-Llama-3.1-8B-Instruct")

File uploaded successfully with ID: file-df09cb37-8eda-414d-b37f-feb69e671f7d

{'name': 'Sofrony-Test-hfjV',
 'base_model': 'meta-llama/Meta-Llama-3.1-8B-Instruct',
 'source': 'file-df09cb37-8eda-414d-b37f-feb69e671f7d',
 'description': 'description',
 'created_at': 1744106413,
 'status': 'validating'}

last_lora = list_loras().json()[-1]  # list loras and get the last one
last_lora

{'type': 'text2text',
 'name': 'meta-llama/Meta-Llama-3.1-8B-Instruct-LoRa:Sofrony-Test-hfjV',
 'status': 'active',
 'status_reason': None,
 'checkpoint_id': None,
 'job_id': None,
 'file_id': 'file-df09cb37-8eda-414d-b37f-feb69e671f7d',
 'url': None,
 'created_at': 1744106413,
 'description': 'description',
 'vendor': 'meta',
 'tags': ['128K context', 'small', 'JSON mode', 'lora'],
 'use_cases': ['lora'],
 'quality': 73,
 'context_window_k': 128,
 'size_b': 8.03}

get_completion(last_lora["name"])

'How can I help you today?'

delete_lora("Sofrony-Test-hfjV").json()

{'name': 'Sofrony-Test-hfjV',
 'base_model': 'meta-llama/Meta-Llama-3.1-8B-Instruct',
 'source': 'file-df09cb37-8eda-414d-b37f-feb69e671f7d',
 'description': 'description',
 'created_at': 1744106413,
 'status': 'deleted'}