docs/implementation-plan.md

# Implementation Plan: `tei-annotator`

Prompt:

> Design a Python library called `tei-annotator` for annotating text with TEI XML tags using a two-stage LLM pipeline. The library should:
>
> **Inputs:**
>
> - A text string to annotate (may already contain partial XML)
> - An injected `call_fn: (str) -> str` for calling an arbitrary inference endpoint
> - An `EndpointCapability` enum indicating whether the endpoint is plain text generation, JSON-constrained, or a native extraction model like GLiNER2
> - A `TEISchema` data structure describing a subset of TEI elements with descriptions, allowed attributes, and legal child elements

> The source text must never be modified by any model. Provide a package structure, all key data structures, and a step-by-step execution flow.

## Package Structure

```
tei_annotator/
├── __init__.py
├── models/
│   ├── schema.py           # TEI element/attribute data structures
│   └── spans.py            # Span manifest data structures
├── detection/
│   └── gliner_detector.py  # Optional local GLiNER first-pass span detection
├── chunking/
│   └── chunker.py          # Overlap-aware text chunker, XML-safe boundaries
├── prompting/
│   ├── builder.py          # Prompt assembly
│   └── templates/
│       ├── text_gen.jinja2         # For plain text-generation endpoints
│       └── json_enforced.jinja2    # For JSON-mode / constrained endpoints
├── inference/
│   └── endpoint.py         # Endpoint wrapper + capability enum
├── postprocessing/
│   ├── resolver.py         # Context-anchor → char offset resolution
│   ├── validator.py        # Span verification + schema validation
│   └── injector.py         # Deterministic XML construction
└── pipeline.py             # Top-level orchestration
```

---

## Data Structures (`models/`)

```python
# schema.py
@dataclass
class TEIAttribute:
    name: str                        # e.g. "ref", "type", "cert"
    description: str
    required: bool = False
    allowed_values: list[str] | None = None   # None = free string

@dataclass
class TEIElement:
    tag: str                         # e.g. "persName"
    description: str                 # from TEI Guidelines
    allowed_children: list[str]      # tags of legal child elements
    attributes: list[TEIAttribute]

@dataclass
class TEISchema:
    elements: list[TEIElement]
    # convenience lookup
    def get(self, tag: str) -> TEIElement | None: ...

# spans.py
@dataclass
class SpanDescriptor:
    element: str
    text: str
    context: str                     # must contain text as substring
    attrs: dict[str, str]
    # always flat — nesting is inferred from offset containment in resolver/injector,
    # not emitted by the model (models produce unreliable nested trees)
    confidence: float | None = None  # passed through from GLiNER

@dataclass
class ResolvedSpan:
    element: str
    start: int
    end: int
    attrs: dict[str, str]
    children: list["ResolvedSpan"]
    fuzzy_match: bool = False        # flagged for human review
```

---

## GLiNER Dependency (`detection/gliner_detector.py`)

`gliner` is a regular package dependency declared in `pyproject.toml` and installed via `uv add gliner`. No manual setup step is needed.

Model weights are fetched from HuggingFace Hub automatically on first use of `GLiNER.from_pretrained(model_id)` and cached in `~/.cache/huggingface/`. If the import fails at runtime (e.g. the optional extra was not installed), the module raises a standard `ImportError` with a clear message — no wrapper needed.

Recommended models (specified as `gliner_model` parameter):

- `urchade/gliner_medium-v2.1` — balanced, Apache 2.0
- `numind/NuNER_Zero` — stronger multi-word entities, MIT (default)
- `knowledgator/gliner-multitask-large-v0.5` — adds relation extraction

All models run on CPU; no GPU required.

---

## Endpoint Abstraction (`inference/endpoint.py`)

```python
class EndpointCapability(Enum):
    TEXT_GENERATION = "text_generation"   # plain LLM, JSON via prompt only
    JSON_ENFORCED   = "json_enforced"     # constrained decoding guaranteed
    EXTRACTION      = "extraction"        # GLiNER2/NuExtract-style native

@dataclass
class EndpointConfig:
    capability: EndpointCapability
    call_fn: Callable[[str], str]
    # call_fn signature: takes a prompt string, returns a response string
    # caller is responsible for auth, model selection, retries
```

