feat(py): citation-request tracker and tool-result tagging

pkg-py/src/commons/_citations.py

@@ -1,26 +1,37 @@
-"""Parsing a citation block and verifying its quote against a trusted corpus.
+"""Citations: verifying a quote against a trusted corpus, and asking for one.
 
 The normalization rules and the matching verdicts are a cross-language contract
-pinned by ``tests/shared/citations.json``; change that fixture, not just this
-file. ``pkg-r/R/citations.R`` implements the same contract for R.
+pinned by ``tests/shared/citations.json``, and where the citation request lands
+by ``tests/shared/citation-request.json``; change those fixtures, not just this
+file. ``pkg-r/R/citations.R`` implements the same contracts for R.
 """
 
 from __future__ import annotations
 
 import html
 import re
 from collections.abc import Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Literal, get_args
 
+from chatlas import ContentToolResult, Turn
+from chatlas.types import ContentText
+
+from ._prompt import read_prompt
+from ._provenance import Tag
+
 __all__ = [
     "CitationDecision",
+    "CitationRequest",
     "CorpusEntry",
     "ParsedCitation",
     "citation_aside_html",
+    "citation_reminder_text",
     "match_citation",
     "normalize_citation",
     "parse_commons_citation",
+    "tool_result",
+    "turn_has_user_message",
 ]
 
 # The minimum length is a guard, not a tuning knob: a fragment this short can
@@ -201,3 +212,66 @@ def citation_aside_html(quote: str, explanation: str, label: str, kind: str) ->
 # Ampersands first, so the entities this generates are not escaped again.
 def _escape_attr(text: str) -> str:
     return text.replace("&", "&").replace('"', """)
+
+
+def tool_result(value: Any, tag: Tag | None = None) -> ContentToolResult:
+    """A tool result carrying the provenance tag of the output it holds.
+
+    The tag is read back off ``extra`` when the turn is classified, so it is
+    set here rather than at the point a result is added to the conversation.
+    """
+    return ContentToolResult(value=value, extra={"commons_tag": tag})
+
+
+def citation_reminder_text() -> str:
+    """The reminder text, from the prompt file both packages ship."""
+    return read_prompt("citation-request.md")
+
+
+@dataclass
+class CitationRequest:
+    """Whether this user turn has carried the citation reminder yet.
+
+    The citation contract lives in the system prompt; this is the nudge that
+    rides on the first tool result of a turn whose output has to be cited.
+    """
+
+    reminder: str = field(default_factory=citation_reminder_text)
+    requested: bool = False
+
+    def add_request(self, result: ContentToolResult) -> ContentToolResult:
+        """Add the reminder to ``result``, unless this turn has asked already.
+
+        An errored result passes through without spending the request: the
+        model is sent the error rather than the value, so a reminder added to
+        the value would never arrive.
+        """
+        if self.requested or result.error is not None:
+            return result
+        self.requested = True
+        result.value = _with_reminder(result.value, self.reminder)
+        return result
+
+    def reset(self) -> None:
+        """Start a new user turn, so the next eligible result asks again."""
+        self.requested = False
+
+
+def _with_reminder(value: Any, reminder: str) -> Any:
+    if isinstance(value, str):
+        return f"{value}\n\n{reminder}"
+    part = ContentText(text=reminder)
+    if isinstance(value, list):
+        return [*value, part]
+    # A tool result can hold something that is neither: it becomes the first
+    # part rather than being reformatted to make room for the reminder.
+    return [value, part]
+
+
+def turn_has_user_message(turn: Turn) -> bool:
+    """Whether a turn asks something new, rather than continuing the tool loop.
+
+    A turn of nothing but tool results is the same question still running, so
+    the reminder stays spent until the person says something.
+    """
+    return any(not isinstance(content, ContentToolResult) for content in turn.contents)

pkg-py/tests/test_citation_request.py

