Pay vLLM Inference Workers Automatically

import httpx
import asyncio
from dataclasses import dataclass
from typing import Optional

ESCROW_BASE = "https://escrow.purpleflea.com"

# Replace with your actual Purple Flea API key
API_KEY = "pf_live_your_key_here"

@dataclass
class EscrowRecord:
    escrow_id: str
    worker_wallet: str
    estimated_cost_usdc: float
    prompt_tokens: int
    max_completion_tokens: int
    model: str

class PurpleFleatEscrowClient:
    def __init__(self, api_key: str, price_per_token: float = 0.000002):
        self.api_key = api_key
        self.price_per_token = price_per_token  # USDC per token
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

    async def create_escrow(
        self,
        worker_wallet: str,
        prompt_tokens: int,
        max_completion_tokens: int,
        model: str,
    ) -> EscrowRecord:
        total_tokens = prompt_tokens + max_completion_tokens
        estimated_cost = total_tokens * self.price_per_token

        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{ESCROW_BASE}/escrow/create",
                headers=self.headers,
                json={
                    "recipient_wallet": worker_wallet,
                    "amount_usdc": estimated_cost,
                    "metadata": {
                        "model": model,
                        "prompt_tokens": prompt_tokens,
                        "max_completion_tokens": max_completion_tokens,
                    },
                },
            )
            resp.raise_for_status()
            data = resp.json()

        return EscrowRecord(
            escrow_id=data["escrow_id"],
            worker_wallet=worker_wallet,
            estimated_cost_usdc=estimated_cost,
            prompt_tokens=prompt_tokens,
            max_completion_tokens=max_completion_tokens,
            model=model,
        )

    async def release_escrow(
        self,
        escrow_id: str,
        actual_completion_tokens: int,
    ) -> dict:
        # Calculates refund automatically if actual < estimated
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{ESCROW_BASE}/escrow/release",
                headers=self.headers,
                json={
                    "escrow_id": escrow_id,
                    "actual_completion_tokens": actual_completion_tokens,
                    "price_per_token": self.price_per_token,
                },
            )
            resp.raise_for_status()
            return resp.json()

    async def cancel_escrow(self, escrow_id: str) -> dict:
        # Called on error — refunds full amount to orchestrator
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{ESCROW_BASE}/escrow/cancel",
                headers=self.headers,
                json={"escrow_id": escrow_id},
            )
            resp.raise_for_status()
            return resp.json()

import asyncio
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
from vllm.utils import random_uuid
from escrow_client import PurpleFleatEscrowClient

# ── Engine init ──────────────────────────────────────────────
engine_args = AsyncEngineArgs(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    tensor_parallel_size=1,
    max_model_len=8192,
)
engine = AsyncLLMEngine.from_engine_args(engine_args)
escrow = PurpleFleatEscrowClient(api_key="pf_live_your_key_here")

# ── Per-request cost tracker ─────────────────────────────────
async def generate_with_billing(
    prompt: str,
    worker_wallet: str,
    max_tokens: int = 512,
    temperature: float = 0.7,
) -> str:
    request_id = random_uuid()
    sampling_params = SamplingParams(
        temperature=temperature,
        max_tokens=max_tokens,
    )

    # Estimate prompt token count (tokenizer call)
    tokenizer = await engine.get_tokenizer()
    prompt_tokens = len(tokenizer.encode(prompt))

    # Lock escrow before work starts
    record = await escrow.create_escrow(
        worker_wallet=worker_wallet,
        prompt_tokens=prompt_tokens,
        max_completion_tokens=max_tokens,
        model="llama-3-8b-instruct",
    )
    print(f"[billing] Escrow {record.escrow_id} locked: ${record.estimated_cost_usdc:.6f} USDC")

    full_output = ""
    completion_tokens = 0

    try:
        # Stream tokens from vLLM
        async for output in engine.generate(prompt, sampling_params, request_id):
            if output.outputs:
                full_output = output.outputs[0].text
                completion_tokens = len(output.outputs[0].token_ids)

        # Release with actual token count — overpayment auto-refunded
        result = await escrow.release_escrow(
            escrow_id=record.escrow_id,
            actual_completion_tokens=completion_tokens,
        )
        print(f"[billing] Released: {completion_tokens} tokens, paid ${result['amount_paid']:.6f} USDC")
        return full_output

    except Exception as e:
        # Cancel escrow on error — full refund
        await escrow.cancel_escrow(record.escrow_id)
        print(f"[billing] Escrow {record.escrow_id} cancelled: {e}")
        raise

# ── Worker settlement loop ───────────────────────────────────
async def batch_inference(requests: list[dict]) -> list[str]:
    """Run multiple inference requests, each with independent escrow."""
    tasks = [
        generate_with_billing(
            prompt=req["prompt"],
            worker_wallet=req["worker_wallet"],
            max_tokens=req.get("max_tokens", 512),
        )
        for req in requests
    ]
    return await asyncio.gather(*tasks, return_exceptions=True)

import time

async def generate_with_performance_bonus(
    prompt: str,
    worker_wallet: str,
    max_tokens: int = 512,
) -> dict:
    """Track throughput and submit quality-gated release."""
    tokenizer = await engine.get_tokenizer()
    prompt_tokens = len(tokenizer.encode(prompt))
    request_id = random_uuid()

    record = await escrow.create_escrow(
        worker_wallet=worker_wallet,
        prompt_tokens=prompt_tokens,
        max_completion_tokens=max_tokens,
        model="llama-3-8b-instruct",
    )

    start_time = time.monotonic()
    first_token_time: Optional[float] = None
    full_output = ""
    completion_tokens = 0

    async for output in engine.generate(
        prompt, SamplingParams(max_tokens=max_tokens), request_id
    ):
        if output.outputs:
            if first_token_time is None:
                first_token_time = time.monotonic()
            full_output = output.outputs[0].text
            completion_tokens = len(output.outputs[0].token_ids)

    elapsed = time.monotonic() - (first_token_time or start_time)
    tps = completion_tokens / elapsed if elapsed > 0 else 0

    # Assess quality with a lightweight judge (optional)
    quality_score = await assess_quality(prompt, full_output)

    result = await escrow.release_escrow(
        escrow_id=record.escrow_id,
        actual_completion_tokens=completion_tokens,
    )

    return {
        "output": full_output,
        "tokens_per_sec": round(tps, 1),
        "quality_score": quality_score,
        "amount_paid_usdc": result["amount_paid"],
        "escrow_id": record.escrow_id,
    }