The `call_fn` injection means the library is agnostic about whether the caller is hitting Anthropic, OpenAI, a local Ollama instance, or Fastino's GLiNER2 API. The library just hands it a string and gets a string back.

---

## Pipeline (`pipeline.py`)

```python
@dataclass
class AnnotationResult:
    xml: str                          # annotated XML string
    fuzzy_spans: list[ResolvedSpan]   # spans flagged for human review

def annotate(
    text: str,                        # may contain existing XML tags
    schema: TEISchema,                # subset of TEI elements in scope
    endpoint: EndpointConfig,         # injected inference dependency
    gliner_model: str | None = "numind/NuNER_Zero",  # None disables GLiNER pass
    chunk_size: int = 1500,           # chars
    chunk_overlap: int = 200,
) -> AnnotationResult:
```

### Execution Flow

```
1. SETUP
   strip existing XML tags from text for processing,
   preserve them as a restoration map for final merge

2. GLINER PASS  (skipped if gliner_model=None, endpoint is EXTRACTION, or text is short)
   map TEISchema elements → flat label list for GLiNER
     e.g. [("persName", "a person's name"), ("placeName", "a place name"), ...]
   chunk text if len(text) > chunk_size (with overlap)
   run gliner.predict_entities() on each chunk
   merge cross-chunk duplicates by span text + context overlap
   output: list[SpanDescriptor] with text + context + element + confidence
   (GLiNER is a pre-filter only; the LLM may reject, correct, or extend its candidates)

3. PROMPT ASSEMBLY  
   select template based on EndpointCapability:
     TEXT_GENERATION:   include JSON structure example + "output only JSON" instruction
     JSON_ENFORCED:     minimal prompt, schema enforced externally
     EXTRACTION:        pass schema directly in endpoint's native format, skip LLM prompt
   inject into prompt:
     - TEIElement descriptions + allowed attributes for in-scope elements
     - GLiNER pre-detected spans as candidates for the model to enrich/correct
     - source text chunk
     - instruction to emit one SpanDescriptor per occurrence, not per unique entity

4. INFERENCE
   call endpoint.call_fn(prompt) → raw response string

5. POSTPROCESSING  (per chunk, then merged)

   a. Parse
      JSON_ENFORCED/EXTRACTION: parse directly
      TEXT_GENERATION: strip markdown fences, parse JSON,
                       on failure: retry once with a correction prompt that includes
                       the original (bad) response and the parse error message,
                       so the model can self-correct rather than starting from scratch

   b. Resolve  (resolver.py)
      for each SpanDescriptor:
        find context string in source → exact match preferred
        find text within context window
        assert source[start:end] == span.text → reject on mismatch
        fuzzy fallback (threshold 0.92) → flag for review

   c. Validate  (validator.py)
      reject spans where text not in source
      check attributes against TEISchema allowed values
      check element is in schema scope

   d. Inject  (injector.py)
      infer nesting from offset containment (child ⊂ parent by [start, end] bounds)
      check inferred nesting: children must be within parent bounds
      sort ResolvedSpans by start offset, handle nesting depth-first
      insert tags into a copy of the original source string
      restore previously existing XML tags from step 1

   e. Final validation
      parse output as XML → reject malformed documents
      optionally validate against full TEI RelaxNG schema via lxml

6. RETURN
   AnnotationResult(
     xml=annotated_xml_string,
     fuzzy_spans=list_of_flagged_resolved_spans,
   )
```

---

## Key Design Constraints

- The source text is **never modified by any model call**. All text in the output comes from the original input; models only contribute tag positions and attributes.
- The **GLiNER pass is optional** (`gliner_model=None` disables it). It is most useful for long texts with `TEXT_GENERATION` endpoints; it is skipped automatically for `EXTRACTION` endpoints or short inputs. When enabled, GLiNER is a pre-filter only — the LLM may reject, correct, or extend its candidates.
- **Span nesting is inferred from offsets**, never emitted by the model. `SpanDescriptor` is always flat; `ResolvedSpan.children` is populated by the injector from containment relationships.
- `call_fn` has **no required signature beyond `(str) -> str`**, making it trivial to swap endpoints, add logging, or inject mock functions for testing.
- Fuzzy-matched spans are **surfaced, not silently accepted** — `AnnotationResult.fuzzy_spans` provides a reviewable list alongside the XML.

---

## Testing Strategy

### Mocking philosophy

