# Label-Driven RAG Evaluation Pipeline

## 1. Objective

Build a separate Python project that uses `labelgen` as its paragraph labeling
backend and exposes a clean RAG-oriented pipeline for evaluation workflows.

This project is responsible for:

1. accepting a paragraph corpus that has already been segmented upstream
2. building and storing paragraph-side label metadata
3. accepting a query string
4. extracting query-side concepts and label IDs with the same logic
5. retrieving paragraphs by greedy query-label coverage
6. producing a generation-ready request package and a structured response record

This project is not responsible for:

- loading raw datasets from specific sources
- paragraph segmentation
- benchmark orchestration
- metric computation
- experiment dashboards

The deliverable for the first milestone is an implementation-ready design
document, not production code.

## 2. Relationship To `labelgen`

The current repository already defines the core labeling primitives:

- paragraph normalization and stable paragraph IDs in [src/labelgen/preprocessing/paragraphs.py](/Users/huruilizhen/Repo/labelgen/src/labelgen/preprocessing/paragraphs.py)
- typed result models in [src/labelgen/types.py](/Users/huruilizhen/Repo/labelgen/src/labelgen/types.py)
- stateful fitting and transformation in [src/labelgen/pipeline/label_generator.py](/Users/huruilizhen/Repo/labelgen/src/labelgen/pipeline/label_generator.py)
- JSON serialization helpers in [src/labelgen/io/serialize.py](/Users/huruilizhen/Repo/labelgen/src/labelgen/io/serialize.py)

The new RAG project should treat `labelgen` as a dependency and should not
duplicate concept extraction, concept filtering, community detection, or label
assignment logic.

Recommended dependency direction:

- `rag-eval-pipeline` depends on `paralabelgen`
- `paralabelgen` remains retrieval-agnostic

## 3. Scope

### In scope

- corpus initialization from `list[str]` or `list[Paragraph]`
- support for serialized `Paragraph` payloads reconstructed by caller code
- paragraph-side label metadata indexing
- query-side concept and label extraction
- greedy retrieval over query label IDs
- prompt-ready context assembly
- structured output models for downstream generation and evaluation
- local serialization of the pipeline state and retrieval traces
- tests and code quality tooling

### Out of scope

- dataset-specific adapters
- model serving infrastructure
- async job scheduling
- evaluation metrics
- judge prompts
- UI, demo app, notebook workflow
- hybrid dense retrieval unless explicitly added later

## 4. Key Design Principles

### 4.1 Keep The Dependency Boundary Clean

`labelgen` remains the only component that knows how to extract concepts and
derive communities and labels. The RAG pipeline consumes its outputs and adds
retrieval and generation orchestration only.

### 4.2 Preserve Interpretability

Every returned answer record must make it possible to inspect:

- query concepts
- query label IDs
- which paragraphs were retrieved
- which query label IDs each retrieved paragraph covered
- what context was finally sent to the generator

### 4.3 Deterministic Retrieval

Given the same fitted pipeline state, query text, and retrieval config, the
retrieval result should be stable. Tie-breaking rules must be explicit.

### 4.4 Offline-First Evaluation Support

The pipeline should work in local evaluation loops. External generation calls
should be injectable rather than hard-wired into the retrieval stack.

## 5. Recommended Project Layout

Use a new repository or top-level package with a structure parallel to the
current library.

```text
rag_eval_pipeline/
  pyproject.toml
  README.md
  docs/
    implementation.md
    public_api.md
  src/
    rag_eval_pipeline/
      __init__.py
      config.py
      types.py
      io/
        __init__.py
        serialize.py
      indexing/
        __init__.py
        corpus_index.py
      retrieval/
        __init__.py
        selector.py
        coverage.py
      generation/
        __init__.py
        prompt_builder.py
        generator.py
      pipeline/
        __init__.py
        rag_pipeline.py
  tests/
    test_pipeline.py
    test_retrieval.py
    test_serialization.py
    test_prompt_builder.py
```

Recommended import boundary:

