Categorization stage (two-pass + single-pass)

First real implementation of the categorization pipeline stage. Ports two of the three legacy Gemini-via-OpenRouter strategies. Introduces the backend's first LLM integration, a per-item DB table, DB-backed categorizer settings, and a Storage abstraction with a local-disk implementation.

TL;DRSummary

§1Context

What exists

• features/pipeline/stages.py:25 — categorization_stage is a return None stub. STAGES/NEXT_STAGE already route SEGMENTATION → CATEGORIZATION → APPRAISAL.
• features/segmentation/ — precedent we mirror. step.py entry, config.py with kind enum + build_segmenter() factory, settings.py over a singleton DB row.
• features/pipeline/models.py — PipelineRun audit row. No token-usage columns today.
• core/settings.py:19 — data_dir: Path = Path("data"). The only storage primitive in the backend; segmentation pokes pathlib.Path directly.
• features/deal_images/ — stores URLs of source photos, not bytes on disk. Not relevant to crop storage.
• No backend code imports any LLM SDK; httpx>=0.27 is the only HTTP dependency.

Legacy categorization (the source we are porting)

• checkin-pipeline/app/steps/classify/two_pass/ — two Gemini vision calls (category → category-scoped subcategory). Legacy default.
• checkin-pipeline/app/steps/classify/single_pass/ — one Gemini vision JSON call returning both fields.
• checkin-pipeline/app/steps/classify/describe_classify/ — vision-describe + text-classify. Not ported in this slice.
• _shared.py — CLASSIFICATION_SCHEMA, resolve_subcategory(), format_size_suffix(), strip_code_fences().
• deal_step.py — orchestrator: registry-based variation loader, per-deal Semaphore(10), token accumulation on pipeline_runs.
• app/gemini/client.py — OpenRouter HTTP client with three wrappers (vision_request, vision_json_request, text_json_request).
• config/categories.toml (11 categories) + subcategories.csv (per-category list with hint metadata).
• app/steps/_prompts.py reads prompt .md files on every call (no caching); injects {categories} and {subcategory_tree}.

§2Scope (in / out)

§3Architecture

Directory layout

backend/src/vcc_backend/
├── core/
│   └── storage.py                   # NEW — Storage Protocol + LocalFsStorage
├── api/
│   └── routers/
│       └── files.py                 # NEW — GET /files/{key:path}
├── integrations/
│   └── openrouter/                  # NEW — first LLM integration
│       ├── __init__.py
│       ├── client.py                # lifted from checkin-pipeline/app/gemini/client.py
│       └── schemas.py               # ChatResponse, JsonResponse, TokenUsage
├── features/
│   ├── segmentation/
│   │   └── step.py                  # EXTENDED — storage I/O + DealItem insert
│   ├── categorization/              # NEW
│   │   ├── __init__.py
│   │   ├── models.py                # DealItem, CategorizerSettings ORM
│   │   ├── schemas.py               # pydantic Settings types
│   │   ├── settings.py              # get_categorizer_settings(session) + upsert
│   │   ├── categorizer.py           # Categorizer Protocol, dataclasses
│   │   ├── two_pass.py              # TwoPassCategorizer
│   │   ├── single_pass.py           # SinglePassCategorizer
│   │   ├── factory.py               # build_categorizer(settings, http)
│   │   ├── taxonomy.py              # loads TOML+CSV; resolve_subcategory()
│   │   ├── step.py                  # run_categorization_stage(session, deal)
│   │   ├── config/
│   │   │   ├── categories.toml
│   │   │   └── subcategories.csv
│   │   └── prompts/
│   │       ├── two_pass.category.md
│   │       ├── two_pass.subcategory.md
│   │       └── single_pass.md
│   └── pipeline/
│       └── stages.py                # EDITED — wire run_categorization_stage

Sequence

§4Data model

New table — deal_items

Lives in features/categorization/models.py. Segmentation is the row creator but does not own the table — same as how PipelineRun lives in features/pipeline/ even though other features write to it.

New table — categorizer_settings (singleton)

Mirrors segmenter_settings. Single-row table; pydantic-serialized settings JSONB plus updated_at / updated_by.

