add files

server.py

@@ -0,0 +1,132 @@
+import os
+import uuid
+import time
+from pathlib import Path
+
+# Suppress verbose logs
+os.environ["GRPC_VERBOSITY"] = "ERROR"
+os.environ["GLOG_minloglevel"] = "3"
+
+import litert_lm
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import HTMLResponse
+from pydantic import BaseModel
+from contextlib import asynccontextmanager
+
+# ── Config ────────────────────────────────────────────────────────────────────
+
+MODEL_PATH   = "gemma-4-E2B-it.litertlm"
+TEMPLATE_DIR = Path(__file__).parent / "templates"
+
+# ── Models ───────────────────────────────────────────────────────────────────
+
+class PromptRequest(BaseModel):
+    prompt: str
+
+# ── State ────────────────────────────────────────────────────────────────────
+
+ml_models = {}
+sessions: dict = {}  # session_id -> conversation object
+
+# ── Lifespan ─────────────────────────────────────────────────────────────────
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.CPU)
+    ml_models["engine"] = engine
+    yield
+    sessions.clear()
+    del ml_models["engine"]
+
+# ── App ───────────────────────────────────────────────────────────────────────
+
+app = FastAPI(title="LiteRT-LM API", lifespan=lifespan)
+
+# ── REST: stateless single-turn ───────────────────────────────────────────────
+
+@app.post("/generate")
+async def generate_text(request: PromptRequest):
+    """Single-turn generation. No memory between calls."""
+    engine = ml_models.get("engine")
+    if not engine:
+        raise HTTPException(status_code=503, detail="Model engine not initialized")
+    try:
+        conversation = engine.create_conversation()
+        t0 = time.perf_counter()
+        result = conversation.send_message(request.prompt)
+        elapsed = time.perf_counter() - t0
+        text = result["content"][0]["text"]
+        num_tokens = len(engine.tokenize(text))
+        tps = round(num_tokens / elapsed, 2) if elapsed > 0 else 0
+        return {
+            "response": text,
+            "tokens": num_tokens,
+            "elapsed_s": round(elapsed, 2),
+            "tokens_per_sec": tps,
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# ── REST: multi-turn chat sessions ────────────────────────────────────────────
+
+@app.post("/chat/new")
+async def new_session():
+    """Create a new chat session. Returns session_id."""
+    engine = ml_models.get("engine")
+    if not engine:
+        raise HTTPException(status_code=503, detail="Model engine not initialized")
+    session_id = str(uuid.uuid4())
+    sessions[session_id] = engine.create_conversation()
+    return {"session_id": session_id}
+
+@app.post("/chat/{session_id}")
+async def chat(session_id: str, request: PromptRequest):
+    """Send a message in an existing session (retains conversation history)."""
+    if session_id not in sessions:
+        raise HTTPException(
+            status_code=404,
+            detail="Session not found. Create one via POST /chat/new",
+        )
+    try:
+        engine = ml_models.get("engine")
+        t0 = time.perf_counter()
+        result = sessions[session_id].send_message(request.prompt)
+        elapsed = time.perf_counter() - t0
+        text = result["content"][0]["text"]
+        num_tokens = len(engine.tokenize(text)) if engine else 0
+        tps = round(num_tokens / elapsed, 2) if elapsed > 0 else 0
+        return {
+            "session_id": session_id,
+            "response": text,
+            "tokens": num_tokens,
+            "elapsed_s": round(elapsed, 2),
+            "tokens_per_sec": tps,
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.delete("/chat/{session_id}")
+async def clear_session(session_id: str):
+    """Delete a session and free its memory."""
+    if session_id not in sessions:
+        raise HTTPException(status_code=404, detail="Session not found")
+    del sessions[session_id]
+    return {"status": "cleared", "session_id": session_id}
+
+@app.get("/chat/sessions/list")
+async def list_sessions():
+    """List all active session IDs."""
+    return {"sessions": list(sessions.keys()), "count": len(sessions)}
+
+# ── WebUI ─────────────────────────────────────────────────────────────────────
+
+@app.get("/", response_class=HTMLResponse)
+async def web_ui():
+    html = (TEMPLATE_DIR / "index.html").read_text(encoding="utf-8")
+    return HTMLResponse(content=html)
+
+# ── Run ───────────────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)