- `indexing/` builds paragraph-side metadata tables from `labelgen`
- `retrieval/` implements greedy query-label coverage
- `generation/` formats prompts and optionally calls an injected answer model
- `pipeline/` exposes the main public entrypoint

## 6. Public API Proposal

### Primary entrypoint

```python
from rag_eval_pipeline import RAGPipeline, RAGPipelineConfig

pipeline = RAGPipeline(config)
pipeline.fit(paragraphs)
result = pipeline.answer(question)
```

### Required methods

```python
fit(paragraphs: list[str] | list[Paragraph]) -> RAGPipeline
answer(question: str) -> RAGAnswerResult
save(path: str | Path) -> None
load(path: str | Path) -> RAGPipeline
```

### Optional convenience methods

```python
build_context(question: str) -> RetrievalResult
answer_with_generator(question: str, generator: AnswerGenerator) -> RAGAnswerResult
```

Recommended API split:

- `fit(...)` prepares corpus-side label metadata and retrieval index
- `build_context(...)` runs only query analysis and retrieval
- `answer(...)` may either require an injected generator at construction time or
  return a result with empty `answer_text` when no generator is configured

## 7. Core Data Contracts

Prefer `dataclass(slots=True)` to stay aligned with `labelgen`.

### 7.1 Corpus-side models

```python
@dataclass(slots=True)
class IndexedParagraph:
    paragraph_id: str
    text: str
    metadata: dict[str, Any] | None
    concept_ids: list[str]
    concept_texts: list[str]
    label_ids: list[str]
    label_display_names: list[str]
```

`paragraph_id`, `text`, and `metadata` should come from normalized
`labelgen.Paragraph` objects. `concept_ids` and `label_ids` should come from the
fitted `LabelGenerationResult`.

### 7.2 Query analysis model

```python
@dataclass(slots=True)
class QueryAnalysis:
    query_text: str
    concepts: list[str]
    concept_ids: list[str]
    label_ids: list[str]
    label_display_names: list[str]
```

The query should go through the same concept extraction and label assignment
path as paragraphs. The simplest design is to reuse a fitted `LabelGenerator`
and call `transform([query])`.

### 7.3 Retrieval trace model

```python
@dataclass(slots=True)
class RetrievedParagraph:
    paragraph_id: str
    text: str
    metadata: dict[str, Any] | None
    matched_label_ids: list[str]
    matched_concept_ids: list[str]
    paragraph_label_ids: list[str]
    paragraph_concept_ids: list[str]
    marginal_gain: int
    retrieval_score: float
```

### 7.4 Final answer model

```python
@dataclass(slots=True)
class RAGAnswerResult:
    question: str
    answer_text: str
    query_analysis: QueryAnalysis
    retrieved_paragraphs: list[RetrievedParagraph]
    prompt_context: str
    metadata: dict[str, Any]
```

Minimum required metadata keys:

- `covered_label_ids`
- `uncovered_label_ids`
- `retrieval_strategy`
- `generator_name`
- `generation_model`

## 8. Fit-Time Workflow

`fit(paragraphs)` should execute the following steps:

1. normalize `list[str] | list[Paragraph]` into `Paragraph` objects by relying
   on `labelgen` input support
2. call `LabelGenerator.fit_transform(paragraphs)`
3. keep the fitted `LabelGenerator` instance in memory
4. build an index keyed by `paragraph_id`
5. build lookup tables:
   - `label_id -> set[paragraph_id]`
   - `paragraph_id -> set[label_id]`
   - `paragraph_id -> set[concept_id]`
   - `label_id -> display_name`
6. persist enough information for later `answer()` and `save()`

The new project should not reimplement paragraph normalization or paragraph ID
generation. It should trust the normalized outputs of `labelgen`.

## 9. Query-Time Workflow

`answer(question)` should execute the following steps:

1. validate that the pipeline is fitted
2. call `LabelGenerator.transform([question])` or an equivalent query-analysis
   wrapper
3. extract query-side concepts and label IDs
4. run greedy retrieval over the query label IDs
5. assemble the retrieved paragraphs into prompt context
6. call an injected generator if configured
7. return a structured `RAGAnswerResult`

