GPUtopia ai fan out hub, sends requests to workers
from ai_spider.app import app

ai_spider/__init__.py
ArcadeLabsInc/queenbee
import asyncio
import contextlib
import hashlib
import json
import logging
import os
import re
import time
from asyncio import Queue
from functools import lru_cache
from json import JSONDecodeError
from threading import RLock
from typing import Iterator, Optional, Generator, cast

import fastapi
import httpx
import starlette.websockets
import websockets
from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket, Request, HTTPException
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.requests import HTTPConnection

from .openai_types import CompletionChunk, ChatCompletion, CreateChatCompletionRequest

from fastapi.middleware.cors import CORSMiddleware

from starlette.middleware.sessions import SessionMiddleware
from sse_starlette.sse import EventSourceResponse

from .stats import StatsContainer, PUNISH_SECS
from .files import app as file_router  # Adjust the import path as needed
from .util import get_bill_to, BILLING_URL, BILLING_TIMEOUT

log = logging.getLogger(__name__)

load_dotenv()

SECRET_KEY = os.environ["SECRET_KEY"]
APP_NAME = os.environ.get("APP_NAME", "GPUTopia QueenBee")
SLOW_SECS = 5
SLOW_TOTAL_SECS = 120
PUNISH_BUSY_SECS = 30

app = FastAPI(
    title=f"{APP_NAME} API",
    description=f"{APP_NAME} inference and worker information API",
    version="1.0",
)

app.include_router(file_router, prefix='/v1', tags=['files'])

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.add_middleware(SessionMiddleware, secret_key=SECRET_KEY)


# change to "error" from "detail" to be compat with openai

@app.exception_handler(StarletteHTTPException)
async def http_exception_handler(request, exc):
    return JSONResponse(
        status_code=exc.status_code,
        content={"error": {"code": exc.status_code, "type": type(exc).__name__, "message": exc.detail}}
    )


@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request, exc):
    return JSONResponse(
        status_code=400,
        content={"error": {"code": 400, "type": type(exc).__name__, "message": str(exc)}}
    )


@app.exception_handler(AssertionError)
async def assertion_handler(request, exc):
    return JSONResponse(
        status_code=400,
        content={"error": {"code": 400, "type": type(exc).__name__, "message": str(exc)}}
    )


MODEL_MAP = [
    {
        "re": r"vicuna.*\b7b\b.*\bq4",
        "hf": "TheBloke/vicuna-7B-v1.5-GGUF:Q4_K_M",
        "web": "webgpu/vicuna-v1-7b-q4f32_0",
    },
    {
        "re": r"llama-2.*\b13b\b.*\bq4",
        "hf": "TheBloke/Llama-2-13B-chat-GGUF:Q4_K_M",
        "web": "webgpu/Llama-2-13b-chat-hf-q4f32_1"
    },
    {
        "re": r"llama-2.*\b17b\b.*\bq4",
        "hf": "TheBloke/Llama-2-7B-Chat-GGUF:Q4_K_M",
        "web": "webgpu/Llama-2-7b-chat-hf-q4f32_1"
    }
]


def alt_models(model) -> dict | None:
    for ent in MODEL_MAP:
        if re.match(ent["re"], model, re.I):
            return ent
    return None


async def check_creds_and_funds(request):
    bill_to_token = get_bill_to(request)

    if bill_to_token == os.environ.get("BYPASS_TOKEN"):
        return True

    command = dict(
        command="check",
        bill_to_token=bill_to_token,
    )

    try:
        async with httpx.AsyncClient() as client:
            res = await client.post(BILLING_URL, json=command, timeout=BILLING_TIMEOUT)
    except Exception as ex:
        raise HTTPException(status_code=500, detail="billing endpoint error: %s" % ex)

    if res.status_code != 200:
        log.error("bill endpoint: %s/%s", res.status_code, res.text)
        raise HTTPException(status_code=500, detail="billing endpoint error: %s/%s" % (res.status_code, res.text))

    js = res.json()

    if js.get("ok"):
        return True

    log.debug("no balance for %s", bill_to_token)
    raise HTTPException(status_code=422, detail="insufficient funds in account, or incorrect auth token")


def get_key(sock):
    if k := sock.info.get("pubkey"):
        return k
    if k := sock.info.get("worker_id"):
        return k
    return sock


g_stats = StatsContainer(key=lambda sock: get_key(sock))



def record_stats(sock, msize, usage, secs):
    g_stats.bump(sock, msize, usage, secs)


async def bill_usage(request, msize: int, usage: dict, worker_info: dict, secs: float):
    # todo: this should bill based on model size * usage
    pay_to_lnurl = worker_info.get("ln_address", worker_info.get("ln_url"))  # todo: ln_url is old.
    pay_to_auth = worker_info.get("auth_key")

    bill_to_token = get_bill_to(request)

    command = dict(
        command="complete",
        bill_to_token=bill_to_token,
        pay_to_lnurl=pay_to_lnurl,
        pay_to_auth=pay_to_auth,
    )

    try:
        async with httpx.AsyncClient() as client:
            res = await client.post(BILLING_URL, json=command, timeout=BILLING_TIMEOUT)

        log.info("bill %s/%s/%s to: (%s), pay to: (%s)", usage, msize, secs, bill_to_token, worker_info)

        if res.status_code != 200:
            log.error("bill endpoint: %s/%s", res.status_code, res.text)
    except Exception as ex:
        log.error(f"billing error ({ex}): {usage}/{msize}/{secs} to: ({bill_to_token}), pay to: ({worker_info})")

    return True


async def check_bill_usage(request, msize: int, js: dict, worker_info: dict, secs: float):
    if js.get("usage"):
        await bill_usage(request, msize, js["usage"], worker_info, secs)


@app.get("/worker/stats", tags=["worker"])
async def worker_stats() -> dict:
    """Simple free/busy total count"""
    mgr = get_reg_mgr()
    return mgr.worker_stats()


@app.get("/worker/detail", tags=["worker"])
async def worker_detail(query: Optional[str] = None) -> list:
    """List of all workers, with anonymized info"""
    mgr = get_reg_mgr()
    return mgr.worker_anon_detail(query=query)


