Review a model’s proposal before saving it

import argparse
from contextlib import closing
from itertools import islice
import json
import math
from pathlib import Path
import sqlite3
from urllib.error import HTTPError
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

BASE = "http://127.0.0.1:11434"
MODEL = "all-minilm:latest"
MAX_PASSAGES = 200


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def api(endpoint, payload=None):
    request = Request(BASE + endpoint)
    if payload is not None:
        request.data = json.dumps(payload).encode("utf-8")
        request.add_header("Content-Type", "application/json")
    opener = build_opener(ProxyHandler({}), NoRedirects())
    with opener.open(request, timeout=120) as response:
        return json.load(response)


def model_digest():
    for model in api("/api/tags")["models"]:
        if model["name"] == MODEL:
            return model["digest"]
    raise ValueError(f"Pull {MODEL} with Ollama first.")


def embed(texts):
    result = api("/api/embed", {
        "model": MODEL,
        "input": texts,
        "truncate": False,
    })["embeddings"]
    if len(result) != len(texts):
        raise ValueError("Ollama returned the wrong number of embeddings.")
    return result


def cosine(left, right):
    if len(left) != len(right):
        raise ValueError("Vector dimensions changed. Rebuild the index.")
    length = math.hypot(*left) * math.hypot(*right)
    return sum(x * y for x, y in zip(left, right)) / length if length else 0


def build_index(folder, database):
    from index_notes import passages

    rows = list(islice(passages(folder), MAX_PASSAGES + 1))
    if not rows or len(rows) > MAX_PASSAGES:
        raise ValueError(f"Use 1–{MAX_PASSAGES} nonempty passages for this example.")
    digest = model_digest()
    stored = []
    for start in range(0, len(rows), 16):
        batch = rows[start:start + 16]
        vectors = embed([text for _, _, text in batch])
        stored.extend((*row, json.dumps(vector)) for row, vector in zip(batch, vectors))
    if model_digest() != digest:
        raise ValueError("The embedding model changed during indexing. Build the index again.")
    with closing(sqlite3.connect(database)) as connection, connection:
        connection.execute("CREATE TABLE IF NOT EXISTS model (name TEXT, digest TEXT)")
        connection.execute(
            "CREATE TABLE IF NOT EXISTS passages "
            "(source TEXT, number INTEGER, text TEXT, vector TEXT)"
        )
        connection.execute("DELETE FROM passages")
        connection.execute("DELETE FROM model")
        connection.execute("INSERT INTO model VALUES (?, ?)", (MODEL, digest))
        connection.executemany("INSERT INTO passages VALUES (?, ?, ?, ?)", stored)
    print(f"Saved {len(stored)} embeddings in {database}.")


def find_notes(database, query):
    if not query.strip() or not database.is_file():
        raise ValueError("Build the index and enter a nonempty question.")
    with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as connection:
        connection.execute("BEGIN")
        identity = connection.execute("SELECT name, digest FROM model").fetchone()
        rows = connection.execute("SELECT source, number, text, vector FROM passages").fetchall()
    digest = model_digest()
    if identity != (MODEL, digest):
        raise ValueError("The embedding model differs from the one in the index. Rebuild it.")
    vector, = embed([query])
    if model_digest() != digest:
        raise ValueError("The embedding model changed during the query. Rebuild the index.")
    matches = [
        (cosine(json.loads(embedding), vector), source, number, text)
        for source, number, text, embedding in rows
    ]
    ranked = sorted(matches, key=lambda match: match[0], reverse=True)
    for score, source, number, text in ranked[:3]:
        print(f"{score:.3f} | {source}, passage {number}\n{text}\n")


def main():
    parser = argparse.ArgumentParser(description="Store local embeddings and search them later.")
    parser.add_argument("--db", type=Path, default=Path("notes-vectors.sqlite3"))
    commands = parser.add_subparsers(dest="command", required=True)
    build = commands.add_parser("build")
    build.add_argument("folder", type=Path)
    find = commands.add_parser("find")
    find.add_argument("query")
    args = parser.parse_args()

    try:
        if args.command == "build":
            build_index(args.folder, args.db)
        else:
            find_notes(args.db, args.query)
    except HTTPError as error:
        raise SystemExit(f"Ollama returned HTTP {error.code}: {error.read().decode()}") from None
    except (OSError, sqlite3.Error, ValueError) as error:
        raise SystemExit(str(error)) from None


