Create app/main.py

app/main.py

@@ -0,0 +1,212 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, Dict, List, Optional
+
+import httpx
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, Field
+from pypdf import PdfReader
+
+PERPLEXITY_URL = "https://api.perplexity.ai/chat/completions"
+
+# Models listed in Perplexity docs for Sonar family.
+SONAR_MODELS = [
+    "sonar",
+    "sonar-pro",
+    "sonar-deep-research",
+    "sonar-reasoning",
+    "sonar-reasoning-pro",
+]
+
+app = FastAPI(title="Perplexity HF Playground", version="1.0.0")
+
+# Allow any origin because UI might be opened from file:// or different domain
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_headers=["*"],
+    allow_methods=["*"],
+)
+
+static_dir = os.path.join(os.path.dirname(__file__), "static")
+app.mount("/static", StaticFiles(directory=static_dir), name="static")
+
+
+class ChatRequest(BaseModel):
+    api_key: str = Field(..., description="Perplexity API key (Bearer token), provided by user in UI.")
+    model: str = Field(default="sonar")
+    messages: List[Dict[str, Any]] = Field(default_factory=list)
+
+    # Common knobs (all optional; forwarded as-is)
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    max_tokens: Optional[int] = None
+    stream: bool = False
+
+    # Perplexity-specific (optional)
+    search_mode: Optional[str] = None  # "web" | "academic" | "sec"
+    reasoning_effort: Optional[str] = None  # ONLY for sonar-deep-research: "low" | "medium" | "high"
+    search_domain_filter: Optional[List[str]] = None
+    return_images: Optional[bool] = None
+    return_related_questions: Optional[bool] = None
+    search_recency_filter: Optional[str] = None
+    disable_search: Optional[bool] = None
+    web_search_options: Optional[Dict[str, Any]] = None
+
+
+def _auth_headers(api_key: str) -> Dict[str, str]:
+    api_key = (api_key or "").strip()
+    if not api_key:
+        raise HTTPException(status_code=400, detail="Missing api_key")
+    return {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+    }
+
+
+@app.get("/")
+def root():
+    return FileResponse(os.path.join(static_dir, "index.html"))
+
+
+@app.get("/api/health")
+def health():
+    return {"ok": True}
+
+
+@app.get("/api/models")
+def models():
+    return {"ok": True, "models": SONAR_MODELS}
+
+
+async def _perplexity_post_json(payload: Dict[str, Any], api_key: str) -> Dict[str, Any]:
+    async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client:
+        resp = await client.post(PERPLEXITY_URL, headers=_auth_headers(api_key), json=payload)
+        if resp.status_code >= 400:
+            # Avoid leaking key; return error details safely
+            raise HTTPException(status_code=resp.status_code, detail=resp.text)
+        return resp.json()
+
+
+@app.post("/api/chat")
+async def chat(req: ChatRequest):
+    if req.model not in SONAR_MODELS:
+        # Keep strict to prevent surprising failures. You can relax later.
+        raise HTTPException(status_code=400, detail=f"Unsupported model: {req.model}")
+
+    if not req.messages:
+        raise HTTPException(status_code=400, detail="messages[] is required")
+
+    # Build payload by forwarding only non-null fields
+    payload: Dict[str, Any] = {
+        "model": req.model,
+        "messages": req.messages,
+        "stream": req.stream,
+    }
+
+    optional_fields = [
+        "temperature",
+        "top_p",
+        "max_tokens",
+        "search_mode",
+        "reasoning_effort",
+        "search_domain_filter",
+        "return_images",
+        "return_related_questions",
+        "search_recency_filter",
+        "disable_search",
+        "web_search_options",
+    ]
+    for f in optional_fields:
+        v = getattr(req, f)
+        if v is not None:
+            payload[f] = v
+
+    if not req.stream:
+        data = await _perplexity_post_json(payload, req.api_key)
+        # Normalized output for frontend convenience
+        content = (
+            (data.get("choices") or [{}])[0].get("message", {}).get("content")
+            if isinstance(data, dict)
+            else None
+        )
+        citations = data.get("citations") if isinstance(data, dict) else None
+        return JSONResponse({"ok": True, "content": content, "citations": citations, "raw": data})
+
+    # STREAMING (SSE passthrough)
+    async def event_generator():
+        headers = _auth_headers(req.api_key)
+        async with httpx.AsyncClient(timeout=None) as client:
+            async with client.stream("POST", PERPLEXITY_URL, headers=headers, json=payload) as r:
+                if r.status_code >= 400:
+                    body = await r.aread()
+                    # SSE error frame
+                    yield f"event: error\ndata: {body.decode('utf-8', errors='ignore')}\n\n"
+                    return
+
+                async for line in r.aiter_lines():
+                    if not line:
+                        continue
+                    # Pass through raw SSE-ish lines. Frontend will parse best-effort.
+                    yield line + "\n"
+
+    return StreamingResponse(event_generator(), media_type="text/event-stream")
+
+
+def _extract_text_from_upload(file: UploadFile, max_chars: int = 120_000) -> str:
+    filename = (file.filename or "").lower()
+    if filename.endswith(".txt") or (file.content_type or "").startswith("text/"):
+        raw = file.file.read()
+        text = raw.decode("utf-8", errors="ignore")
+        return text[:max_chars]
+
+    if filename.endswith(".pdf") or file.content_type == "application/pdf":
+        reader = PdfReader(file.file)
+        chunks: List[str] = []
+        for page in reader.pages:
+            t = page.extract_text() or ""
+            if t.strip():
+                chunks.append(t)
+            if sum(len(c) for c in chunks) > max_chars:
+                break
+        return ("\n\n".join(chunks))[:max_chars]
+
+    raise HTTPException(status_code=400, detail="Only .txt and .pdf supported for upload in this build.")
+
+
+@app.post("/api/chat-upload")
+async def chat_upload(
+    api_key: str = Form(...),
+    model: str = Form("sonar"),
+    prompt: str = Form(...),
+    system_prompt: str = Form("You are a helpful assistant."),
+    stream: bool = Form(False),
+    file: Optional[UploadFile] = File(None),
+):
+    if model not in SONAR_MODELS:
+        raise HTTPException(status_code=400, detail=f"Unsupported model: {model}")
+
+    file_text = ""
+    if file is not None:
+        file_text = _extract_text_from_upload(file)
+
+    user_content = prompt
+    if file_text.strip():
+        user_content += "\n\n---\nAttached file text (extracted):\n" + file_text
+
+    req = ChatRequest(
+        api_key=api_key,
+        model=model,
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_content},
+        ],
+        stream=stream,
+    )
+    return await chat(req)