File size: 6,251 Bytes
e2921fc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """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()
|