Move scaffold/scaffold/rivet_serve.py -> scaffold/rivet_serve.py

scaffold/rivet_serve.py

@@ -0,0 +1,184 @@
+"""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()