@app.post("/v1/chat/completions")
async def create_chat_completion(
        request: Request,
        body: CreateChatCompletionRequest,
) -> ChatCompletion:
    """Openai compatible chat completion endpoint."""
    await check_creds_and_funds(request)

    worker_type = worker_type_from_model_name(body.model)

    msize = get_model_size(body.model)
    mgr = get_reg_mgr()
    gpu_filter = body.gpu_filter

    try:
        try:
            with mgr.get_socket_for_inference(msize, worker_type, gpu_filter) as ws:
                return await do_inference(request, body, ws)
        except (fastapi.WebSocketDisconnect, HTTPException) as ex:
            if type(ex) is HTTPException and "gguf" in ex.detail:
                raise
            log.error("try again: %s: ", repr(ex))
            await asyncio.sleep(0.5)
            with mgr.get_socket_for_inference(msize, worker_type, gpu_filter) as ws:
                return await do_inference(request, body, ws)
    except HTTPException as ex:
        log.error("inference failed : %s", repr(ex))
        raise
    except TimeoutError as ex:
        log.error("inference failed : %s", repr(ex))
        raise HTTPException(408, detail=repr(ex))
    except AssertionError as ex:
        log.error("inference failed : %s", repr(ex))
        raise HTTPException(400, detail=repr(ex))
    except Exception as ex:
        log.exception("unknown error : %s", repr(ex))
        raise HTTPException(500, detail=repr(ex))


def worker_type_from_model_name(model):
    worker_type = "cli"
    if model.startswith("webgpu/"):
        worker_type = "web"
    elif alt_models(model):
        worker_type = "any"
    return worker_type


def augment_reply(body: CreateChatCompletionRequest, js, prev_js={}):
    # todo: this should be used to VALIDATE the reply, not "fix" it!
    # todo: this all happens because the web-worker is total hack, need to clean it up

    if not js.get("model"):
        js["model"] = prev_js.get("model", body.model)

    if not js.get("object"):
        js["object"] = "chat.completion"

    if not js.get("created"):
        js["created"] = prev_js.get("created", int(time.time()))

    if not js.get("id"):
        js["id"] = prev_js.get("id", os.urandom(16).hex())

    inp = sum(len(msg.content) for msg in body.messages) // 3
    c0 = js["choices"][0]

    if c0.get("finish_reason"):
        if not js.get("usage"):
            msg = c0.get("message", {}).get("content", "")
            out = len(msg) // 3
            js["usage"] = dict(
                prompt_tokens=int(inp),
                completion_tokens=int(out),
                total_tokens=int(inp + out),
            )


def get_stream_final(body: CreateChatCompletionRequest, prev_js, content_len):
    js = {"choices": [{"finish_reason": "stop", "content": ""}]}
    augment_reply(body, js, prev_js)
    inp = js["usage"]["prompt_tokens"]
    out = content_len // 3
    js["usage"].update(dict(
        completion_tokens=int(out),
        total_tokens=int(inp + out),
    ))
    return js


def punish_failure(ws: "QueueSocket", reason, secs=PUNISH_SECS):
    log.info("punishing: %s, worker: %s", reason, ws.info)
    g_stats.punish(ws, secs)


def is_web_worker(info):
    return info.get("worker_version", "") == "web" or sum([1 for _ in info.get("web_gpus", [])])


def adjust_model_for_worker(model, info) -> str:
    model = model.strip()

    want_type = worker_type_from_model_name(model)

    is_web = is_web_worker(info)

    assert not (want_type == "web" and not is_web), f"invalid model for worker: {model}"

    if want_type != "any":
        return model

    alt = alt_models(model)

    assert alt, f"invalid model name {model}"

    if is_web:
        return alt["web"]

    return alt["hf"]


async def do_inference(request: Request, body: CreateChatCompletionRequest, ws: "QueueSocket"):
    # be sure we don't alter the original request, so it can be retried
    body = body.model_copy()

    body.model = adjust_model_for_worker(body.model, ws.info)

    msize = get_model_size(body.model)

    if not msize:
        err = "unknown model size: %s" % body.model
        log.error(err)
        raise HTTPException(status_code=400, detail=err)

    while ws.results.qsize():
        ws.results.get_nowait()

    await ws.queue.put({
        "openai_url": "/v1/chat/completions",
        "openai_req": body.model_dump(mode="json")
    })

    if ws.info.get("current_model", "") == body.model:
        already_loaded = True
    else:
        ws.info["current_model"] = body.model
        # None == unsure.  We're not sure if it's already loaded, because we have no protocol for that yet
        already_loaded = None

    start_time = time.monotonic()

    if body.stream:
        async def stream() -> Iterator[CompletionChunk]:
            prev_js = {}
            total_content_len = 0
            prev_time = start_time
            while True:
                try:
                    js: dict = await asyncio.wait_for(ws.results.get(), timeout=body.timeout)

                    log.debug("got msg %s", js)

                    if not js and prev_js:
                        fin = get_stream_final(body, prev_js, total_content_len)
                        yield json.dumps(fin)
                        # bill when stream is done, for now, could actually charge per stream, but whatever
                        end_time = time.monotonic()
                        asyncio.create_task(check_bill_usage(request, msize, fin, ws.info, end_time - start_time))
                        record_stats(ws, msize, fin.get("usage"), end_time - start_time)
                        break

                    c0 = js["choices"][0]
                    # fix bug:                    
                    if c0.get("message") and not c0.get("delta"):
                        c0["delta"] = c0.pop("message")

                    c_len = len(c0.get("delta", {}).get("content", ""))
                    cur_time = time.monotonic()
                    token_time = cur_time - prev_time
                    prev_time = cur_time

                    # first token can be long (load time), but between tokens should be fast!
                    if token_time > SLOW_SECS and total_content_len > 0:
                        punish_failure(ws, "slow worker, try again: %s" % token_time)
                        raise HTTPException(status_code=438, detail="slow worker, try again")

                    total_content_len += c_len

                    augment_reply(body, js, prev_js)

                    js["object"] = "chat.completion.chunk"

                    prev_js = js

                    yield json.dumps(js)

                    if c0.get("finish_reason"):
                        asyncio.create_task(check_bill_usage(request, msize, js, ws.info, cur_time - start_time))
                        break

                    if js.get("error"):
                        log.info("got an error: %s", js["error"])
                        raise HTTPException(status_code=400, detail=json.dumps(js))
                except Exception as ex:
                    if isinstance(ex, (KeyError, IndexError)):
                        punish_failure(ws, repr(ex))
                    log.exception("error during stream")
                    yield json.dumps({"error": repr(ex), "error_type": type(ex).__name__})
                    break

        return EventSourceResponse(stream())
    else:
        js = await asyncio.wait_for(ws.results.get(), timeout=body.timeout)
        if err := js.get("error"):
            if err == "busy":
                # don't punish a busy worker for long
                punish_failure(ws, "error: %s" % err, PUNISH_BUSY_SECS)
            else:
                punish_failure(ws, "error: %s" % err)
            raise HTTPException(status_code=400, detail=json.dumps(js))
        end_time = time.monotonic()
        augment_reply(body, js)
        asyncio.create_task(check_bill_usage(request, msize, js, ws.info, end_time - start_time))
        if usage := js.get("usage"):
            record_stats(ws, msize, usage, end_time - start_time)
        if already_loaded and (end_time - start_time) > SLOW_TOTAL_SECS:
            # mark bad... too slow
            punish_failure(ws, "slow inference: %s" % int(end_time - start_time))
        return js


