Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """ | |
| STARBOY PRIME Local OpenAI-Compatible Proxy Sidecar. | |
| Exposes standard OpenAI endpoint on localhost and routes to ZeroGPU Space. | |
| """ | |
| import os | |
| import sys | |
| import json | |
| import time | |
| import argparse | |
| import uvicorn | |
| from fastapi import FastAPI, Request, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse, StreamingResponse | |
| from pydantic import BaseModel | |
| from typing import List, Optional | |
| from safe_client import SafeZeroGPUClient, ZeroGPUQuotaError, ZeroGPURateLimitError | |
| UPSTREAM_URL = os.environ.get("STARBOY_UPSTREAM_URL", "https://starboy001-antigravity-cloud.hf.space") | |
| UPSTREAM_KEY = os.environ.get("STARBOY_API_KEY", "starboy2026") | |
| DEFAULT_MODEL = os.environ.get("MODEL_ID", "BlossomsAI/Qwen2.5-Coder-7B-Instruct-Uncensored") | |
| client = SafeZeroGPUClient(base_url=UPSTREAM_URL, api_key=UPSTREAM_KEY, model=DEFAULT_MODEL) | |
| proxy_app = FastAPI( | |
| title="STARBOY PRIME OpenAI Proxy", | |
| description="Local OpenAI-compatible sidecar for ZeroGPU Flash LLM", | |
| version="1.0.0" | |
| ) | |
| proxy_app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| class ChatMessage(BaseModel): | |
| role: str | |
| content: str | |
| class ChatCompletionRequest(BaseModel): | |
| model: Optional[str] = None | |
| messages: List[ChatMessage] | |
| max_tokens: Optional[int] = 512 | |
| max_new_tokens: Optional[int] = None | |
| temperature: Optional[float] = 0.7 | |
| top_p: Optional[float] = 0.9 | |
| stream: Optional[bool] = False | |
| async def health(): | |
| upstream_health = client.health_check() | |
| return { | |
| "status": "healthy", | |
| "proxy": "starboy-openai-proxy", | |
| "upstream_url": UPSTREAM_URL, | |
| "upstream_status": upstream_health.get("status", "unknown"), | |
| "timestamp": time.time() | |
| } | |
| async def list_models(): | |
| return { | |
| "object": "list", | |
| "data": [ | |
| { | |
| "id": DEFAULT_MODEL, | |
| "object": "model", | |
| "created": 1700000000, | |
| "owned_by": "starboy-prime" | |
| } | |
| ] | |
| } | |
| async def chat_completions(req: ChatCompletionRequest): | |
| max_tokens = req.max_new_tokens if req.max_new_tokens is not None else (req.max_tokens or 512) | |
| temp = req.temperature if req.temperature is not None else 0.7 | |
| top_p = req.top_p if req.top_p is not None else 0.9 | |
| messages_payload = [{"role": m.role, "content": m.content} for m in req.messages] | |
| model_name = req.model or DEFAULT_MODEL | |
| if req.stream: | |
| def stream_generator(): | |
| created = int(time.time()) | |
| req_id = f"chatcmpl-{os.urandom(6).hex()}" | |
| try: | |
| for token in client.stream_chat(messages_payload, max_tokens=max_tokens, temperature=temp, top_p=top_p): | |
| chunk_obj = { | |
| "id": req_id, | |
| "object": "chat.completion.chunk", | |
| "created": created, | |
| "model": model_name, | |
| "choices": [{"index": 0, "delta": {"content": token}, "finish_reason": None}] | |
| } | |
| yield f"data: {json.dumps(chunk_obj)}\n\n" | |
| final_chunk = { | |
| "id": req_id, | |
| "object": "chat.completion.chunk", | |
| "created": created, | |
| "model": model_name, | |
| "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}] | |
| } | |
| yield f"data: {json.dumps(final_chunk)}\n\n" | |
| yield "data: [DONE]\n\n" | |
| except ZeroGPUQuotaError as q_err: | |
| err_chunk = {"error": {"message": str(q_err), "type": "quota_exceeded", "code": 402}} | |
| yield f"data: {json.dumps(err_chunk)}\n\n" | |
| except Exception as e: | |
| err_chunk = {"error": {"message": str(e), "type": "upstream_error", "code": 500}} | |
| yield f"data: {json.dumps(err_chunk)}\n\n" | |
| return StreamingResponse( | |
| stream_generator(), | |
| media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"} | |
| ) | |
| try: | |
| content = client.chat_completion(messages_payload, max_tokens=max_tokens, temperature=temp, top_p=top_p) | |
| created = int(time.time()) | |
| req_id = f"chatcmpl-{os.urandom(6).hex()}" | |
| return JSONResponse({ | |
| "id": req_id, | |
| "object": "chat.completion", | |
| "created": created, | |
| "model": model_name, | |
| "choices": [ | |
| { | |
| "index": 0, | |
| "message": {"role": "assistant", "content": content}, | |
| "finish_reason": "stop" | |
| } | |
| ], | |
| "usage": { | |
| "prompt_tokens": len(str(messages_payload)) // 4, | |
| "completion_tokens": len(content) // 4, | |
| "total_tokens": (len(str(messages_payload)) + len(content)) // 4 | |
| } | |
| }) | |
| except ZeroGPUQuotaError as q_err: | |
| raise HTTPException(status_code=402, detail=str(q_err)) | |
| except ZeroGPURateLimitError as r_err: | |
| raise HTTPException(status_code=429, detail=str(r_err)) | |
| except Exception as exc: | |
| raise HTTPException(status_code=500, detail=str(exc)) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="STARBOY PRIME OpenAI Proxy Sidecar") | |
| parser.add_argument("--host", default="127.0.0.1", help="Host to bind to") | |
| parser.add_argument("--port", type=int, default=8080, help="Port to bind to") | |
| args = parser.parse_args() | |
| uvicorn.run(proxy_app, host=args.host, port=args.port) | |