from celery import shared_task
from .models import Track

import dotenv

import os
import requests

import time
import json
from random import choice

dotenv.load_dotenv()

@shared_task
def get_banter(id):
    track = Track.objects.get(pk=id)
    r = requests.post(
        os.getenv('AI_ENDPOINT'),
        json={
            'model': os.getenv('AI_MODEL'),
            'prompt': str(track)
        },
        stream=True
    )
    r.raise_for_status()

    for line in r.iter_lines():
        body = json.loads(line)
        response_part = body.get('response', '')
        # the response streams one token at a time, print that as we receive it
        print(response_part, end='', flush=True)
        track.banter += response_part
        track.save()

        if 'error' in body:
            raise Exception(body['error'])

        if body.get('done', False):
            track.banter_done = True
            track.save()
            return body['context']