**Always mock `call_fn` and the GLiNER detector in unit tests.** Do not use a real GLiNER model as a substitute for a remote LLM endpoint — GLiNER is a span-labelling model that cannot produce JSON responses; it cannot exercise the parse/resolve/inject pipeline. Using a real model also makes tests slow (~seconds per inference on CPU), non-deterministic across versions and hardware, and dependent on a 400MB+ download.

The `call_fn: (str) -> str` design makes mocking trivial — a lambda returning a hardcoded JSON string is sufficient. No mock framework is needed.

### Test layers

**Layer 1 — Unit tests** (always run, <1s total, fully mocked):

```
tests/
├── test_chunker.py        # chunker unit tests
├── test_resolver.py       # resolver unit tests
├── test_validator.py      # validator unit tests
├── test_injector.py       # injector unit tests
├── test_builder.py        # prompt builder unit tests
├── test_parser.py         # JSON parse + retry unit tests
└── test_pipeline.py       # full pipeline smoke test (mocked call_fn + GLiNER)
```

**Layer 2 — Integration tests** (opt-in, gated by `pytest -m integration`):

```
tests/integration/
├── test_gliner_detector.py   # real GLiNER model, real HuggingFace download
└── test_pipeline_e2e.py      # full annotate() with real GLiNER + mocked call_fn
```

Integration tests are excluded from CI by default via `pyproject.toml`:

```toml
[tool.pytest.ini_options]
addopts = "-m 'not integration'"
```

### TDD cycle (red → green → refactor)

Each module is written test-first. Write a failing test, implement the minimum code to pass, refactor.

### Key test cases per module

**`chunker.py`**

- Short text below `chunk_size` → single chunk, offset 0
- Long text → multiple chunks with correct `start_offset` per chunk
- Span exactly at a chunk boundary → appears in both chunks with correct global offset
- Input with existing XML tags → chunk boundaries never split a tag

**`resolver.py`**

- Exact context match → `ResolvedSpan` with `fuzzy_match=False`
- Context not found in source → span rejected
- `source[start:end] != span.text` → span rejected
- Context found but span text not within it → span rejected
- Context found, span text found, score < 0.92 → span rejected
- Context found, span text found, 0.92 ≤ score < 1.0 → `fuzzy_match=True`
- Multiple occurrences of context → first match used, or rejection if ambiguous

**`validator.py`**

- Element not in schema → span rejected
- Attribute not in schema → span rejected
- Attribute value not in `allowed_values` → span rejected
- Valid span → passes through unchanged

**`injector.py`**

- Two non-overlapping spans → both tags inserted correctly
- Span B offset-contained in span A → B is child of A in output
- Overlapping (non-nesting) spans → reject or flatten with warning
- Restored XML tags from step 1 do not conflict with injected tags

**`builder.py`**

- `TEXT_GENERATION` capability → prompt contains JSON example and "output only JSON" instruction
- `JSON_ENFORCED` capability → prompt is minimal (no JSON scaffolding)
- GLiNER candidates present → candidates appear in prompt
- GLiNER candidates absent (pass skipped) → prompt has no candidate section

**`parser.py` (parse + retry logic)**

- Valid JSON response → parsed to `list[SpanDescriptor]` without retry
- Markdown-fenced JSON → fences stripped, parsed correctly
- Invalid JSON on first attempt → retry triggered with correction prompt that includes original bad response + parse error message
- Invalid JSON on second attempt → exception raised, chunk skipped

**`pipeline.py` (smoke test)**

```python
def mock_call_fn(prompt: str) -> str:
    return json.dumps([
        {"element": "persName", "text": "John Smith",
         "context": "...said John Smith yesterday...", "attrs": {}}
    ])

def test_annotate_smoke():
    schema = TEISchema(elements=[
        TEIElement(tag="persName", description="a person's name",
                   allowed_children=[], attributes=[])
    ])
    endpoint = EndpointConfig(
        capability=EndpointCapability.JSON_ENFORCED,
        call_fn=mock_call_fn,
    )
    result = annotate(
        text="He said John Smith yesterday.",
        schema=schema,
        endpoint=endpoint,
        gliner_model=None,   # disable GLiNER in unit tests
    )
    assert "persName" in result.xml
    assert "John Smith" in result.xml
    assert result.xml.count("John Smith") == 1   # text not duplicated
```

