Update puter_server.py

puter_server.py

@@ -1,134 +1,390 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse, StreamingResponse
-import requests
+#!/usr/bin/env python3
+"""
+Puter.com Reverse OpenAI-Compatible API Server
+
+Accepts OpenAI Chat Completions requests and forwards them to:
+  POST https://api.puter.com/drivers/call
+
+with payload:
+  {
+    "interface": "puter-chat-completion",
+    "driver": "xai",
+    "test_mode": false,
+    "method": "complete",
+    "args": {
+      "messages": [{"content": "..."}],
+      "model": "x-ai/grok-4.1-fast",
+      "stream": true
+    }
+  }
+"""
 import json
+import time
 import uuid
-from typing import Dict, Any, List
+import logging
+from typing import Any, Dict, List, Optional, Union, AsyncGenerator
 
-from config import (
-    PUTER_API_URL,
-    PUTER_HEADERS,
-    PUTER_AUTH_BEARER,
-    MODEL_MAPPING,
-)
+import requests
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel, Field
+
+try:
+    from .config import (
+        PUTER_HEADERS,
+        PUTER_AUTH_BEARER,
+        SERVER_CONFIG,
+        MODEL_MAPPING,
+    )
+except ImportError:
+    from config import (
+        PUTER_HEADERS,
+        PUTER_AUTH_BEARER,
+        SERVER_CONFIG,
+        MODEL_MAPPING,
+    )
 
-app = FastAPI()
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
 
+PUTER_URL = "https://api.puter.com/drivers/call"
+REQUEST_TIMEOUT = 120
 
-# ---------- Helpers ----------
 
-def build_puter_payload(openai_req: Dict[str, Any]) -> Dict[str, Any]:
-    messages = [{"content": m["content"]} for m in openai_req["messages"]]
+# ===== OpenAI-compatible models =====
+class OpenAIMessage(BaseModel):
+    role: Optional[str] = Field(default=None, description="Role")
+    content: Optional[Union[str, List[Dict[str, Any]]]] = None
+    name: Optional[str] = None
+    function_call: Optional[Dict[str, Any]] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    tool_call_id: Optional[str] = None
 
-    model_key = openai_req.get("model", "default")
-    mapping = MODEL_MAPPING.get(model_key, MODEL_MAPPING["default"])
+    def get_text(self) -> str:
+        if isinstance(self.content, str):
+            return self.content
+        if isinstance(self.content, list):
+            parts: List[str] = []
+            for item in self.content:
+                if isinstance(item, dict) and item.get("type") == "text":
+                    parts.append(item.get("text", ""))
+            return "".join(parts)
+        return str(self.content) if self.content is not None else ""
 
-    args: Dict[str, Any] = {
-        "model": mapping["puter_model"],
-        "messages": messages,
-        "stream": True,
-    }
+    class Config:
+        extra = "allow"
+
+
+class OpenAIFunction(BaseModel):
+    name: str
+    description: Optional[str] = None
+    parameters: Optional[Dict[str, Any]] = None
+
+    class Config:
+        extra = "allow"
+
+
+class OpenAITool(BaseModel):
+    type: str = Field(default="function")
+    function: Optional[OpenAIFunction] = None
+
+    class Config:
+        extra = "allow"
 
