Building Trustworthy Grounded Applications

import re

# Straight and curly quotes each close their own kind, so a quote nested in another
# can't shift which marks pair up.
_QUOTED = re.compile(r'"([^"]*)"|“([^”]*)”')
# A quote embedded in a sentence is routinely re-punctuated at its edges.
_EDGE = " .,;:!?"


def quoted_spans(text: str, min_length: int = 12) -> list[str]:
    """Quoted spans, paired in order; by default short ones such as a single quoted word are skipped."""
    spans = [straight or curly for straight, curly in _QUOTED.findall(text)]
    return [span for span in spans if len(span) >= min_length]


def quotes_balanced(text: str) -> bool:
    """False when a quotation mark has no partner, so some quoted text could not be checked."""
    depth = 0  # each curly closing mark must follow an opening one
    for ch in text:
        if ch == "“":
            depth += 1
        elif ch == "”":
            depth -= 1
            if depth < 0:
                return False
    return depth == 0 and text.count('"') % 2 == 0


def normalize(text: str) -> str:
    """Fold typographic variation so a real quote is not flagged over punctuation."""
    text = text.replace("’", "'").replace("‘", "'")
    text = text.replace("“", '"').replace("”", '"')
    text = text.replace("…", "...")
    return re.sub(r"\s+", " ", text).strip().lower()


def unverified_quotes(answer: str, sources: list[str]) -> list[str]:
    """Return every quoted span in `answer` that is supported by no single source."""
    normalized = [normalize(s) for s in sources]
    unverified = []
    for span in quoted_spans(answer):
        # Split on ellipses so "A ... B" verifies if both halves are present.
        parts = [p for p in (p.strip(_EDGE) for p in normalize(span).split("...")) if p]
        # Every part must appear in the SAME source. Searching the sources
        # concatenated would let a quote stitched from two documents verify,
        # which is exactly the fabrication this check exists to catch.
        if not parts or not any(all(part in src for part in parts) for src in normalized):
            unverified.append(span)
    return unverified

// Straight and curly quotes each close their own kind, so a quote nested in another
// can't shift which marks pair up.
const QUOTED = /"([^"]*)"|“([^”]*)”/g;
// A quote embedded in a sentence is routinely re-punctuated at its edges.
const EDGE = /^[\s.,;:!?]+|[\s.,;:!?]+$/g;

/** Quoted spans, paired in order; by default short ones such as a single quoted word are skipped. */
export function quotedSpans(text: string, minLength = 12): string[] {
  const spans = [...text.matchAll(QUOTED)].map((m) => m[1] ?? m[2]);
  return spans.filter((span) => span.length >= minLength);
}