def get_model_size(model_mame):
    mod = model_mame
    m = re.search(r"(\d)+b(.*)", mod.lower())
    # todo: actually have a nice mapping of model sizes
    if m:
        msize = int(m[1])
        mod = m[2]
    else:
        msize = 13
    m = re.search(r"[Qq](\d)+[_f.-]", mod.lower())
    if m:
        quant = int(m[1])
        if quant == 2:
            msize *= 0.4
        elif quant == 3:
            msize *= 0.5
        elif quant == 4:
            msize *= 0.6
        elif quant == 5:
            msize *= 0.7
        elif quant == 6:
            msize *= 0.8
        elif quant == 8:
            msize *= 1.0
        else:
            # f16
            msize *= 2
    return msize


class QueueSocket(WebSocket):
    queue: Queue
    results: Queue
    info: dict


def anon_info(ent, **fin):
    info = ent.info
    fin["worker_version"] = info.get("worker_version")
    nv_gpu_cnt = sum([1 for _ in info.get("nv_gpus", [])])
    cl_gpu_cnt = sum([1 for _ in info.get("cl_gpus", [])])
    web_gpu_cnt = sum([1 for _ in info.get("web_gpus", [])])
    fin["gpu_cnt"] = max(nv_gpu_cnt, cl_gpu_cnt, web_gpu_cnt)
    fin["perf"] = g_stats.perf(ent, 5)
    fin["cnt"] = g_stats.cnt(ent)
    return fin


class WorkerManager:
    def __init__(self):
        self.lock = RLock()
        self.socks = dict[QueueSocket, dict]()
        # todo: clean this up
        # im using it *or* someone said they were busy/not-busy
        self.busy = dict()
        # im actually using it
        self.very_busy = set()

    def register_js(self, sock: QueueSocket, info: dict):
        self.socks[sock] = info

    def drop_worker(self, sock):
        self.socks.pop(sock, None)
        self.busy.pop(sock, None)

    @contextlib.contextmanager
    def get_socket_for_inference(self, msize: int, worker_type: str, gpu_filter={}) -> Generator[
        QueueSocket, None, None]:
        # msize is params adjusted by quant level with a heuristic

        # nv gpu memory is reported in MB
        gpu_needed = msize * 1000

        disk_needed = msize * 1000 * 1.5

        # cpu memory is reported in bytes, it's ok to have less... cuz llama.cpp is good about that
        cpu_needed = msize * 1000000000 * 0.75

        if worker_type in ("any", "web"):
            cpu_needed = min(cpu_needed, 8000000000)

        with self.lock:
            good = []
            close = []
            for sock, info in self.socks.items():
                cpu_vram = info.get("vram", 0)
                disk_space = info.get("disk_space", 0)
                nv_gpu_ram = sum([el.get("memory", 0) for el in info.get("nv_gpus", [])])
                cl_gpu_ram = sum([el.get("memory", 0) for el in info.get("cl_gpus", [])])
                have_web_gpus = is_web_worker(info)

                if wid := gpu_filter.get("pubkey", gpu_filter.get("worker_id")):
                    # used for the autopay cron
                    if info.get("worker_id") != wid:
                        continue

                if uid := gpu_filter.get("user_id"):
                    if (info.get("auth_key") != "uid:" + str(uid)) and info.get("user_id") != uid:
                        continue

                if worker_type in ("any", "web"):
                    if cpu_needed < cpu_vram and have_web_gpus:
                        # very little ability to check here
                        # todo: end the whole self reporting thing and just bench
                        good.append(sock)

                if worker_type in ("any", "cli"):
                    if gpu_needed < nv_gpu_ram and cpu_needed < cpu_vram and disk_needed < disk_space:
                        good.append(sock)
                    elif gpu_needed < cl_gpu_ram and cpu_needed < cpu_vram and disk_needed < disk_space:
                        good.append(sock)
                    elif gpu_needed < nv_gpu_ram * 1.2 and cpu_needed < cpu_vram and disk_needed < disk_space:
                        close.append(sock)
                    elif gpu_needed < cl_gpu_ram * 1.2 and cpu_needed < cpu_vram and disk_needed < disk_space:
                        close.append(sock)

            if not good and not close:
                assert False, "No workers available"

            if len(good):
                num = g_stats.pick_best(good, msize)
                choice = good[num]
            elif len(close):
                num = g_stats.pick_best(close, msize)
                choice = close[num]

            info = self.socks.pop(choice, None)
            self.busy[choice] = info
            self.very_busy.add(choice)

        choice.info = info

        try:
            yield choice
        finally:
            with self.lock:
                self.socks[choice] = info
                self.busy.pop(choice, None)
                self.very_busy.discard(choice)

    def set_busy(self, sock, val):
        with self.lock:
            if sock in self.very_busy:
                return
            if val:
                info = self.socks.pop(sock, None)
                if info:
                    self.busy[sock] = info
            else:
                info = self.busy.pop(sock, None)
                if info:
                    self.socks[sock] = info

    def worker_stats(self):
        connected = len(self.socks)
        busy = len(self.busy)
        return dict(
            connected=connected,
            busy=busy
        )

    def worker_detail(self):
        all = {}
        for ent in self.socks:
            all[id(ent)] = ent.info.copy()
            all[id(ent)]["busy"] = False
            all[id(ent)]["perf"] = g_stats.perf(ent, 5)
        for ent in self.busy:
            all[id(ent)] = ent.info.copy()
            all[id(ent)]["busy"] = True
            all[id(ent)]["perf"] = g_stats.perf(ent, 5)
        return all

    def worker_anon_detail(self, *, query):
        all = {}
        for ent in self.socks:
            if query and not self.filter_match(ent.info, query):
                continue
            all[id(ent)] = anon_info(ent, busy=False)
        for ent in self.busy:
            if query and not self.filter_match(ent.info, query):
                continue
            all[id(ent)] = anon_info(ent, busy=True)
        return list(all.values())

    @staticmethod
    def filter_match(info, query):
        if query and info.get("auth_key") == query:
            return True
        if query and info.get("worker_id") == query:
            return True
        if query and info.get("pubkey") == query:
            return True
        return False