Important implementation note:

The query is not part of the training corpus. The current `labelgen`
implementation already supports `transform()` against fitted communities, which
makes it the correct entrypoint for query analysis.

## 10. Retrieval Algorithm Proposal

### 10.1 Problem framing

Given:

- a query label set `Q`
- a set of indexed paragraphs `P`
- each paragraph `p` with label set `L(p)`

Select an ordered paragraph list that greedily covers `Q`.

### 10.2 Baseline greedy algorithm

Initialize:

- `remaining = set(query_label_ids)`
- `selected = []`

Iterate until termination:

1. for each unselected paragraph, compute `gain = len(L(p) ∩ remaining)`
2. discard paragraphs with `gain == 0`
3. select the paragraph with highest `gain`
4. tie-break in this order:
   - larger total overlap on query concept IDs if available
   - larger `len(L(p))`
   - lexicographically smaller `paragraph_id`
5. add paragraph to `selected`
6. remove covered labels from `remaining`

Termination conditions:

- `remaining` becomes empty
- `max_paragraphs` reached
- no paragraph has positive gain

### 10.3 Why this baseline is acceptable

- easy to reason about
- stable and testable
- aligned with your current description
- good enough for a first retrieval baseline in an evaluation pipeline

### 10.4 Required trace outputs

For each selected paragraph, record:

- labels newly covered at this step
- labels already covered before this step
- marginal gain
- optional concept overlap count

### 10.5 Known limitations

- may ignore paragraph ordering in original documents
- may over-select paragraphs with noisy broad labels
- does not optimize answer utility directly
- does not use dense semantic similarity

These are acceptable for v0.

## 11. Context Construction

After retrieval, build a plain-text context block for answer generation.

Recommended default format:

```text
[Paragraph 1 | id=<paragraph_id>]
<paragraph text>

[Paragraph 2 | id=<paragraph_id>]
<paragraph text>
```

Keep prompt construction in a dedicated module so the evaluation team can swap
templates without touching retrieval logic.

## 12. Generator Boundary

The project should define an abstraction instead of hard-coding one model API.

```python
class AnswerGenerator(Protocol):
    def generate(self, question: str, context: str) -> GeneratedAnswer: ...
```

Recommended generated answer contract:

```python
@dataclass(slots=True)
class GeneratedAnswer:
    text: str
    metadata: dict[str, Any]
```

This keeps model calling outside the core retrieval logic and matches the stated
boundary that the team only needs to hand downstream evaluation code the answer
text plus retrieval metadata.

## 13. Serialization Requirements

Support JSON serialization for:

- pipeline config
- fitted label generator state
- retrieval index artifacts
- answer result records

Recommended persistence approach:

1. reuse `LabelGenerator.save(...)` for the underlying fitted label model
2. store RAG-pipeline-specific artifacts in adjacent JSON files
3. keep paths explicit and human-inspectable

Suggested on-disk layout:

```text
pipeline_dir/
  config.json
  label_generator.json
  corpus_index.json
```

## 14. Configuration Proposal

```python
@dataclass(slots=True)
class RetrievalConfig:
    max_paragraphs: int = 8
    require_full_label_coverage: bool = False
    allow_label_free_fallback: bool = True

@dataclass(slots=True)
class PromptConfig:
    include_paragraph_ids: bool = True
    include_label_annotations: bool = False
    max_context_characters: int | None = None

@dataclass(slots=True)
class RAGPipelineConfig:
    labelgen: LabelGeneratorConfig = field(default_factory=LabelGeneratorConfig)
    retrieval: RetrievalConfig = field(default_factory=RetrievalConfig)
    prompt: PromptConfig = field(default_factory=PromptConfig)
```

One important boundary rule:

The new project may configure `LabelGeneratorConfig`, but it must not patch
`labelgen` internals directly.

## 15. Quality Gates

Mirror the current repository style in [pyproject.toml](/Users/huruilizhen/Repo/labelgen/pyproject.toml).

Minimum toolchain:

