scaffold/rivet_serve.py

"""Rivet Server — the code assistant endpoint.

Wraps Ollama with the discipline gate and context loader.
Faculty hit this endpoint; Rivet responds with architecture-aware,
discipline-gated suggestions.

Usage:
    python rivet_serve.py                    # Start server on port 8100
    python rivet_serve.py --port 8200       # Custom port
    python rivet_serve.py --model qwen3.5:27b  # Use fallback model
"""

import argparse
import json
import time
from http.server import HTTPServer, BaseHTTPRequestHandler

import requests

from context_loader import load_context
from discipline_gate import run_gate, GateResult


# Config
OLLAMA_URL = "http://localhost:11434"
DEFAULT_MODEL = "rivet"
FALLBACK_MODEL = "qwen3.5:27b"

# Load context once at startup
print("Loading Rivet context...", flush=True)
CONTEXT = load_context()
SYSTEM_PROMPT = CONTEXT.to_system_context() + """

---

# YOUR ROLE

You are Rivet, a senior engineer embedded with the Multiverse Campus team.

## Rules
1. Never suggest a destructive migration. Staging and prod share the database.
2. Every code suggestion includes: what it changes, what it could break, and what tests verify it.
3. State your confidence level: HIGH (traced the full path), MEDIUM (read the code), LOW (reasoning from architecture).
4. Auth changes require explicit callout: "This touches authentication. Review with security before merging."
5. Flag known vulnerability patterns proactively.
6. You are a colleague, not the lead. Suggest, don't decree.
7. If you cannot verify your suggestion compiles, say so.
"""

print(f"Context loaded: ~{CONTEXT.token_estimate} tokens", flush=True)


def query_ollama(prompt: str, model: str = DEFAULT_MODEL) -> str:
    """Send a prompt to Ollama and get the response."""
    try:
        resp = requests.post(
            f"{OLLAMA_URL}/api/generate",
            json={
                "model": model,
                "prompt": prompt,
                "system": SYSTEM_PROMPT,
                "stream": False,
                "options": {
                    "temperature": 0.3,
                    "top_p": 0.9,
                    "num_ctx": 32768,
                },
            },
            timeout=120,
        )
        resp.raise_for_status()
        return resp.json().get("response", "")
    except requests.exceptions.ConnectionError:
        return f"ERROR: Cannot connect to Ollama at {OLLAMA_URL}. Is it running?"
    except requests.exceptions.Timeout:
        return "ERROR: Ollama request timed out (120s). Try a shorter question or check the model."
    except Exception as e:
        return f"ERROR: {e}"


class RivetHandler(BaseHTTPRequestHandler):
    model = DEFAULT_MODEL

    def do_POST(self):
        if self.path == "/ask":
            content_length = int(self.headers.get("Content-Length", 0))
            body = json.loads(self.rfile.read(content_length))
            question = body.get("question", "")
            user = body.get("user", "anonymous")

            t0 = time.time()

            # Get response from model
            response = query_ollama(question, model=self.model)

            # Run discipline gate on the response
            gate_result = run_gate(response, context=question)
            gate_warnings = gate_result.format_warnings()

            # Compose final response
            final_response = response
            if gate_warnings:
                final_response = gate_warnings + "\n\n---\n\n" + response
            if not gate_result.passed:
                final_response = (
                    "🛑 **BLOCKED by Discipline Gate**\n\n"
                    "The suggested approach was blocked for safety reasons:\n"
                    + gate_warnings
                    + "\n\nPlease rephrase your request or ask for a safe alternative."
                )

            elapsed = time.time() - t0

            result = {
                "response": final_response,
                "gate_passed": gate_result.passed,
                "flags": gate_result.flags,
                "confidence": gate_result.confidence.value,
                "user": user,
                "elapsed_seconds": round(elapsed, 1),
            }

            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps(result).encode())

        elif self.path == "/health":
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({
                "status": "ok",
                "model": self.model,
                "context_tokens": CONTEXT.token_estimate,
            }).encode())

        else:
            self.send_response(404)
            self.end_headers()

    def do_GET(self):
        if self.path == "/health":
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            self.wfile.write(json.dumps({
                "status": "ok",
                "model": self.model,
                "context_tokens": CONTEXT.token_estimate,
            }).encode())
        else:
            self.send_response(404)
            self.end_headers()

    def log_message(self, format, *args):
        print(f"[rivet] {args[0]}", flush=True)


def main():
    parser = argparse.ArgumentParser(description="Rivet Code Assistant Server")
    parser.add_argument("--port", type=int, default=8100)
    parser.add_argument("--model", default=DEFAULT_MODEL)
    args = parser.parse_args()

    RivetHandler.model = args.model

    server = HTTPServer(("0.0.0.0", args.port), RivetHandler)
    print(f"\n🔩 Rivet listening on port {args.port}", flush=True)
    print(f"   Model: {args.model}", flush=True)
    print(f"   Context: ~{CONTEXT.token_estimate} tokens loaded", flush=True)
    print(f"   POST /ask {{\"question\": \"...\", \"user\": \"...\"}}", flush=True)
    print(f"   GET  /health", flush=True)
    print(flush=True)

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n🔩 Rivet shutting down.", flush=True)
        server.shutdown()


if __name__ == "__main__":
    main()