g_reg_mgr: Optional[WorkerManager] = None


def get_reg_mgr() -> WorkerManager:
    global g_reg_mgr
    if not g_reg_mgr:
        g_reg_mgr = WorkerManager()
    return g_reg_mgr


def get_ip(request: HTTPConnection):
    ip = request.client.host
    if not (ip.startswith("192.168") or ip.startswith("10.10")):
        return ip
    if ip := request.headers.get("x-real-ip"):
        return ip
    if ip := request.headers.get("x-forwarded-for"):
        return ip
    return ip


def validate_worker_info(js):
    pk = js.get("pubkey", None)
    sig = js.pop("sig", None)
    # todo: raise an error if invalid sig


@app.websocket("/worker")
async def worker_connect(websocket: WebSocket):
    # request dependencies don't work with websocket, so just roll our own
    try:
        await websocket.accept()
        js = await websocket.receive_json()
    except (websockets.ConnectionClosedOK, RuntimeError, starlette.websockets.WebSocketDisconnect) as ex:
        log.debug("aborted connection before message: %s", repr(ex))
        return

    websocket = cast(QueueSocket, websocket)

    validate_worker_info(js)

    # turn it into a queuesocket by adding "info" and a queue
    websocket.info = js

    # get the source ip, for long-term punishment of bad actors
    req = HTTPConnection(websocket.scope)
    websocket.info["ip"] = get_ip(req)

    websocket.queue = Queue()
    websocket.results = Queue()

    # debug: log everything
    log.debug("connected: %s", websocket.info)

    mgr = get_reg_mgr()
    mgr.register_js(sock=websocket, info=js)

    while True:
        try:
            pending = [asyncio.create_task(ent) for ent in [websocket.queue.get(), websocket.receive_json()]]
            done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
            for fut in done:
                try:
                    action = await fut
                    # this is either a queued result, a new request or a busy statement
                    # could distinguish them above with some intermediate functions that return tuples, but no need yet
                    if "openai_req" in action:
                        log.info("action %s", action)
                        while not websocket.results.empty():
                            websocket.results.get_nowait()
                        await websocket.send_json(action)
                    elif "busy" in action:
                        log.info("action %s", action)
                        mgr.set_busy(websocket, action.get("busy"))
                    else:
                        log.debug("action %s", action)
                        await websocket.results.put(action)
                except JSONDecodeError:
                    log.warning("punish worker failure")
                    punish_failure(websocket, "json decode error")
                    # continue so we don't get a new ws, and lose stats
                    # todo: do this by inbound ip instead of ws handle, so they persist across connections
                    # then we can disconnect here if we want, and even block reconn for a while
                except (websockets.ConnectionClosedOK, RuntimeError, starlette.websockets.WebSocketDisconnect):
                    # disconnect means drop out of the loop
                    raise
                except Exception:
                    # other exceptions could be my logic error, try again until disconnected
                    log.exception("exception in loop")
                finally:
                    # clean up futures
                    for ent in pending:
                        ent.cancel()
        except (websockets.ConnectionClosedOK, RuntimeError, starlette.websockets.WebSocketDisconnect):
            log.info("dropped worker")
            break
    mgr.drop_worker(websocket)

ai_spider/app.py
ArcadeLabsInc/queenbee
import os
from functools import wraps

import httpx
from botocore.exceptions import ClientError
from dotenv import load_dotenv
from fastapi import File, UploadFile, HTTPException, Depends, Request, APIRouter
from fastapi.responses import StreamingResponse
import boto3

from ai_spider.util import get_bill_to, BILLING_URL, BILLING_TIMEOUT

app = APIRouter()

load_dotenv()

g_s3 = None


def s3():
    global g_s3
    if not g_s3:
        g_s3 = boto3.client('s3',
                            aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
                            aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"))

    return g_s3

bucket_name = 'gputopia-user-files'


def handle_aws_exceptions(route_function):
    @wraps(route_function)
    async def wrapper(*args, **kwargs):
        try:
            return await route_function(*args, **kwargs)
        except ClientError as e:
            error_code = e.response['Error']['Code']
            if error_code == "NoSuchKey":
                raise HTTPException(status_code=404, detail=str(e))
            else:
                raise HTTPException(status_code=500, detail=str(e))
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))
    return wrapper

async def check_bearer_token(request: Request) -> str:
    bill_to_token = get_bill_to(request)

    if bill_to_token == os.environ.get("BYPASS_TOKEN"):
        return "bypass"

    command = dict(
        command="check",
        bill_to_token=bill_to_token,
    )

    try:
        res = httpx.post(BILLING_URL, json=command, timeout=BILLING_TIMEOUT)
    except Exception as ex:
        raise HTTPException(status_code=500, detail="billing endpoint error: %s" % ex)

    js = res.json()

    if js.get("user_id"):
        return js.get("user_id")

    raise HTTPException(status_code=401, detail="Invalid token")


@app.get("/v1/files")
@handle_aws_exceptions
async def list_files(user_id: str = Depends(check_bearer_token)):
    user_folder = f"{user_id}/"
    file_objects = s3().list_objects(Bucket=bucket_name, Prefix=user_folder)['Contents']
    return {"data": [{"id": obj["Key"][len(user_folder):], "bytes": obj["Size"], "created_at": obj["LastModified"].timestamp()} for obj in
                     file_objects], "object": "list"}