---

## Implementation Status

**Completed 2026-02-28** — full implementation per the plan above.

### What was built

All modules in the package structure were implemented:

| File | Notes |
| --- | --- |
| `tei_annotator/models/schema.py` | `TEIAttribute`, `TEIElement`, `TEISchema` dataclasses |
| `tei_annotator/models/spans.py` | `SpanDescriptor`, `ResolvedSpan` dataclasses |
| `tei_annotator/inference/endpoint.py` | `EndpointCapability` enum, `EndpointConfig` dataclass |
| `tei_annotator/chunking/chunker.py` | `chunk_text()` — overlap chunker, XML-safe boundaries |
| `tei_annotator/detection/gliner_detector.py` | `detect_spans()` — optional, raises `ImportError` if `[gliner]` extra not installed |
| `tei_annotator/prompting/builder.py` | `build_prompt()` + `make_correction_prompt()` |
| `tei_annotator/prompting/templates/text_gen.jinja2` | Verbose prompt with JSON example, "output only JSON" instruction |
| `tei_annotator/prompting/templates/json_enforced.jinja2` | Minimal prompt for constrained-decoding endpoints |
| `tei_annotator/postprocessing/parser.py` | `parse_response()` — fence stripping, one-shot self-correction retry |
| `tei_annotator/postprocessing/resolver.py` | `resolve_spans()` — context-anchor → char offset, rapidfuzz fuzzy fallback at threshold 0.92 |
| `tei_annotator/postprocessing/validator.py` | `validate_spans()` — element, attribute name, allowed-value checks |
| `tei_annotator/postprocessing/injector.py` | `inject_xml()` — stack-based nesting tree, recursive tag insertion |
| `tei_annotator/pipeline.py` | `annotate()` — full orchestration, tag strip/restore, deduplication across chunks, lxml final validation |

### Dependencies added

Runtime: `jinja2`, `lxml`, `rapidfuzz`. Optional extra `[gliner]` for GLiNER support. Dev: `pytest`, `pytest-cov`.

### Tests

- **63 unit tests** (Layer 1) — fully mocked, run in < 0.1 s via `uv run pytest`
- **9 integration tests** (Layer 2, no GLiNER) — complex resolver/injector/pipeline scenarios, run via `uv run pytest --override-ini="addopts=" -m integration tests/integration/test_pipeline_e2e.py -k "not real_gliner"`
- **1 GLiNER integration test** — requires `[gliner]` extra and HuggingFace model download

### Smoke script

`scripts/smoke_test_llm.py` — end-to-end test with real LLM calls (no GLiNER). Verified against:

- **Google Gemini 2.0 Flash** (`GEMINI_API_KEY` from `.env`)
- **KISSKI `llama-3.3-70b-instruct`** (`KISSKI_API_KEY` from `.env`, OpenAI-compatible API at `https://chat-ai.academiccloud.de/v1`)

Run with `uv run scripts/smoke_test_llm.py`.

### Key implementation notes

- The `_strip_existing_tags` / `_restore_existing_tags` pair in `pipeline.py` preserves original markup by tracking plain-text offsets of each stripped tag and re-inserting them after annotation.
- `_build_nesting_tree` in `injector.py` uses a sort-by-(start-asc, length-desc) + stack algorithm; partial overlaps are dropped with a `warnings.warn`.
- The resolver does an exact `str.find` first; fuzzy search (sliding-window rapidfuzz) is only attempted if exact fails and rapidfuzz is installed.
- `parse_response` passes `call_fn` and `make_correction_prompt` only for `TEXT_GENERATION` endpoints; `JSON_ENFORCED` and `EXTRACTION` never retry.

---

## RNG Schema Parsing (`tei_annotator/tei.py`)

**Added 2026-02-28** — `create_schema()` parses a RELAX NG schema file and produces a `TEISchema` for use with `annotate()`.

### Function signature

```python
def create_schema(
    schema_path: str | Path,
    element: str = "text",
    depth: int = 1,
) -> TEISchema
```

- `schema_path` — path to a `.rng` file (e.g. `schema/tei-bib.rng`).
- `element` — name of the root TEI element to start from (default: `"text"`).
- `depth` — how many levels of descendant elements to include. `depth=1` adds the root element **and** its direct children; `depth=2` adds grandchildren too; `depth=0` adds only the root itself.
- Raises `ValueError` if the element is not found in the schema.