class TwoPassConfig(BaseModel):
    model: str = "google/gemini-2.5-flash"
    timeout_seconds: float = 30.0
    temperature: float = 0.0

class SinglePassConfig(BaseModel):
    model: str = "google/gemini-2.5-flash"
    timeout_seconds: float = 30.0
    temperature: float = 0.0

class CategorizerSettings(BaseModel):
    kind: Literal["two_pass", "single_pass"] = "two_pass"
    concurrency: int = 10
    two_pass: TwoPassConfig = Field(default_factory=TwoPassConfig)
    single_pass: SinglePassConfig | None = None

Modified table — pipeline_runs

All nullable; older stages and old rows stay NULL. No backfill.

Migration 0016_categorization.py

One Alembic revision. Up creates deal_items, categorizer_settings, and the three columns on pipeline_runs. Down drops in reverse. Append deal_items, categorizer_settings to TRUNCATE_TABLES in backend/tests/conftest.py.

§5Storage abstraction

Storage protocol

# core/storage.py
class Storage(Protocol):
    async def read_bytes(self, key: str) -> bytes: ...
    async def write_bytes(self, key: str, data: bytes) -> None: ...
    async def exists(self, key: str) -> bool: ...
    async def delete_prefix(self, prefix: str) -> None: ...
    def signed_url(self, key: str, ttl_seconds: int = 3600) -> str: ...

Keys are POSIX-style strings (e.g. deals/{deal_id}/crops/IMG_0421/item_03.png). Never contain backslashes; never start with /. The protocol is intentionally minimal — no streaming, no metadata. Add when needed.

LocalFsStorage(root: Path)

• root is Settings.data_dir (default Path("data")).
• read_bytes / write_bytes via aiofiles; parent dirs created on write.
• delete_prefix recursively removes root / prefix; safe against .. escape (validated).
• signed_url(key) returns f"/files/{quote(key)}". The ttl_seconds argument is accepted but ignored (documented).

GET /files/{key:path} proxy route

Resolves the key under data_dir, validates it stays under data_dir, streams the file with sniffed Content-Type. 404 for missing keys, 400 for invalid keys. Documented as dev/local only — in production with Azure, the proxy is unmounted and the Azure backend issues real signed URLs.

Selection

Settings.storage_kind: Literal["local"] = "local" (env: VCC_STORAGE_KIND). A second pydantic-settings field group per kind (only local today). core/storage.py exposes get_storage() — lru_cached, mirrors get_db_engine.

§6OpenRouter client (integrations/openrouter/)

Lift

Copy checkin-pipeline/app/gemini/client.py to integrations/openrouter/client.py. Rename methods onto a class:

• gemini_vision_request → OpenRouterClient.vision_request(image, system, model, …) → ChatResponse
• gemini_vision_json_request → OpenRouterClient.vision_json_request(image, system, schema, model, …) → JsonResponse
• gemini_text_json_request → OpenRouterClient.text_json_request(prompt, schema, model, …) → JsonResponse

Class takes an httpx.AsyncClient and API key in __init__; instances constructed per-stage (cheap, no global state).

Typed errors

# integrations/openrouter/client.py
class OpenRouterError(Exception):
    """Base. Carries status + first 500 bytes of body."""
class OpenRouterAuthError(OpenRouterError):    pass   # 401, 403
class OpenRouterQuotaError(OpenRouterError):   pass   # 429, 402
class OpenRouterServerError(OpenRouterError):  pass   # 5xx
class OpenRouterParseError(OpenRouterError):   pass   # JSON-mode response didn't parse

The categorization stage uses these to distinguish per-item failures (caught, item → failed) from stage-level failures (auth/quota, re-raised past the per-item try/except).

Schemas

# integrations/openrouter/schemas.py
class TokenUsage(BaseModel):
    input_tokens: int
    output_tokens: int

class ChatResponse(BaseModel):
    text: str
    usage: TokenUsage
    model: str

class JsonResponse(BaseModel):
    data: dict[str, Any]      # already JSON-parsed; raises on bad JSON
    usage: TokenUsage
    model: str