@app.post("/v1/files")
@handle_aws_exceptions
async def upload_file(file: UploadFile = File(...), purpose: str = "", user_id: str = Depends(check_bearer_token)):
    user_folder = f"{user_id}/"
    s3().upload_fileobj(file.file, bucket_name, f"{user_folder}{file.filename}")
    # You may also need to store additional metadata, such as purpose, in a database
    return {"id": file.filename, "object": "file", "bytes": file.size, "filename": file.filename, "purpose": purpose,
            "status": "uploaded"}


@app.delete("/v1/files/{file_id}")
@handle_aws_exceptions
async def delete_file(file_id: str, user_id: str = Depends(check_bearer_token)):
    try:
        user_folder = f"{user_id}/"
        s3().delete_object(Bucket=bucket_name, Key=f"{user_folder}{file_id}")
        return {"id": file_id, "object": "file", "deleted": True}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/files/{file_id}")
@handle_aws_exceptions
async def retrieve_file(file_id: str, user_id: str = Depends(check_bearer_token)):
    try:
        user_folder = f"{user_id}/"
        # Retrieve file metadata from S3 and potentially other metadata from the database
        file_meta = s3().head_object(Bucket=bucket_name, Key=f"{user_folder}{file_id}")
        return {"id": file_id, "object": "file", "bytes": file_meta["ContentLength"],
                "created_at": file_meta["LastModified"].timestamp(), "filename": file_id}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/v1/files/{file_id}/content")
@handle_aws_exceptions
async def retrieve_file_content(file_id: str, user_id: str = Depends(check_bearer_token)):
    user_folder = f"{user_id}/"
    file_content = s3().get_object(Bucket=bucket_name, Key=f"{user_folder}{file_id}")['Body']
    return StreamingResponse(content=file_content, media_type="application/octet-stream")

ai_spider/files.py
ArcadeLabsInc/queenbee
from typing import List, Optional, Dict, Union

from pydantic import BaseModel, Field
from typing_extensions import TypedDict, NotRequired, Literal


class EmbeddingUsage(TypedDict):
    prompt_tokens: int
    total_tokens: int


class EmbeddingData(TypedDict):
    index: int
    object: str
    embedding: List[float]


class Embedding(TypedDict):
    object: Literal["list"]
    model: str
    data: List[EmbeddingData]
    usage: EmbeddingUsage


class CompletionLogprobs(TypedDict):
    text_offset: List[int]
    token_logprobs: List[Optional[float]]
    tokens: List[str]
    top_logprobs: List[Optional[Dict[str, float]]]


class CompletionChoice(TypedDict):
    text: str
    index: int
    logprobs: Optional[CompletionLogprobs]
    finish_reason: Optional[str]


class CompletionUsage(TypedDict):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class CompletionChunk(TypedDict):
    id: str
    object: Literal["text_completion"]
    created: int
    model: str
    choices: List[CompletionChoice]


class Completion(TypedDict):
    id: str
    object: Literal["text_completion"]
    created: int
    model: str
    choices: List[CompletionChoice]
    usage: CompletionUsage


class ChatCompletionMessage(TypedDict):
    role: Literal["assistant", "user", "system"]
    content: str
    user: NotRequired[str]


class ChatCompletionChoice(TypedDict):
    index: int
    message: ChatCompletionMessage
    finish_reason: Optional[str]


class ChatCompletion(TypedDict):
    id: str
    object: Literal["chat.completion"]
    created: int
    model: str
    choices: List[ChatCompletionChoice]
    usage: CompletionUsage
    ln_url: NotRequired[str]


class ChatCompletionChunkDeltaEmpty(TypedDict):
    pass


class ChatCompletionChunkDelta(TypedDict):
    role: NotRequired[Literal["assistant"]]
    content: NotRequired[str]


class ChatCompletionChunkChoice(TypedDict):
    index: int
    delta: Union[ChatCompletionChunkDelta, ChatCompletionChunkDeltaEmpty]
    finish_reason: Optional[str]


class ChatCompletionChunk(TypedDict):
    id: str
    model: str
    object: Literal["chat.completion.chunk"]
    created: int
    choices: List[ChatCompletionChunkChoice]


class ChatCompletionRequestMessage(BaseModel):
    role: Literal["system", "user", "assistant"] = Field(
        default="user", description="The role of the message."
    )
    content: str = Field(default="", description="The content of the message.")


model_field = Field(description="The model to use for generating completions.", default=None)
max_tokens_field = Field(
    default=200, ge=1, description="The maximum number of tokens to generate."
)
temperature_field = Field(
    default=0.8,
    ge=0.0,
    le=2.0,
    description="Adjust the randomness of the generated text.\n\n"
                + "Temperature is a hyperparameter that controls the randomness of the generated text. It affects the probability distribution of the model's output tokens. A higher temperature (e.g., 1.5) makes the output more random and creative, while a lower temperature (e.g., 0.5) makes the output more focused, deterministic, and conservative. The default value is 0.8, which provides a balance between randomness and determinism. At the extreme, a temperature of 0 will always pick the most likely next token, leading to identical outputs in each run.",
)
top_p_field = Field(
    default=0.95,
    ge=0.0,
    le=1.0,
    description="Limit the next token selection to a subset of tokens with a cumulative probability above a threshold P.\n\n"
                + "Top-p sampling, also known as nucleus sampling, is another text generation method that selects the next token from a subset of tokens that together have a cumulative probability of at least p. This method provides a balance between diversity and quality by considering both the probabilities of tokens and the number of tokens to sample from. A higher value for top_p (e.g., 0.95) will lead to more diverse text, while a lower value (e.g., 0.5) will generate more focused and conservative text.",
)
stop_field = Field(
    default=None,
    description="A list of tokens at which to stop generation. If None, no stop tokens are used.",
)
stream_field = Field(
    default=False,
    description="Whether to stream the results as they are generated. Useful for chatbots.",
)
top_k_field = Field(
    default=40,
    ge=0,
    description="Limit the next token selection to the K most probable tokens.\n\n"
                + "Top-k sampling is a text generation method that selects the next token only from the top k most likely tokens predicted by the model. It helps reduce the risk of generating low-probability or nonsensical tokens, but it may also limit the diversity of the output. A higher value for top_k (e.g., 100) will consider more tokens and lead to more diverse text, while a lower value (e.g., 10) will focus on the most probable tokens and generate more conservative text.",
)
repeat_penalty_field = Field(
    default=1.1,
    ge=0.0,
    description="A penalty applied to each token that is already generated. This helps prevent the model from repeating itself.\n\n"
                + "Repeat penalty is a hyperparameter used to penalize the repetition of token sequences during text generation. It helps prevent the model from generating repetitive or monotonous text. A higher value (e.g., 1.5) will penalize repetitions more strongly, while a lower value (e.g., 0.9) will be more lenient.",
)
presence_penalty_field = Field(
    default=0.0,
    ge=-2.0,
    le=2.0,
    description="Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics.",
)
frequency_penalty_field = Field(
    default=0.0,
    ge=-2.0,
    le=2.0,
    description="Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.",
)
mirostat_mode_field = Field(
    default=0,
    ge=0,
    le=2,
    description="Enable Mirostat constant-perplexity algorithm of the specified version (1 or 2; 0 = disabled)",
)
mirostat_tau_field = Field(
    default=5.0,
    ge=0.0,
    le=10.0,
    description="Mirostat target entropy, i.e. the target perplexity - lower values produce focused and coherent text, larger values produce more diverse and less coherent text",
)
mirostat_eta_field = Field(
    default=0.1, ge=0.001, le=1.0, description="Mirostat learning rate"
)


