splitbit_llm/server/server.py

"""FastAPI Server for SplitBit LLM.

Endpoints:
- GET  /          — Chat UI
- GET  /jarvis    — Jarvis voice UI
- POST /v1/chat   — Chat completion
- POST /v1/chat/stream — SSE streaming chat
- POST /v1/voice/stream — SSE streaming optimized for TTS (sentence-by-sentence)
- GET  /v1/stats  — Model stats
- GET  /v1/health — Health check
- No authentication required (100% local)
- CORS enabled
"""

from __future__ import annotations

import json
import logging
from typing import Any

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from pydantic import BaseModel

from ..harness.harness import SplitBitHarness
from .ui import CHAT_UI_HTML, JARVIS_UI_HTML

logger = logging.getLogger(__name__)


class ChatRequest(BaseModel):
    message: str
    channel: str = "web"
    session_id: str = ""
    max_tokens: int | None = None
    temperature: float | None = None


def create_app(harness: SplitBitHarness | None = None) -> FastAPI:
    """Create and configure the FastAPI app."""
    app = FastAPI(title="SplitBit LLM", version="0.1.0")

    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    # Initialize harness if not provided
    if harness is None:
        harness = SplitBitHarness()

    @app.get("/", response_class=HTMLResponse)
    async def chat_ui():
        return CHAT_UI_HTML

    @app.get("/jarvis", response_class=HTMLResponse)
    async def jarvis_ui():
        return JARVIS_UI_HTML

    @app.get("/v1/health")
    async def health():
        return {"status": "ok", "model": "splitbit-llm", "version": "0.1.0"}

    @app.post("/v1/chat")
    async def chat(req: ChatRequest):
        result = harness.chat(
            message=req.message,
            channel=req.channel,
            session_id=req.session_id,
            max_tokens=req.max_tokens,
            temperature=req.temperature,
        )
        return JSONResponse(result)

    @app.post("/v1/chat/stream")
    async def chat_stream(req: ChatRequest):
        def generate():
            for chunk in harness.chat_stream(
                message=req.message,
                channel=req.channel,
                session_id=req.session_id,
            ):
                yield f"data: {json.dumps({'chunk': chunk})}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")

    @app.post("/v1/voice/stream")
    async def voice_stream(req: ChatRequest):
        def generate():
            for sentence in harness.chat_stream_sentences(
                message=req.message,
                channel="voice",
                session_id=req.session_id,
            ):
                yield f"data: {json.dumps({'sentence': sentence})}\n\n"
            yield "data: [DONE]\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")

    @app.get("/v1/stats")
    async def stats():
        return JSONResponse(harness.get_stats())

    @app.get("/v1/tools")
    async def tools():
        return {"tools": harness.tools.list_tools()}

    @app.post("/v1/image/generate")
    async def generate_image(request: Request):
        body = await request.json()
        prompt = body.get("prompt", "")
        width = body.get("width", 256)
        height = body.get("height", 256)
        result = harness.generate_image(prompt, width=width, height=height)
        return JSONResponse(result)

    @app.get("/v1/image/generate")
    async def generate_image_get(prompt: str, width: int = 256, height: int = 256):
        result = harness.generate_image(prompt, width=width, height=height)
        return JSONResponse(result)

    @app.post("/v1/connectors/register")
    async def register_connector(request: Request):
        body = await request.json()
        harness.register_connector(
            name=body.get("name", ""),
            base_url=body.get("base_url", ""),
            api_key=body.get("api_key", ""),
            auth_type=body.get("auth_type", "api_key"),
        )
        return {"success": True, "name": body.get("name")}

    @app.post("/v1/connectors/call")
    async def call_connector(request: Request):
        body = await request.json()
        result = harness.api_call(
            name=body.get("name", ""),
            method=body.get("method", "GET"),
            endpoint=body.get("endpoint", ""),
            data=body.get("data"),
        )
        return JSONResponse(result)

    @app.get("/v1/connectors")
    async def list_connectors():
        return JSONResponse(harness.connectors.get_stats())

    @app.post("/v1/webhook/{path:path}")
    async def webhook_incoming(path: str, request: Request):
        body = await request.json()
        signature = request.headers.get("X-Webhook-Signature", "")
        result = harness.webhooks.handle_request(f"/{path}", body, signature)
        return JSONResponse(result)

    @app.post("/v1/daemon/start")
    async def start_daemon():
        harness.start_daemon()
        return {"success": True, "message": "Always-on daemon started"}

    @app.post("/v1/daemon/stop")
    async def stop_daemon():
        harness.stop_daemon()
        return {"success": True, "message": "Always-on daemon stopped"}

    @app.get("/v1/daemon/status")
    async def daemon_status():
        return JSONResponse(harness.daemon.get_stats())

    @app.post("/v1/goals/create")
    async def create_goal(request: Request):
        body = await request.json()
        goal_id = harness.create_goal(
            title=body.get("title", ""),
            description=body.get("description", ""),
            priority=body.get("priority", "high"),
        )
        return {"success": True, "goal_id": goal_id}

    @app.get("/v1/goals")
    async def list_goals():
        return JSONResponse({"goals": harness.get_goals()})

    @app.get("/v1/agents")
    async def agent_status():
        return JSONResponse(harness.get_agent_status())

    @app.get("/v1/identity")
    async def get_identity():
        return JSONResponse(harness.identity.get_stats())

    @app.post("/v1/identity/name")
    async def set_identity_name(request: Request):
        body = await request.json()
        name = body.get("name", "")
        if name:
            harness.identity.set_name(name)
            return {"success": True, "name": name, "greeting": harness.identity.get_greeting()}
        return {"success": False, "error": "No name provided"}

    @app.get("/v1/identity/greeting")
    async def get_greeting():
        return {"greeting": harness.identity.get_greeting(), "first_run": harness.identity.is_first_run()}

    @app.get("/v1/fast-cache")
    async def fast_cache_stats():
        return JSONResponse(harness.fast_cache.get_stats())

    @app.post("/v1/fast-cache/clear")
    async def clear_fast_cache():
        import os as _os
        db_path = _os.path.join(harness.data_dir, "fast_cache.db")
        import sqlite3
        with sqlite3.connect(db_path) as conn:
            conn.execute("DELETE FROM reply_cache")
        return {"success": True, "message": "Fast reply cache cleared"}

    @app.get("/v1/vault")
    async def vault_stats():
        return JSONResponse(harness.vault.get_stats())

    @app.post("/v1/vault/cleanup")
    async def vault_cleanup():
        result = harness.vault.force_cleanup()
        return JSONResponse(result)

    @app.get("/v1/vault/artifacts")
    async def list_artifacts():
        return JSONResponse({"artifacts": harness.vault.list_artifacts()})

    @app.post("/v1/mesh/converse")
    async def mesh_converse(request: Request):
        body = await request.json()
        mode = body.get("mode", "")
        topic = body.get("topic", "")
        result = harness.conversation_mesh.run_conversation(mode=mode, topic=topic)
        return JSONResponse(result)

    @app.get("/v1/mesh/pools")
    async def mesh_pools():
        return JSONResponse({"pools": harness.conversation_mesh.get_skill_pools()})

    @app.get("/v1/mesh/categories")
    async def mesh_categories():
        return JSONResponse({
            "categories": harness.conversation_mesh.get_categories(),
            "auto_categories": harness.conversation_mesh.get_auto_categories(),
        })

    @app.get("/v1/mesh/stats")
    async def mesh_stats():
        return JSONResponse(harness.conversation_mesh.get_stats())

    @app.get("/v1/subscription")
    async def subscription_status():
        return JSONResponse(harness.subscription.get_stats())

    @app.post("/v1/subscription/subscribe")
    async def subscription_subscribe(request: Request):
        body = await request.json()
        ref = body.get("payment_reference", "")
        token = body.get("token", "USDT")
        user_id = body.get("user_id", "splitbit-user")
        result = harness.subscription.subscribe(payment_reference=ref, token=token, user_id=user_id)
        return JSONResponse(result)

    @app.post("/v1/subscription/trial")
    async def subscription_trial():
        result = harness.subscription.start_trial()
        return JSONResponse(result)

    @app.post("/v1/subscription/unsubscribe")
    async def subscription_unsubscribe():
        result = harness.subscription.unsubscribe()
        return JSONResponse(result)

    @app.get("/v1/subscription/access")
    async def subscription_access():
        return JSONResponse(harness.subscription.check_access())

    @app.post("/v1/subscription/unlock")
    async def subscription_unlock(request: Request):
        body = await request.json()
        password = body.get("password", "")
        result = harness.subscription.founder_unlock(password)
        return JSONResponse(result)

    @app.get("/v1/subscription/payment")
    async def subscription_payment():
        return JSONResponse(harness.subscription.get_payment_instructions())

    @app.post("/v1/subscription/verify")
    async def subscription_verify(request: Request):
        body = await request.json()
        deposit_id = body.get("deposit_id", "")
        if not deposit_id:
            return JSONResponse({"status": "error", "message": "deposit_id required"}, status_code=400)
        result = harness.subscription.verify_payment(deposit_id)
        return JSONResponse(result)

    @app.get("/v1/subscription/auto-transfer")
    async def auto_transfer_stats():
        return JSONResponse(harness.subscription.get_auto_transfer_stats())

    @app.post("/v1/subscription/auto-transfer/process")
    async def auto_transfer_process():
        result = harness.subscription.process_auto_transfers()
        return JSONResponse(result)

    @app.post("/v1/subscription/auto-transfer/toggle")
    async def auto_transfer_toggle(request: Request):
        body = await request.json()
        enabled = body.get("enabled", True)
        result = harness.subscription.set_auto_transfer(enabled)
        return JSONResponse(result)

    @app.get("/v1/subscription/bank")
    async def bank_info():
        return JSONResponse(harness.subscription.get_bank_info())

    return app


def run_server(host: str = "0.0.0.0", port: int = 8548, harness: SplitBitHarness | None = None) -> None:
    """Run the SplitBit LLM server."""
    import uvicorn
    app = create_app(harness=harness)
    logger.info("Starting SplitBit LLM server on %s:%d", host, port)
    uvicorn.run(app, host=host, port=port, log_level="info")