Astraea: WA family-law multi-agent site (Flask + Ollama semantic RAG)

app.py

@@ -0,0 +1,192 @@
+# -*- coding: utf-8 -*-
+"""Astraea — a constellation of Washington State family-law specialist agents.
+
+Flask app + semantic RAG backend. LLM (granite4.2 RAG lane) and embeddings
+(nomic-embed-text-v2-moe) run on the Ollama host (nightmare).
+"""
+import json
+import re
+import time
+import urllib.request
+
+from flask import Flask, jsonify, render_template, request
+
+import rag
+from agents import (
+    AGENTS, AGENT_BY_ID, OLLAMA_URL, RAG_MODEL, GENERAL_MODEL, DISCLAIMER,
+)
+
+app = Flask(__name__)
+
+
+def _ollama_chat(model, messages, num_ctx=16384, num_predict=1024, temperature=0.1, timeout=180):
+    payload = {
+        "model": model,
+        "messages": messages,
+        "stream": False,
+        "think": False,
+        "options": {
+            "temperature": temperature,
+            "num_predict": num_predict,
+            "num_ctx": num_ctx,
+        },
+    }
+    req = urllib.request.Request(
+        f"{OLLAMA_URL}/api/chat",
+        data=json.dumps(payload).encode("utf-8"),
+        headers={"Content-Type": "application/json"},
+    )
+    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
+    with opener.open(req, timeout=timeout) as r:
+        d = json.loads(r.read().decode("utf-8"))
+    msg = d.get("message", {})
+    content = msg.get("content") or msg.get("thinking") or ""
+    return content.strip()
+
+
+def _strip_model_sources(text):
+    """Remove a model-written trailing 'Sources:'/'References:' block (defensive)."""
+    text = re.split(r"\n\s*(?:Sources|References|Citations)\s*:\s*\n", text, flags=re.I)[0]
+    text = re.split(r"\n\s*Sources?\s*$", text, flags=re.I)[0]
+    return text.strip()
+
+
+def _build_answer(agent, message, history):
+    top = rag.retrieve(message, agent["id"], top_k=5)
+    if not top:
+        return {
+            "answer": "I couldn't find relevant Washington law on that in my reference "
+                      "library yet. Try rephrasing, or ask the Navigator to point you to "
+                      "the right specialist.",
+            "citations": [],
+            "grounded": False,
+        }
+
+    context_blocks = []
+    citations = []
+    seen = set()
+    for i, (score, c) in enumerate(top, 1):
+        key = (c["source"], c["title"])
+        if key in seen:
+            continue
+        seen.add(key)
+        context_blocks.append(f"[{i}] ({c['source']} — {c['title']})\n{c['text']}")
+        citations.append({
+            "source": c["source"],
+            "title": c["title"],
+            "score": round(score, 3),
+            "snippet": c["text"][:260] + ("…" if len(c["text"]) > 260 else ""),
+        })
+
+    system = (
+        f"{agent['system']}\n\n"
+        "Rules:\n"
+        "- Answer the user's question DIRECTLY and concisely. Begin your answer immediately — "
+        "do NOT restate the question, do NOT narrate your reasoning, and do NOT say what you "
+        "are about to do.\n"
+        "- Answer using ONLY the reference documents provided below.\n"
+        "- Cite the RCW section (or source) inline for every legal claim, e.g. (RCW 26.09.030).\n"
+        "- If the documents do not contain the answer, say so clearly and suggest which "
+        "specialist or official resource to consult.\n"
+        "- Be practical and plain-English, specific to Washington State.\n"
+        "- Do NOT write a 'Sources' list at the end; cite inline only."
+    )
+    user = (
+        "REFERENCE DOCUMENTS:\n\n"
+        + "\n\n".join(context_blocks)
+        + f"\n\nUSER QUESTION: {message}"
+    )
+
+    messages = [{"role": "system", "content": system}]
+    if history:
+        for turn in history[-6:]:
+            if turn.get("role") in ("user", "assistant"):
+                messages.append({"role": turn["role"], "content": turn["content"]})
+    messages.append({"role": "user", "content": user})
+
+    try:
+        answer = _ollama_chat(RAG_MODEL, messages)
+    except Exception as e:
+        # Fall back to the general model if the RAG lane is unavailable.
+        try:
+            answer = _ollama_chat(GENERAL_MODEL, messages, num_ctx=16384)
+        except Exception:
+            return {
+                "answer": "The legal engine is temporarily unavailable. Please try again in a "
+                          "moment. (Backend LLM could not be reached.)",
+                "citations": [],
+                "grounded": False,
+            }
+
+    answer = _strip_model_sources(answer)
+    return {
+        "answer": answer,
+        "citations": citations,
+        "grounded": True,
+        "model": RAG_MODEL,
+    }
+
+
+@app.route("/")
+def index():
+    agents_public = [
+        {k: a[k] for k in ("id", "name", "emoji", "tagline", "description", "accent")}
+        for a in AGENTS
+    ]
+    return render_template("index.html", agents=agents_public, disclaimer=DISCLAIMER)
+
+
+@app.route("/api/agents")
+def api_agents():
+    return jsonify([
+        {k: a[k] for k in ("id", "name", "emoji", "tagline", "description", "accent")}
+        for a in AGENTS
+    ])
+
+
+@app.route("/api/chat", methods=["POST"])
+def api_chat():
+    data = request.get_json(force=True, silent=True) or {}
+    agent_id = (data.get("agent_id") or "navigator").strip()
+    message = (data.get("message") or "").strip()
+    history = data.get("history") or []
+    agent = AGENT_BY_ID.get(agent_id)
+    if not agent:
+        return jsonify({"error": "Unknown agent"}), 400
+    if not message:
+        return jsonify({"error": "Message required"}), 400
+    if len(message) > 4000:
+        message = message[:4000]
+
+    t0 = time.time()
+    result = _build_answer(agent, message, history)
+    result["latency_ms"] = round((time.time() - t0) * 1000)
+    result["agent"] = agent_id
+    return jsonify(result)
+
+
+@app.route("/api/health")
+def health():
+    return jsonify({"ok": True, "agent_count": len(AGENTS), "index": rag.index_status()})
+
+
+@app.route("/api/index/status")
+def index_status():
+    return jsonify(rag.index_status())
+
+
+@app.route("/api/reindex", methods=["POST"])
+def reindex():
+    try:
+        idx = rag.build_index()
+        return jsonify({"ok": True, "chunks": len(idx)})
+    except Exception as e:
+        return jsonify({"ok": False, "error": str(e)}), 500
+
+
+# Kick off the background index build once at import (before serving).
+rag.start_background_index()
+
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=5000, threaded=True)