-    # Forward OpenAI-compatible params
-    for key in [
-        "temperature",
-        "max_tokens",
-        "top_p",
-        "stop",
-        "presence_penalty",
-        "frequency_penalty",
-        "user",
-    ]:
-        if key in openai_req:
-            args[key] = openai_req[key]
-
-    return {
+
+class OpenAIChatRequest(BaseModel):
+    model: str
+    messages: List[OpenAIMessage]
+    max_tokens: Optional[int] = None
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    n: Optional[int] = 1
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
+    presence_penalty: Optional[float] = None
+    frequency_penalty: Optional[float] = None
+    logit_bias: Optional[Dict[str, float]] = None
+    user: Optional[str] = None
+    tools: Optional[List[OpenAITool]] = None
+    tool_choice: Optional[Union[str, Dict[str, Any]]] = None
+    functions: Optional[List[OpenAIFunction]] = None
+    function_call: Optional[Union[str, Dict[str, Any]]] = None
+
+    class Config:
+        extra = "allow"
+
+
+class OpenAIChoice(BaseModel):
+    index: int = 0
+    message: Dict[str, Any]
+    finish_reason: Optional[str] = None
+
+
+class OpenAIChatResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[OpenAIChoice]
+    usage: Optional[Dict[str, int]] = None
+
+
+class OpenAIStreamChoice(BaseModel):
+    index: int = 0
+    delta: Dict[str, Any]
+    finish_reason: Optional[str] = None
+
+
+class OpenAIStreamChunk(BaseModel):
+    id: str
+    object: str = "chat.completion.chunk"
+    created: int
+    model: str
+    choices: List[OpenAIStreamChoice]
+
+
+def _build_puter_payload(openai_req: OpenAIChatRequest) -> Dict[str, Any]:
+    # Map OpenAI messages to Puter format: only 'content' is used
+    mapped_messages: List[Dict[str, str]] = []
+    for m in openai_req.messages:
+        txt = m.get_text()
+        mapped_messages.append({"content": txt})
+
+    # Model mapping: map OpenAI model key -> (driver, puter_model)
+    mapping = MODEL_MAPPING.get(openai_req.model) or MODEL_MAPPING.get("default")
+    driver = mapping["driver"]
+    puter_model = mapping["puter_model"]
+
+    payload: Dict[str, Any] = {
         "interface": "puter-chat-completion",
-        "driver": mapping["driver"],
-        "method": "complete",
+        "driver": driver,
         "test_mode": False,
-        "args": args,
+        "method": "complete",
+        "args": {
+            "messages": mapped_messages,
+            "model": puter_model,
+            "stream": True,  # always request streaming upstream; we aggregate if needed
+        },
     }
+    return payload
 
 
-def stream_to_openai(resp):
-    completion_id = f"chatcmpl-{uuid.uuid4().hex}"
+def _headers_with_auth() -> Dict[str, str]:
+    h = dict(PUTER_HEADERS)
+    h["authorization"] = f"Bearer {PuterAuth.token}"
+    return h
 
-    for line in resp.iter_lines():
-        if not line:
-            continue
 
-        data = json.loads(line.decode("utf-8"))
+class PuterAuth:
+    token: str = PUTER_AUTH_BEARER
 
-        token = data.get("text")
-        if not token:
-            continue
 
-        chunk = {
-            "id": completion_id,
-            "object": "chat.completion.chunk",
-            "choices": [
-                {
-                    "delta": {"content": token},
-                    "index": 0,
-                    "finish_reason": None,
-                }
-            ],
-        }
-        yield f"data: {json.dumps(chunk)}\n\n"
+async def _stream_openai_chunks(openai_req: OpenAIChatRequest, request_id: str) -> AsyncGenerator[str, None]:
+    headers = _headers_with_auth()
+    payload = _build_puter_payload(openai_req)
 
-    yield "data: [DONE]\n\n"
+    with requests.Session() as sess:
+        try:
+            resp = sess.post(
+                PUTER_URL,
+                headers=headers,
+                json=payload,
+                stream=True,
+                timeout=REQUEST_TIMEOUT,
+            )
+        except requests.RequestException as e:
+            raise HTTPException(status_code=502, detail=f"Upstream connection error: {e}")
 
+        if resp.status_code != 200:
+            detail = resp.text[:500]
+            raise HTTPException(status_code=502, detail=f"Upstream error {resp.status_code}: {detail}")
 
-# ---------- Routes ----------
+        created = int(time.time())
 
-@app.post("/v1/chat/completions")
-async def chat_completions(req: Request):
-    openai_req = await req.json()
-    stream = openai_req.get("stream", False)
+        # Initial role chunk
+        initial = OpenAIStreamChunk(
+            id=request_id,
+            created=created,
+            model=openai_req.model,
+            choices=[OpenAIStreamChoice(index=0, delta={"role": "assistant"}, finish_reason=None)],
+        )
+        yield f"data: {initial.model_dump_json()}\n\n"
 
-    payload = build_puter_payload(openai_req)
+        # Stream content
+        for raw in resp.iter_lines():
+            if not raw:
+                continue
+            try:
+                line = raw.decode("utf-8", errors="ignore")
+            except Exception:
+                continue
 
-    headers = {
-        **PUTER_HEADERS,
-        "authorization": f"Bearer {PUTER_AUTH_BEARER}",
-    }
+            text_piece: Optional[str] = None
+            # Many APIs stream JSON lines; try to parse
+            try:
+                obj = json.loads(line)
+                # Common keys
+                for k in ("delta", "text", "content", "output"):
+                    if isinstance(obj.get(k), str) and obj.get(k):
+                        text_piece = obj.get(k)
+                        break
+            except Exception:
+                # Fallback to raw text
+                if line and line != "[DONE]":
+                    text_piece = line
+
+            if not text_piece:
+                continue
+
+            chunk = OpenAIStreamChunk(
+                id=request_id,
+                created=created,
+                model=openai_req.model,
+                choices=[OpenAIStreamChoice(index=0, delta={"content": text_piece}, finish_reason=None)],
+            )
+            yield f"data: {chunk.model_dump_json()}\n\n"
+
+        final = OpenAIStreamChunk(
+            id=request_id,
+            created=created,
+            model=openai_req.model,
+            choices=[OpenAIStreamChoice(index=0, delta={}, finish_reason="stop")],
+        )
+        yield f"data: {final.model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
 
-    resp = requests.post(
-        PUTER_API_URL,
-        headers=headers,
-        json=payload,
-        stream=True,
-        timeout=300,
-    )
 
-    if stream:
+def _complete_non_streaming(openai_req: OpenAIChatRequest) -> str:
+    headers = _headers_with_auth()
+    payload = _build_puter_payload(openai_req)
+    payload["args"]["stream"] = True
+
+    with requests.Session() as sess:
+        try:
+            resp = sess.post(
+                PUTER_URL,
+                headers=headers,
+                json=payload,
+                stream=True,
+                timeout=REQUEST_TIMEOUT,
+            )
+        except requests.RequestException as e:
+            raise HTTPException(status_code=502, detail=f"Upstream connection error: {e}")
+
+        if resp.status_code != 200:
+            detail = resp.text[:500]
+            raise HTTPException(status_code=502, detail=f"Upstream error {resp.status_code}: {detail}")
+
+        parts: List[str] = []
+        for raw in resp.iter_lines():
+            if not raw:
+                continue
+            try:
+                line = raw.decode("utf-8", errors="ignore")
+            except Exception:
+                continue
+            try:
+                obj = json.loads(line)
+                for k in ("delta", "text", "content", "output"):
+                    if isinstance(obj.get(k), str) and obj.get(k):
+                        parts.append(obj.get(k))
+                        break
+            except Exception:
+                if line and line != "[DONE]":
+                    parts.append(line)
+        return "".join(parts)
+
+
+# ===== FastAPI app =====
+app = FastAPI(
+    title="Puter Reverse OpenAI API",
+    version="1.0.0",
+    description="OpenAI-compatible API proxying to api.puter.com"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.get("/")
+async def root():
+    return {"message": "Puter Reverse OpenAI API", "status": "running", "version": "1.0.0"}
+
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy", "timestamp": int(time.time())}
+
+
+@app.get("/v1/models")
+async def models():
+    created = int(time.time())
+    data = []
+    for key in [k for k in MODEL_MAPPING.keys() if k != "default"]:
+        data.append({"id": key, "object": "model", "created": created, "owned_by": "puter"})
+    if not data:
+        data.append({"id": "x-ai/grok-4.1-fast", "object": "model", "created": created, "owned_by": "puter"})
+    return {"object": "list", "data": data}
+
+
+@app.post("/v1/chat/completions")
+async def chat(request: OpenAIChatRequest):
+    req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+    logger.info(f"[{req_id}] model={request.model}, stream={bool(request.stream)}")
+
+    if bool(request.stream):
         return StreamingResponse(
-            stream_to_openai(resp),
+            _stream_openai_chunks(request, req_id),
             media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+                "Access-Control-Allow-Origin": "*",
+                "Access-Control-Allow-Headers": "*",
+            },
         )
 
-    # Non-stream response aggregation
-    full_text = ""
-    for line in resp.iter_lines():
-        if not line:
-            continue
-        data = json.loads(line.decode("utf-8"))
-        full_text += data.get("text", "")
-
-    return JSONResponse(
-        {
-            "id": f"chatcmpl-{uuid.uuid4().hex}",
-            "object": "chat.completion",
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": full_text,
-                    },
-                    "finish_reason": "stop",
-                }
-            ],
-        }
+    content = _complete_non_streaming(request)
+    created = int(time.time())
+    response = OpenAIChatResponse(
+        id=req_id,
+        created=created,
+        model=request.model,
+        choices=[OpenAIChoice(index=0, message={"role": "assistant", "content": content}, finish_reason="stop")],
+        usage={
+            "prompt_tokens": len(" ".join([m.get_text() for m in request.messages]).split()),
+            "completion_tokens": len(content.split()),
+            "total_tokens": len(" ".join([m.get_text() for m in request.messages]).split()) + len(content.split()),
+        },
     )
+    return response
+
+
+@app.post("/v1/chat/completions/raw")
+async def raw(req: Request):
+    body = await req.body()
+    try:
+        obj = json.loads(body)
+        _ = OpenAIChatRequest(**obj)
+        return {"valid": True}
+    except Exception as e:
+        return JSONResponse(status_code=422, content={"valid": False, "error": str(e)})
+
+
+if __name__ == "__main__":
+    try:
+        import uvicorn
+        host = SERVER_CONFIG.get("host", "0.0.0.0")
+        port = int(SERVER_CONFIG.get("port", 8781))
+        logger.info(f"Starting Puter Reverse API on {host}:{port}")
+        uvicorn.run(app, host=host, port=port, log_level="info")
+    except Exception as e:
+        logger.error(f"Failed to start server: {e}")
+