/** False when a quotation mark has no partner, so some quoted text could not be checked. */
export function quotesBalanced(text: string): boolean {
  let depth = 0; // each curly closing mark must follow an opening one
  for (const ch of text) {
    if (ch === "“") depth += 1;
    else if (ch === "”" && --depth < 0) return false;
  }
  return depth === 0 && (text.match(/"/g) ?? []).length % 2 === 0;
}

/** Fold typographic variation so a real quote is not flagged over punctuation. */
export function normalize(text: string): string {
  return text
    .replace(/[‘’]/g, "'")
    .replace(/[“”]/g, '"')
    .replace(/…/g, "...")
    .replace(/\s+/g, " ")
    .trim()
    .toLowerCase();
}

/** Return every quoted span in `answer` that is supported by no single source. */
export function unverifiedQuotes(answer: string, sources: string[]): string[] {
  const normalized = sources.map(normalize);
  const unverified: string[] = [];
  for (const span of quotedSpans(answer)) {
    // Split on ellipses so "A ... B" verifies if both halves are present.
    const parts = normalize(span)
      .split("...")
      .map((p) => p.replace(EDGE, ""))
      .filter(Boolean);
    // Every part must appear in the SAME source. Searching the sources
    // concatenated would let a quote stitched from two documents verify,
    // which is exactly the fabrication this check exists to catch.
    const supported =
      parts.length > 0 &&
      normalized.some((src) => parts.every((part) => src.includes(part)));
    if (!supported) {
      unverified.push(span);
    }
  }
  return unverified;
}

MAX_TURNS = 20
MAX_CHARS = 60_000


def trim_history(messages: list[dict], max_turns: int = MAX_TURNS,
                 max_chars: int = MAX_CHARS) -> list[dict]:
    """Trim oldest turns first, by count and by size.

    Size matters as much as count: verified answers carrying quotations and
    citations run long, so a turn-count cap alone can still overflow limits.
    """
    system = [m for m in messages if m["role"] == "system"]
    rest = [m for m in messages if m["role"] != "system"]

    rest = rest[-max_turns:]
    while rest and sum(len(m["content"]) for m in rest) > max_chars:
        rest.pop(0)

    return system + rest

export interface ChatMessage {
  role: "system" | "user" | "assistant";
  content: string;
}

export const MAX_TURNS = 20;
export const MAX_CHARS = 60_000;

/**
 * Trim oldest turns first, by count and by size.
 *
 * Size matters as much as count: verified answers carrying quotations and
 * citations run long, so a turn-count cap alone can still overflow limits.
 */
export function trimHistory(
  messages: ChatMessage[],
  maxTurns = MAX_TURNS,
  maxChars = MAX_CHARS,
): ChatMessage[] {
  const system = messages.filter((m) => m.role === "system");
  let rest = messages.filter((m) => m.role !== "system").slice(-maxTurns);

  while (rest.length && rest.reduce((n, m) => n + m.content.length, 0) > maxChars) {
    rest = rest.slice(1);
  }

  return [...system, ...rest];
}

def answer_with_verification(client, question: str, publisher: str) -> dict:
    """Generate, verify, widen once, then decline rather than risk a fabrication.

    Returns a dict with `status` of "verified", "unverified", "unsupported", or "ungrounded".
    Only "verified" and "unverified" are safe to display as an answer, and only
    "verified" may say its quotations were checked.
    """
    unverified: list[str] = []

    for sources_limit in (5, 10):  # one escalation, not a loop
        result = client.grounded_completion(
            question, rag_publisher=publisher,
            sources_limit=sources_limit, include_citations=True,
        )

        # Check grounding FIRST. With no sources there is nothing to compare
        # against, so every span below would fail containment automatically.
        # Grounding never fails the request: an ungrounded answer still arrives
        # fluent and complete, so this flag must be checked explicitly.
        if not result["sources_returned"]:
            return {"status": "ungrounded", "text": None}

        answer = result["choices"][0]["message"]["content"]
        sources = [s for c in result["citations"] for s in c["snippets"]]

        unverified = unverified_quotes(answer, sources) + unverified_titles(answer, sources)
        if not unverified:
            # No quotation, a quote too short to check, or an unpaired mark: nothing was verified.
            spans = quoted_spans(answer)
            checked = bool(spans) and quotes_balanced(answer) and len(spans) == len(quoted_spans(answer, 1))
            status = "verified" if checked else "unverified"
            return {"status": status, "text": answer, "citations": result["citations"]}

    return {"status": "unsupported", "text": None, "unverified": unverified}

/**
 * Generate, verify, widen once, then decline rather than risk a fabrication.
 *
 * `status` is "verified", "unverified", "unsupported", or "ungrounded".
 * Only "verified" and "unverified" are safe to display as an answer, and only
 * "verified" may say its quotations were checked.
 */
export async function answerWithVerification(
  client: GroundedClient,
  question: string,
  publisher: string,
) {
  let unverified: string[] = [];

  for (const sourcesLimit of [5, 10]) {  // one escalation, not a loop
    const result = await client.groundedCompletion(question, {
      rag_publisher: publisher,
      sources_limit: sourcesLimit,
      include_citations: true,
    });

    // Check grounding FIRST. With no sources there is nothing to compare
    // against, so every span below would fail containment automatically.
    // Grounding never fails the request: an ungrounded answer still arrives
    // fluent and complete, so this flag must be checked explicitly.
    if (!result.sources_returned) {
      return { status: "ungrounded" as const, text: null };
    }

    const answer = result.choices[0].message.content;
    const sources = result.citations.flatMap((c) => c.snippets);

    unverified = [
      ...unverifiedQuotes(answer, sources),
      ...unverifiedTitles(answer, sources),
    ];
    if (unverified.length === 0) {
      // No quotation, a quote too short to check, or an unpaired mark: nothing was verified.
      const checked =
        quotedSpans(answer).length > 0 &&
        quotesBalanced(answer) &&
        quotedSpans(answer).length === quotedSpans(answer, 1).length;
      const status = checked ? ("verified" as const) : ("unverified" as const);
      return { status, text: answer, citations: result.citations };
    }
  }

  return { status: "unsupported" as const, text: null, unverified };
}