- `ruff` for lint and formatting
- `pyright` in `strict` mode
- `pytest`
- `pytest-cov` optional but recommended

Recommended `pyproject.toml` baseline:

```toml
[tool.pytest.ini_options]
minversion = "8.0"
testpaths = ["tests"]
addopts = [
  "--strict-config",
  "--strict-markers",
]

[tool.ruff]
target-version = "py311"
line-length = 100
src = ["src", "tests"]

[tool.ruff.lint]
select = ["E", "F", "I", "B", "UP"]

[tool.ruff.format]
quote-style = "double"
indent-style = "space"

[tool.pyright]
include = ["src", "tests"]
pythonVersion = "3.11"
typeCheckingMode = "strict"
venvPath = "."
venv = ".venv"
```

## 16. Test Plan

### Unit tests

- fit supports both `list[str]` and `list[Paragraph]`
- serialized `Paragraph` payloads can be reconstructed by caller and accepted
- query analysis uses fitted label communities
- greedy retrieval covers all reachable query labels
- tie-breaking is deterministic
- empty query label set triggers the configured fallback path
- save/load preserves retrieval behavior

### Integration tests

- end-to-end fit -> answer with a stub generator
- end-to-end fit -> build_context without generator
- answer result includes required metadata fields

### Non-goals for tests

- external dataset loading
- external benchmark scripts
- live model provider reliability

## 17. Capability Boundaries

This project owns:

- indexing
- retrieval
- prompt context construction
- structured answer packaging

This project does not own:

- dataset adapters
- evaluation metrics
- benchmark runners
- label extraction internals

Any feature request that changes concept extraction, community detection, or
paragraph label assignment should normally be implemented in `labelgen`, not in
this RAG pipeline.

## 18. Open Decisions To Confirm

The following points are still under-specified and should be confirmed before
implementation begins.

### 18.1 Query-side concepts field

Question:
Should `query_analysis.concepts` return normalized concept texts only, or both
surface forms and normalized forms?

Recommendation:
Return normalized concept texts in the top-level summary field and keep full raw
mention details in `metadata` only if needed later.

### 18.2 Query-side label semantics

Question:
Should query-side labels be represented as the flat `label_ids` assigned by
`labelgen.ParagraphLabels`, or as a higher-order set-of-sets representation?

Recommendation:
For v0, define query-side labels as the flat `label_ids` assigned to the query
by `LabelGenerator.transform([query])`. Do not introduce second-order label-set
structures until they are required by a concrete retrieval need.

### 18.3 Fallback when the query has no labels

Question:
If query analysis yields zero labels, should retrieval return no paragraphs, or
fall back to another strategy?

Recommendation:
Support a config flag and default to a simple fallback. The safest v0 fallback is
to return an empty retrieval result and leave richer fallback retrieval for a
later milestone unless you already know another deterministic rule you want.

### 18.4 Generator placement

Question:
Should `answer()` always call an injected generator, or should the core pipeline
only build retrieval context and let another layer call the LLM?

Recommendation:
Support both by keeping `build_context()` and `answer_with_generator()` separate.
That gives your evaluation teammate a stable seam.

### 18.5 Retrieval target

Question:
Do you want greedy coverage over:

- only query label IDs
- query label IDs plus evidence concept IDs
- only evidence concept IDs

Recommendation:
Start with query label IDs only. Use concept overlap strictly as a tie-breaker.
This matches the current description and avoids overfitting to extractor noise.

## 19. Recommended Implementation Order

1. define config and result dataclasses
2. implement fit-time corpus indexing on top of `labelgen`
3. implement deterministic greedy retrieval
4. implement prompt context builder
5. implement serialization
6. add stub-generator integration
7. add tests and quality gates

## 20. Immediate Next-Step Questions

Before coding, I recommend locking down these three decisions:

1. Is this RAG pipeline a standalone package/repository, or a new subpackage in
   the current repo?
2. For v0, should query-side labels be treated as flat `query_label_ids`, as
   proposed above?
3. When query labels cannot fully cover any paragraph set, should retrieval stop
   with partial coverage, or should there be a non-label fallback path?