if __name__ == "__main__":
    main()

import argparse
from contextlib import closing
import json
from pathlib import Path
import re
import sqlite3

from semantic_notes import api

SLOTS = {
    "scanner-1400": "Flatbed scanner at 14:00",
    "scanner-1500": "Flatbed scanner at 15:00",
    "projector-1400": "Projector at 14:00",
}


def requested_slots(request):
    # This small demo supports explicit 24-hour times such as 14:00.
    times = re.findall(r"(?<![\w:])[0-9]{1,2}:[0-9]{2}(?![\w:])", request)
    if len(times) != 1:
        return []
    return [slot for slot, label in SLOTS.items() if label.endswith(" " + times[0])]


def suggest_slot(request):
    choices = "\n".join(f"{key}: {label}" for key, label in SLOTS.items())
    reply = api("/api/chat", {
        "model": "qwen2.5:1.5b",
        "messages": [
            {"role": "system", "content": (
                "Choose a slot only when the request explicitly specifies one available item "
                "and its time. Use null for an unavailable item, missing time, or ambiguity.\n"
                f"Slots:\n{choices}"
            )},
            {"role": "user", "content": request},
        ],
        "format": {
            "type": "object",
            "properties": {"slot": {"enum": [*SLOTS, None]}},
            "required": ["slot"],
            "additionalProperties": False,
        },
        "options": {"temperature": 0, "num_ctx": 2048, "num_predict": 80},
        "stream": False,
    })
    proposal = json.loads(reply["message"]["content"])
    if not isinstance(proposal, dict) or set(proposal) != {"slot"}:
        raise ValueError("The model did not return a slot proposal.")
    slot = proposal["slot"]
    if slot is not None and (not isinstance(slot, str) or slot not in SLOTS):
        raise ValueError("The proposed slot is outside the available choices.")
    return slot


def reserve(database, slot):
    if slot not in SLOTS:
        raise ValueError("Unknown slot.")
    with closing(sqlite3.connect(database)) as db, db:
        db.execute("CREATE TABLE IF NOT EXISTS reservations (slot TEXT PRIMARY KEY)")
        result = db.execute(
            "INSERT INTO reservations(slot) VALUES (?) ON CONFLICT(slot) DO NOTHING", (slot,)
        )
        return result.rowcount == 1


def review(database, request):
    candidates = requested_slots(request)
    if not candidates:
        print("Please supply one available item and a specific time. Nothing was saved.")
        return
    slot = suggest_slot(request)
    if slot not in candidates:
        print("Please supply one available item and a specific time. Nothing was saved.")
        return
    print(f"Original request: {request!r}\nProposed reservation: {SLOTS[slot]}")
    try:
        answer = input("Type APPROVE to save this reservation: ")
    except (EOFError, KeyboardInterrupt):
        answer = ""
        print()
    if answer.strip() != "APPROVE":
        print("Cancelled. Nothing was saved.")
        return
    if reserve(database, slot):
        print(f"Saved: {SLOTS[slot]}")
    else:
        print("That slot is already reserved. No new record was saved.")


def list_reservations(database):
    rows = []
    if database.is_file():
        with closing(sqlite3.connect(database.resolve().as_uri() + "?mode=ro", uri=True)) as db:
            rows = db.execute("SELECT slot FROM reservations ORDER BY slot").fetchall()
    for (slot,) in rows:
        print(SLOTS[slot])
    if not rows:
        print("No reservations.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Review a proposal for a fictional reservation.")
    parser.add_argument("--db", type=Path, default=Path("demo-reservations.sqlite3"))
    commands = parser.add_subparsers(dest="command", required=True)
    propose = commands.add_parser("propose")
    propose.add_argument("request")
    commands.add_parser("list")
    args = parser.parse_args()
    try:
        if args.command == "propose":
            review(args.db, args.request)
        else:
            list_reservations(args.db)
    except (OSError, KeyError, ValueError, sqlite3.Error) as error:
        raise SystemExit(str(error)) from None