class CreateChatCompletionRequest(BaseModel):
    messages: List[ChatCompletionRequestMessage] = Field(
        default=[], description="A list of messages to generate completions for."
    )
    max_tokens: int = max_tokens_field
    temperature: float = temperature_field
    top_p: float = top_p_field
    mirostat_mode: int = mirostat_mode_field
    mirostat_tau: float = mirostat_tau_field
    mirostat_eta: float = mirostat_eta_field
    stop: Optional[List[str]] = stop_field
    stream: bool = stream_field
    presence_penalty: Optional[float] = presence_penalty_field
    frequency_penalty: Optional[float] = frequency_penalty_field
    logit_bias: Optional[Dict[str, float]] = Field(None)

    # ignored or currently unsupported
    model: Optional[str] = model_field
    n: Optional[int] = 1
    user: Optional[str] = Field(None)

    # llama.cpp specific parameters
    top_k: int = top_k_field
    repeat_penalty: float = repeat_penalty_field
    logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None)
    timeout: int = 60 * 10
    gpu_filter: dict = {}

ai_spider/openai_types.py
ArcadeLabsInc/queenbee
import math
import random
import time
from collections import defaultdict

STATS_EMA_ALPHA = 0.9
PUNISH_SECS = 60 * 15
PUNISH_BAD_PERF = 9999


class StatsBin:
    def __init__(self, alpha):
        self.val = None
        self.alpha = alpha

    def bump(self, secs):
        if self.val is None:
            self.val = secs
        else:
            self.val = self.alpha * self.val + (1 - self.alpha) * secs


class StatsWorker:
    def __init__(self, alpha):
        self.stats: dict[int, StatsBin] = defaultdict(lambda: StatsBin(alpha))
        self.bad = None
        self.cnt = 0

    def bump(self, msize, usage, secs):
        self.cnt += 1
        toks = usage.get("total_tokens")
        # similar-sized models are lumped together
        msize_bin = round(math.sqrt(msize))
        # similar-sized token counts are lumped together too
        self.stats[msize_bin].bump(secs / toks)
        # all is forgiven
        self.bad = None

    def perf(self, msize):
        if self.bad and self.bad > time.monotonic():
            return PUNISH_BAD_PERF
        if not self.stats:
            return None
        msize_bin = round(math.sqrt(msize))
        exact = self.stats.get(msize_bin)
        if exact:
            return exact.val
        close_bin = sorted(self.stats.keys())[-1]
        close = self.stats.get(close_bin)
        # adjustment for model size is (msize/closest_size) ** 1.5
        if msize_bin < close_bin:
            approx = close.val * (msize / (close_bin ** 2.0)) ** 1.5
        else:
            # more conservative for upsizing
            approx = close.val * (msize / (close_bin ** 2.0)) ** 1.8
        return approx

    def punish(self, secs):
        self.bad = time.monotonic() + secs


# 2 == very skewed (prefer first, but the rest are skewed to the front)
# 1 == not skewed (prefer first, but the rest are even)

POWER = 2


class StatsContainer:
    def __init__(self, alpha=STATS_EMA_ALPHA, key=None):
        self.stats: dict[str, StatsWorker] = defaultdict(lambda: StatsWorker(alpha))
        self.all = StatsWorker(alpha)
        ident = lambda k: k
        self.key_func = key or ident

    def bump(self, key, msize, usage, secs):
        key = self.key_func(key)
        self.stats[key].bump(msize, usage, secs)
        self.all.bump(msize, usage, secs)

    def perf(self, key, msize):
        key = self.key_func(key)
        wrk = self.stats.get(key)
        if wrk:
            ret = wrk.perf(msize)
            if ret:
                return ret
        # assume average perf
        return self.all.perf(msize) or 0

    def pick_best(self, choices, msize):
        ordered = sorted(enumerate(choices), key=lambda ent: self.perf(ent[1], msize))
        # simple skewed range between 0 and 0.5
        pick = int(max(random.random() ** POWER - 0.5, 0) * (len(choices) * 2))
        return ordered[pick][0]

    def punish(self, key, secs=PUNISH_SECS):
        key = self.key_func(key)
        wrk = self.stats.get(key)
        if wrk:
            wrk.punish(secs)

    def cnt(self, key):
        key = self.key_func(key)
        wrk = self.stats.get(key)
        if wrk:
            return wrk.cnt
        return 0


ai_spider/stats.py
ArcadeLabsInc/queenbee
import os

from dotenv import load_dotenv

load_dotenv()


def get_bill_to(request):
    req_user = request.headers.get("Authorization")
    bill_to_token = ""
    if req_user and " " in req_user:
        bill_to_token = req_user.split(" ")[1]
    return bill_to_token


BILLING_URL = os.environ["BILLING_URL"]
BILLING_TIMEOUT = 20

ai_spider/util.py
ArcadeLabsInc/queenbee
from starlette.websockets import WebSocket

from ai_spider.app import alt_models, adjust_model_for_worker, WorkerManager, get_model_size