### How it works

1. Parses the `.rng` with lxml, builds a `{define_name: element_node}` lookup table.
2. Builds a reverse map from TEI element names (e.g. `"persName"`) to their RNG define names (e.g. `"tbibpersName"`).
3. BFS from the requested element, collecting `TEIElement` entries level by level up to `depth`.
4. For each element:
   - **description** — extracted from the first `<a:documentation>` child inside the RNG `<element>` node.
   - **`allowed_children`** — content-model `<ref>` nodes are expanded recursively through macro/model groups; attribute-group refs (names containing `"att."`) are skipped; element-bearing defines are short-circuited (just record the element name, don't recurse — correctly handles self-referential elements like `idno`).
   - **`attributes`** — attribute-group refs are followed recursively; inline `<attribute>` elements are also collected; `required` is inferred from the immediate parent (`<optional>` / `<zeroOrMore>` → False, otherwise True); `allowed_values` comes from `<choice><value>` enumeration inside the attribute definition.
5. Deduplicates children and attributes (preserving order) before constructing each `TEIElement`.

### Tests (`tests/test_tei.py`)

15 unit tests — run in < 0.05 s, no network, no model downloads:

- `idno` (leaf-like, self-referential, enumerated `type` values: `ISBN`, `ISSN`, `DOI`, …)
- `biblStruct` (explicit named children: `analytic`, `monogr`, `series`, `relatedItem`, `citedRange`; plus model-group expansion: `note`, `ptr`)
- `depth=0` vs `depth=1` behaviour
- Duplicate element / attribute detection
- Unknown element raises `ValueError`

Total unit tests after this addition: **78** (63 original + 15 new).

---

## Evaluation Module (`tei_annotator/evaluation/`)

**Added 2026-02-28** — compares annotator output against a gold-standard TEI XML file to compute precision, recall, and F1 score.

### Package structure

```
tei_annotator/evaluation/
├── __init__.py        # public API exports
├── extractor.py       # EvaluationSpan, extract_spans(), spans_from_xml_string()
├── metrics.py         # MatchMode, SpanMatch, ElementMetrics, EvaluationResult,
│                      # match_spans(), compute_metrics(), aggregate()
└── evaluator.py       # evaluate_bibl(), evaluate_file()
```

### Algorithm

For each gold-standard element (e.g. `<bibl>`):

1. **Extract gold spans** — walk the element tree, record `(tag, start, end, text)` for every descendant element using absolute char offsets in the element's plain text (`"".join(element.itertext())`).
2. **Strip tags** — the same plain text is passed to `annotate()`.
3. **Annotate** — run the full pipeline with the injected `EndpointConfig`.
4. **Extract predicted spans** — parse `result.xml` as `<_root>…</_root>`, then run the same span extractor.
5. **Match** — greedily pair (gold, pred) spans by score (highest first); each span matched at most once.
6. **Compute metrics** — count TP/FP/FN per element type; derive precision, recall, F1.

### Match modes (`MatchMode`)

| Mode | A match if… |
| --- | --- |
| `TEXT` (default) | same element tag + normalised text content |
| `EXACT` | same element tag + identical `(start, end)` offsets |
| `OVERLAP` | same element tag + IoU ≥ `overlap_threshold` (default 0.5) |

### Key entry points

```python
# Single element evaluation (lxml element)
result = evaluate_bibl(gold_element, schema, endpoint, gliner_model=None)
print(result.report())

# Full file evaluation
per_record, overall = evaluate_file(
    "tests/fixtures/blbl-examples.tei.xml",
    schema=schema,
    endpoint=endpoint,
    max_items=10,   # optional — first N bibls only
)
print(overall.report())
```

`EvaluationResult` exposes:

- `micro_precision / micro_recall / micro_f1` — aggregate counts, then compute rates
- `macro_precision / macro_recall / macro_f1` — average per-element rates
- `per_element: dict[str, ElementMetrics]` — per-element breakdown
- `matched / unmatched_gold / unmatched_pred` — full span lists for inspection
- `report()` — human-readable summary string

### Evaluation Tests

39 new unit tests in `tests/test_evaluation.py` — all mocked, run in < 0.1 s.

Total unit tests: **117** (78 + 39).