Settings additions (core/settings.py)

openrouter_api_key:        SecretStr | None = None
openrouter_base_url:       str = "https://openrouter.ai/api/v1"
openrouter_default_model:  str = "google/gemini-2.5-flash"
storage_kind:              Literal["local"] = "local"

Worker startup adds VCC_OPENROUTER_API_KEY to the existing required-env validation block (next to HUBSPOT_ACCESS_TOKEN).

§7Categorizer strategies

Categorizer protocol

# features/categorization/categorizer.py
class SizeHint(BaseModel):
    width_cm:  float
    height_cm: float

class ItemClassification(BaseModel):
    category:               str | None
    category_confidence:    float | None
    subcategory:            str | None
    subcategory_id:         str | None
    subcategory_confidence: float | None

class Categorizer(Protocol):
    async def classify(
        self, *, image_bytes: bytes, size_hint: SizeHint | None
    ) -> tuple[ItemClassification, TokenUsage]: ...

TwoPassCategorizer

class TwoPassCategorizer:
    def __init__(self, client, taxonomy, cfg): ...

    async def classify(self, *, image_bytes, size_hint):
        # 1. category call
        sys1 = render_prompt("two_pass.category.md",
            categories=self.taxonomy.categories_block(),
            size_suffix=format_size_suffix(size_hint))
        r1 = await self.client.vision_json_request(image_bytes, sys1, CATEGORY_SCHEMA,
            model=self.cfg.model, timeout=self.cfg.timeout_seconds, temperature=self.cfg.temperature)
        category = r1.data["category"]; cat_conf = r1.data["confidence"]

        # 2. subcategory call (scoped to chosen category)
        sys2 = render_prompt("two_pass.subcategory.md",
            subcategory_block=self.taxonomy.subcategory_block(category),
            category=category, size_suffix=format_size_suffix(size_hint))
        r2 = await self.client.vision_json_request(image_bytes, sys2, SUBCATEGORY_SCHEMA,
            model=self.cfg.model, timeout=self.cfg.timeout_seconds, temperature=self.cfg.temperature)
        subcat_label = r2.data["subcategory"]; sub_conf = r2.data["confidence"]
        subcat_id = self.taxonomy.resolve_subcategory(category, subcat_label)

        return (ItemClassification(category=category, category_confidence=cat_conf,
                                   subcategory=subcat_label, subcategory_id=subcat_id,
                                   subcategory_confidence=sub_conf),
                TokenUsage(input_tokens =r1.usage.input_tokens  + r2.usage.input_tokens,
                           output_tokens=r1.usage.output_tokens + r2.usage.output_tokens))

SinglePassCategorizer

class SinglePassCategorizer:
    async def classify(self, *, image_bytes, size_hint):
        sys = render_prompt("single_pass.md",
            tree=self.taxonomy.tree_block(), size_suffix=format_size_suffix(size_hint))
        r = await self.client.vision_json_request(image_bytes, sys, SINGLE_PASS_SCHEMA,
            model=self.cfg.model, timeout=self.cfg.timeout_seconds, temperature=self.cfg.temperature)
        cat = r.data["category"]; subcat_label = r.data["subcategory"]; conf = r.data["confidence"]
        subcat_id = self.taxonomy.resolve_subcategory(cat, subcat_label)
        # single-pass returns one confidence; copy to both fields for shape parity
        return (ItemClassification(category=cat, category_confidence=conf,
                                   subcategory=subcat_label, subcategory_id=subcat_id,
                                   subcategory_confidence=conf),
                r.usage)

build_categorizer(settings, http)

# features/categorization/factory.py
async def build_categorizer(settings: CategorizerSettings, http: httpx.AsyncClient) -> Categorizer:
    cfg = get_settings()
    if not cfg.openrouter_api_key:
        raise RuntimeError("VCC_OPENROUTER_API_KEY not configured")
    client = OpenRouterClient(http, api_key=cfg.openrouter_api_key.get_secret_value(),
                              base_url=cfg.openrouter_base_url)
    taxonomy = Taxonomy.load()  # reads TOML+CSV each call
    if settings.kind == "two_pass":
        return TwoPassCategorizer(client, taxonomy, settings.two_pass)
    if settings.kind == "single_pass":
        return SinglePassCategorizer(client, taxonomy, settings.single_pass or SinglePassConfig())
    raise ValueError(f"unknown categorizer kind: {settings.kind}")