def test_alt_models():
    alts = alt_models("vicuna-7B-v1.5-GGUF:Q4_K_M")
    assert alts
    assert alts["hf"] == "TheBloke/vicuna-7B-v1.5-GGUF:Q4_K_M"

    alts = alt_models("vicuna-7B-q4")
    assert alts
    assert alts["hf"] == "TheBloke/vicuna-7B-v1.5-GGUF:Q4_K_M"

    alts = alt_models("webgpu/vicuna-7B-q4")
    assert not alts


def test_adjust():
    info = {"web_gpus": [{"name": "3090"}]}
    ok = adjust_model_for_worker("vicuna-7B-v1.5-GGUF:Q4_K_M", info)
    assert ok == "webgpu/vicuna-v1-7b-q4f32_0"
    info = {"nv_gpus": [{"name": "3090"}]}
    ok = adjust_model_for_worker("vicuna-7B-v1.5-GGUF:Q4_K_M", info)
    assert ok == "TheBloke/vicuna-7B-v1.5-GGUF:Q4_K_M"


def test_get_sock():
    sock1 = lambda: None
    sock1.info = {}
    info = {"web_gpus": [{"name": "3090"}], "disk_space": 50000, "vram": 8000000000}
    sock2 = lambda: None
    sock2.info = {}
    info2 = {"nv_gpus": [{"name": "3090", "memory": 5000}], "disk_space": 50000, "vram": 8000000000}
    mgr = WorkerManager()
    siz = get_model_size("vicuna-7B-v1.5-GGUF:Q4_K_M")
    assert siz < 5
    mgr.register_js(sock1, info)
    mgr.register_js(sock2, info2)
    with mgr.get_socket_for_inference(siz, "cli", {}) as ws:
        assert ws == sock2
    with mgr.get_socket_for_inference(siz, "web", {}) as ws:
        assert ws == sock1
    with mgr.get_socket_for_inference(siz, "any", {}) as ws:
        adj = adjust_model_for_worker("vicuna-7B-v1.5-GGUF:Q4_K_M", ws.info)
        assert "/" in adj

tests/test_alt_models.py
ArcadeLabsInc/queenbee
from collections import defaultdict

from ai_spider.stats import StatsContainer


def test_stats():
    s = StatsContainer()
    s.bump("id", 7, dict(total_tokens=2), 1)
    s.bump("id", 7, dict(total_tokens=2), 1)
    assert s.perf("id", 7) == 1 / 2
    # worse than 1/2 but better than 1 for no-info on 13 gb model
    assert 1 / 2 < s.perf("id", 13) < 1


def test_stats_k():
    s = StatsContainer(key=lambda e: e["id"])
    s.bump({"id": 1}, 7, dict(total_tokens=2), 1)
    s.bump({"id": 1, "whatever": 5}, 7, dict(total_tokens=2), 1)
    s.bump({"id": 2}, 7, dict(total_tokens=4), 1)
    assert s.perf({"id": 1, "ignore": 2}, 7) == 1 / 2
    assert s.perf({"id": 2}, 7) == 1 / 4


def usually(f):
    ent = defaultdict(lambda: 0)
    for _ in range(1000):
        z = f()
        ent[z] += 1
    vals = sorted(ent.items(), key=lambda kv: kv[1], reverse=True)

    # more than double the 2nd best
    assert vals[0][1] > vals[1][1] * 2

    return vals[0][0]


def test_stats_pick():
    s = StatsContainer()
    assert s.pick_best(["id1"], 1) == 0

    s.bump("id1", 7, dict(total_tokens=100), 1)
    s.bump("id2", 7, dict(total_tokens=100), 10)
    s.bump("id3", 7, dict(total_tokens=100), 100)
    assert usually(lambda: s.pick_best(["id1", "id2", "id3"], 13)) == 0
    assert usually(lambda: s.pick_best(["id3", "id2", "id1"], 13)) == 2

    s.punish("id1")
    
    assert usually(lambda: s.pick_best(["id3", "id2", "id1"], 13)) == 1


tests/test_stats.py
ArcadeLabsInc/queenbee
import os

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from moto import mock_s3
import boto3
from ai_spider.files import app as router, bucket_name
from util import set_bypass_token

set_bypass_token()
app = FastAPI()
app.include_router(router)

client = TestClient(app)


@pytest.fixture
def aws_credentials():
    """Mocked AWS Credentials for moto."""
    return {
        "aws_access_key_id": "testing",
        "aws_secret_access_key": "testing",
        "aws_session_token": "testing",
    }


@pytest.fixture
def s3_client(aws_credentials):
    with mock_s3():
        cli = boto3.client('s3', **aws_credentials)
        cli.create_bucket(Bucket=bucket_name)
        yield cli


def test_file_operations(s3_client):
    # Upload a file
    token = os.environ["BYPASS_TOKEN"]
    headers = {"authorization": "bearer: " + token}

    response = client.post("/v1/files", files={"file": ("test_file.txt", "some content")}, data={"purpose": "fine-tune"}, headers=headers)
    assert response.status_code == 200
    file_id = response.json()["id"]

    # List files and check if the uploaded file is listed
    response = client.get("/v1/files", headers=headers)
    assert response.status_code == 200
    assert any(f["id"] == file_id for f in response.json()["data"])

    # Get the content of the uploaded file
    response = client.get(f"/v1/files/{file_id}/content", headers=headers)
    assert response.status_code == 200
    assert response.content == b"some content"

    # Delete the uploaded file
    response = client.delete(f"/v1/files/{file_id}", headers=headers)
    assert response.status_code == 200
    assert response.json()["deleted"] == True

    # Check if the file is deleted by trying to fetch its content
    response = client.get(f"/v1/files/{file_id}/content", headers=headers)
    assert response.status_code == 404

tests/test_upload.py
ArcadeLabsInc/queenbee
import asyncio
import contextlib
import json
import os

import time
from dataclasses import dataclass
from multiprocessing import Process
from typing import Any
from unittest.mock import patch, AsyncMock

from ai_worker.main import WorkerMain, Config
from dotenv import load_dotenv
from httpx_sse import connect_sse

from ai_spider.app import app, get_reg_mgr, g_stats
from util import set_bypass_token

from threading import Thread
import httpx
import pytest
from uvicorn import Config as UVConfig, Server
import logging as log

from ai_spider.util import BILLING_URL, BILLING_TIMEOUT