@@ -0,0 +1,89 @@
+"""The citation-request tracker and the tag a tool result carries."""
+
+from typing import Any
+
+from chatlas import ContentToolResult, UserTurn
+from chatlas.types import ContentText
+
+from commons._citations import CitationRequest, tool_result, turn_has_user_message
+from commons._prompt import read_prompt
+from commons._provenance import Tag
+
+from ._shared import load_shared_fixture
+
+
+def _fixture_value(spec: dict[str, Any]) -> Any:
+    if spec["kind"] == "text":
+        return spec["text"]
+    return [ContentText(text=part) for part in spec["parts"]]
+
+
+def _value_shape(value: Any) -> dict[str, Any]:
+    if isinstance(value, str):
+        return {"kind": "text", "text": value}
+    return {"kind": "parts", "parts": [part.text for part in value]}
+
+
+def test_tool_result_carries_its_provenance_tag():
+    result = tool_result("6 rows", tag=Tag.B)
+
+    assert isinstance(result, ContentToolResult)
+    assert result.value == "6 rows"
+    assert result.extra == {"commons_tag": Tag.B}
+
+
+def test_shared_citation_request_cases():
+    section = load_shared_fixture("citation-request")["requests"]
+    assert section["cases"]
+
+    for case in section["cases"]:
+        tracker = CitationRequest(reminder=section["reminder"])
+        for index, step in enumerate(case["steps"]):
+            if step["action"] == "reset":
+                tracker.reset()
+                continue
+            result = tracker.add_request(tool_result(_fixture_value(step["value"]), tag=Tag.B))
+            assert _value_shape(result.value) == step["expected"], (
+                f"{case['name']} step {index}"
+            )
+
+
+def test_a_value_that_is_neither_text_nor_parts_becomes_parts():
+    frame = {"rows": 6}
+    result = CitationRequest(reminder="REMINDER").add_request(
+        tool_result(frame, tag=Tag.B)
+    )
+
+    assert result.value[0] == frame
+    assert result.value[1].text == "REMINDER"
+
+
+def test_an_errored_result_does_not_spend_the_request():
+    tracker = CitationRequest(reminder="REMINDER")
+    failed = tool_result("6 rows", tag=Tag.B)
+    failed.error = ValueError("boom")
+
+    assert tracker.add_request(failed).value == "6 rows"
+
+    result = tracker.add_request(tool_result("3 rows", tag=Tag.B))
+    assert result.value == "3 rows\n\nREMINDER"
+
+
+def test_shared_reset_cases():
+    section = load_shared_fixture("citation-request")["resets"]
+    assert section["cases"]
+
+    for case in section["cases"]:
+        contents = [
+            ContentText(text="a question")
+            if kind == "text"
+            else tool_result("6 rows", tag=Tag.B)
+            for kind in case["contents"]
+        ]
+        turn = UserTurn(contents)
+
+        assert turn_has_user_message(turn) is case["resets"], case["name"]
+
+
+def test_the_reminder_defaults_to_the_shipped_prompt_text():
+    assert CitationRequest().reminder == read_prompt("citation-request.md")

pkg-r/tests/testthat/fixtures/shared/citation-request.json

@@ -0,0 +1,118 @@
+{
+  "description": "How the citation request reaches the model. The citation contract itself lives in the system prompt; a short reminder rides along on the first tool result of a user turn whose output has to be cited. The source is tests/shared/citation-request.json; the copy under pkg-r/tests/testthat/fixtures/shared/ is generated by scripts/sync-shared.sh. Edit the source and re-run that script.",
+  "requests": {
+    "description": "Each case drives one tracker through a sequence of steps. `add` passes a tool result through the tracker and states the value it must come back with; `reset` is what a new user message does to the tracker. The reminder is a stand-in for the real text, which lives in prompts/citation-request.md, so these cases pin placement rather than wording. A value is either the text a tool returned or a list of content parts, and the reminder joins each shape in its own way.",
+    "reminder": "SHORT CITATION REMINDER",
+    "cases": [
+      {
+        "name": "one reminder per user turn",
+        "steps": [
+          {
+            "action": "add",
+            "value": {
+              "kind": "text",
+              "text": "6 rows"
+            },
+            "expected": {
+              "kind": "text",
+              "text": "6 rows\n\nSHORT CITATION REMINDER"
+            }
+          },
+          {
+            "action": "add",
+            "value": {
+              "kind": "text",
+              "text": "3 rows"
+            },
+            "expected": {
+              "kind": "text",
+              "text": "3 rows"
+            }
+          },
+          {
+            "action": "reset"
+          },
+          {
+            "action": "add",
+            "value": {
+              "kind": "text",
+              "text": "2 rows"
+            },
+            "expected": {
+              "kind": "text",
+              "text": "2 rows\n\nSHORT CITATION REMINDER"
+            }
+          }
+        ]
+      },
+      {
+        "name": "content parts take the reminder as another part",
+        "steps": [
+          {
+            "action": "add",
+            "value": {
+              "kind": "parts",
+              "parts": [
+                "output"
+              ]
+            },
+            "expected": {
+              "kind": "parts",
+              "parts": [
+                "output",
+                "SHORT CITATION REMINDER"
+              ]
+            }
+          },
+          {
+            "action": "add",
+            "value": {
+              "kind": "parts",
+              "parts": [
+                "more output"
+              ]
+            },
+            "expected": {
+              "kind": "parts",
+              "parts": [
+                "more output"
+              ]
+            }
+          }
+        ]
+      }
+    ]
+  },
+  "resets": {
+    "description": "Which turns clear the flag, so the next eligible tool result asks for a citation again. A turn that carries a real user message starts a new question; a turn of nothing but tool results is the same question still running, and an empty turn is neither.",
+    "cases": [
+      {
+        "name": "a user message",
+        "contents": [
+          "text"
+        ],
+        "resets": true
+      },
+      {
+        "name": "tool results only",
+        "contents": [
+          "tool_result"
+        ],
+        "resets": false
+      },
+      {
+        "name": "a user message after tool results",
+        "contents": [
+          "tool_result",
+          "text"
+        ],
+        "resets": true
+      },
+      {
+        "name": "no contents",
+        "contents": [],
+        "resets": false
+      }
+    ]
+  }
+}