Taxonomy loader

• Taxonomy.load() reads config/categories.toml + config/subcategories.csv from disk every call (no cache — matches legacy).
• categories_block() renders category names + descriptions.
• subcategory_block(category) renders category-scoped subcategory list with hint metadata.
• tree_block() renders the full category→subcategory tree.
• resolve_subcategory(category, label) case-insensitive substring match, exact preferred; returns canonical CSV id or None.

Prompts

Copied from checkin-pipeline/app/steps/classify/{two_pass,single_pass}/*.md, placeholder names normalized:

• prompts/two_pass.category.md — {categories}, {size_suffix}
• prompts/two_pass.subcategory.md — {category}, {subcategory_block}, {size_suffix}
• prompts/single_pass.md — {tree}, {size_suffix}

render_prompt(name, **kwargs) does Template(text).safe_substitute(**kwargs) against the file contents — re-read each call.

§8Stage entry (features/categorization/step.py)

async def run_categorization_stage(session: AsyncSession, deal: Deal) -> None:
    items = (await session.scalars(
        select(DealItem).where(DealItem.deal_id == deal.id,
                               DealItem.status.in_([ItemStatus.PENDING, ItemStatus.FAILED]))
    )).all()
    if not items:
        logger.info("categorization_skipped_no_items", extra={"deal_id": str(deal.id)})
        return

    settings   = await get_categorizer_settings(session)
    storage    = get_storage()
    totals     = {"input_tokens": 0, "output_tokens": 0}

    async with httpx.AsyncClient() as http:
        categorizer = await build_categorizer(settings, http)
        sem = asyncio.Semaphore(settings.concurrency)

        async def classify_one(item: DealItem) -> None:
            async with sem:
                try:
                    image_bytes = await storage.read_bytes(item.crop_key)
                    result, usage = await categorizer.classify(image_bytes=image_bytes, size_hint=None)
                    item.category               = result.category
                    item.category_confidence    = result.category_confidence
                    item.subcategory            = result.subcategory
                    item.subcategory_id         = result.subcategory_id
                    item.subcategory_confidence = result.subcategory_confidence
                    item.status                 = ItemStatus.CLASSIFIED
                    item.error                  = None
                    totals["input_tokens"]  += usage.input_tokens
                    totals["output_tokens"] += usage.output_tokens
                except Exception as exc:
                    item.status = ItemStatus.FAILED
                    item.error  = repr(exc)[:2000]
                    logger.warning("categorization_item_failed",
                                   extra={"deal_id": str(deal.id),
                                          "item_id": str(item.id),
                                          "error": repr(exc)})
                finally:
                    await session.flush()  # keep ORM state consistent within the open tx

        async with asyncio.TaskGroup() as tg:
            for item in items:
                tg.create_task(classify_one(item))

    # Record aggregate usage onto the in-progress PipelineRun row
    run = await _current_pipeline_run(session, deal, PipelineStage.CATEGORIZATION)
    if run is not None:
        run.input_tokens  = totals["input_tokens"]
        run.output_tokens = totals["output_tokens"]
        run.model         = _resolve_model(settings)
    # commit happens in the orchestrator

_current_pipeline_run(...) selects the in-flight pipeline_runs row (status=running) for this deal+stage; the orchestrator at features/pipeline/orchestrator.py:31-38 creates and flushes that row before calling the stage. _resolve_model(settings) returns settings.two_pass.model or (settings.single_pass or SinglePassConfig()).model based on kind.

§9Error handling

• Per-item LLM failure (timeout, OpenRouter 5xx, JSON parse error): item marked failed, error captured, stage proceeds. Re-runs reprocess pending + failed.
• Per-item storage read failure: same as per-item LLM failure.
• Settings / build_categorizer / get_storage failure: bubble up. Orchestrator rolls back the stage transaction. Deal stays at SEGMENTATION. Procrastinate retries with backoff.
• OpenRouter auth (401) / quota (429): bubble up after the first item — not worth burning items on a known-bad config. OpenRouterClient raises typed OpenRouterAuthError / OpenRouterQuotaError on these statuses; classify_one re-raises them past the try/except. Other exceptions are caught.
• TaskGroup cancellation: if any task re-raises (auth/quota cases), the TaskGroup cancels the rest. In-flight item updates already flushed survive in memory; the stage transaction then rolls back. Items reprocess on retry.

§10Concurrency

• asyncio.Semaphore(settings.concurrency) per deal (default 10).
• No global cap across deals. Worker is single-process today; concurrent deals are not a real concern.
• When the worker scales, a global cap will likely live on the OpenRouterClient — out of slice.

§11Configuration

Environment

Runtime (DB-backed)

categorizer_settings.settings: kind, concurrency, per-kind model / timeout_seconds / temperature.

Repo (no env)

features/categorization/config/categories.toml, config/subcategories.csv, prompts/*.md — copied verbatim from legacy. Edited via PR.

§12Testing

1 · Unit — strategies (test_two_pass.py, test_single_pass.py)

Mock OpenRouterClient.vision_json_request to return canned JsonResponses. Assert:

• Two-pass makes exactly two calls; single-pass makes one.
• Prompt text contains injected {categories} / {subcategory_block} / {tree}.
• size_suffix appended when size_hint provided, omitted when not.
• Known subcategory label resolves to right id; case-insensitive match works; unknown label → subcategory_id=None (label preserved).
• TokenUsage summed across passes (two-pass) or passed through (single-pass).

2 · Unit — taxonomy (test_taxonomy.py)

Load the real shipped TOML + CSV. Assert 11 categories load, known subcategory resolves both case-sensitively and case-insensitively, rendered blocks for each category are non-empty.

3 · Integration — stage (test_step.py, real Postgres, mocked LLM + storage)

Uses the existing testcontainers conftest. Per test:

• Create a Deal, insert 3 DealItem rows.
• Patch OpenRouterClient.vision_json_request to return [success, raises httpx.ReadTimeout, success].
• Patch Storage.read_bytes (monkeypatch at get_storage() cache) to return a canned 64×64 PNG.
• Run run_categorization_stage(session, deal).
• Assert: 2 items classified with fields set, 1 item failed with classification NULL and error set, PipelineRun for CATEGORIZATION has aggregate tokens + model, stage returned without raising.

Second test (test_step_rerun.py): pre-seed one failed and one classified item; run again; assert only failed retried, classified untouched.

4 · Unit — storage (tests/core/test_storage.py)

LocalFsStorage round-trip under tmp_path; delete_prefix recurses; signed_url returns /files/{key} with URL-quoting; path-escape attempt (..) rejected.

Fixtures

tests/features/categorization/fixtures/sample.png — single 64×64 RGB PNG checked into the repo. ~200 bytes.

§13Decisions made (during brainstorm)

1. DealItem rows come from segmentation. Cleaner contract than letting categorization scan disk. Categorization reads WHERE status IN ('pending','failed').
2. DB-backed CategorizerSettings with kind enum (two_pass | single_pass). Mirrors SegmenterSettings; consistent with the existing factory pattern.
3. LLM client lives in integrations/openrouter/. Parallels integrations/hubspot/; predictable home for appraisal next.
4. Categories / subcategories / prompts: copy legacy files as-is under features/categorization/. Re-read each call. Edits live in dev.
5. Token tracking: nullable columns on pipeline_runs (input_tokens, output_tokens, model). Cheap; useful immediately. No separate llm_calls table.
6. Per-deal concurrency: 10, configurable. Matches legacy.
7. Storage abstraction lands in this slice; LocalFsStorage only. Azure backend is a follow-up that only implements the protocol.
8. Describe-then-classify is not ported. The variation registry pattern leaves the seam.
9. Per-item failures don't fail the stage. Items get status='failed' + error; re-runs reprocess them. Stage-level failures (auth, quota, config) bubble up for orchestrator backoff/retry.