set_bypass_token()
load_dotenv()


@dataclass
class SPServer:
    url: str
    httpx: Any


@pytest.fixture(scope="module")
def sp_server():
    config = UVConfig(app=app, host="127.0.0.1", port=0, loop="asyncio")

    with patch("ai_spider.app.httpx") as cli:
        server = Server(config=config)
        thread = Thread(target=server.run)
        thread.daemon = True  # Daemon threads are abruptly stopped at shutdown
        thread.start()

        # uvicorn has no way to wait fo start
        while not server.started:
            time.sleep(.1)

        port = server.servers[0].sockets[0].getsockname()[1]

        yield SPServer(url=f"ws://127.0.0.1:{port}", httpx=cli)

    server.shutdown()


@pytest.mark.asyncio
async def test_websocket_fail(sp_server):
    ws_uri = f"{sp_server.url}/worker"
    with spawn_worker(ws_uri):

        with httpx.Client(timeout=30) as client:
            res = client.post(f"{sp_server.url}/v1/chat/completions", json={
                "model": "TheBloke/WizardLM-7B-uncensored-GGML:q4_ZZ",
                "messages": [
                    {"role": "system", "content": "you are a helpful assistant"},
                    {"role": "user", "content": "write a frog story"}
                ],
                "max_tokens": 20
            }, headers={
                "authorization": "bearer: " + os.environ["BYPASS_TOKEN"]
            }, timeout=1000)

            assert res.status_code >= 400
            js = res.json()
            assert js.get("error")


@pytest.mark.asyncio
@patch("ai_spider.app.SLOW_SECS", 0)
async def test_websocket_slow(sp_server):
    ws_uri = f"{sp_server.url}/worker"
    with spawn_worker(ws_uri):
        with httpx.Client(timeout=30) as client:
            with connect_sse(client, "POST", f"{sp_server.url}/v1/chat/completions", json={
                "model": "TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M",
                "stream": True,
                "messages": [
                    {"role": "system", "content": "you are a helpful assistant"},
                    {"role": "user", "content": "write a story about a frog"}
                ],
                "max_tokens": 100
            }, headers={
                "authorization": "bearer: " + os.environ["BYPASS_TOKEN"]
            }, timeout=1000) as sse:
                events = [ev for ev in sse.iter_sse()]
                assert len(events) > 2
                assert json.loads(events[-1].data).get("error")


def wait_for(func, timeout=5):
    s = time.monotonic()
    last_ex = None

    def try_func():
        nonlocal last_ex
        try:
            if func():
                return func()
        except Exception as ex:
            last_ex = ex

    while not try_func() and time.monotonic() < s + timeout:
        time.sleep(0.1)

    if last_ex:
        raise last_ex


@pytest.mark.asyncio
async def test_websocket_conn(sp_server):
    token = os.environ["BYPASS_TOKEN"]
    ws_uri = f"{sp_server.url}/worker"
    sp_server.httpx.reset_mock()
    with spawn_worker(ws_uri, 2):
        with httpx.Client(timeout=30) as client:
            res = client.post(f"{sp_server.url}/v1/chat/completions", json={
                "model": "TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M",
                "messages": [
                    {"role": "system", "content": "You are a helpful assistant"},
                    {"role": "user", "content": "Write a two sentence frog story"}
                ],
                "max_tokens": 20
            }, headers={
                "authorization": "bearer: " + token
            }, timeout=1000)

            log.info("got completion")

            assert res.status_code == 200
            js = res.json()
            assert not js.get("error")
            assert js.get("usage")
            async with sp_server.httpx.AsyncClient() as cli:
                post = cli.post
            wait_for(lambda: post.called)
            post.assert_called_with(BILLING_URL,
                                    json=dict(command="complete", bill_to_token=token,
                                              pay_to_lnurl='DONT_PAY_ME', pay_to_auth=''),
                                    timeout=BILLING_TIMEOUT)

            sock = list(g_stats.stats.keys())[0]
            perf1 = g_stats.perf(sock, 7)
            assert perf1 < 10

            log.info("try again, with slow patched")

            with patch("ai_spider.app.SLOW_TOTAL_SECS", 0):
                # slow total means punish worker, but only if the model is loaded
                res = client.post(f"{sp_server.url}/v1/chat/completions", json={
                    "model": "TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M",
                    "messages": [
                        {"role": "system", "content": "You are a helpful assistant"},
                        {"role": "user", "content": "Write a two sentence frog story"}
                    ],
                    "max_tokens": 20
                }, headers={
                    "authorization": "bearer: " + token
                }, timeout=1000)
                perf2 = g_stats.perf(sock, 7)
                assert perf2 > 999
                assert res.status_code == 200


def wm_run(ws_uri, loops=1):
    wm = WorkerMain(Config(queen_url=ws_uri, loops=loops))
    asyncio.run(wm.run())


@contextlib.contextmanager
def spawn_worker(ws_uri, loops=1):
    thread = Process(target=wm_run, daemon=True, args=(ws_uri, loops))
    thread.start()

    while not get_reg_mgr().socks:
        time.sleep(0.1)

    yield

    thread.join()

    while get_reg_mgr().socks:
        time.sleep(0.1)


@pytest.mark.asyncio
async def test_websocket_stream(sp_server):
    ws_uri = f"{sp_server.url}/worker"
    with spawn_worker(ws_uri):
        with httpx.Client(timeout=30) as client:
            with connect_sse(client, "POST", f"{sp_server.url}/v1/chat/completions", json={
                "model": "TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M",
                "stream": True,
                "messages": [
                    {"role": "system", "content": "you are a helpful assistant"},
                    {"role": "user", "content": "write a story about a frog"}
                ],
                "max_tokens": 100
            }, headers={
                "authorization": "bearer: " + os.environ["BYPASS_TOKEN"]
            }, timeout=1000) as sse:
                events = [ev for ev in sse.iter_sse()]
                assert len(events) > 2
                assert json.loads(events[-1].data).get("usage").get("completion_tokens")

tests/test_worker.py
ArcadeLabsInc/queenbee
import os


def set_bypass_token():
    if "BYPASS_TOKEN" not in os.environ:
        os.environ.setdefault("BYPASS_TOKEN", "ABCDE12345")

tests/util.py
ArcadeLabsInc/queenbee