File size: 1,412 Bytes
44c2f50 c6cfb1f 44c2f50 c6cfb1f 44c2f50 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """Schema registry — maps document-type strings to Pydantic schemas.
The API endpoint uses this to look up which schema to extract into based on
the `doc_type` argument the caller provides. Adding a new domain (e.g. filings
in v2) is a matter of registering it here.
"""
from __future__ import annotations
from pydantic import BaseModel
from src.schemas.filing import Filing
from src.schemas.invoice import Invoice
from src.schemas.receipt import Receipt
# Registry: doc_type string -> Pydantic schema class
_REGISTRY: dict[str, type[BaseModel]] = {
"invoice": Invoice,
"receipt": Receipt,
"filing": Filing, # SEC 10-K (v2). Same key handles 10-K and 10-K/A.
}
def get_schema(doc_type: str) -> type[BaseModel]:
"""Look up a schema class by document type. Raises KeyError if unknown."""
key = doc_type.strip().lower()
if key not in _REGISTRY:
available = ", ".join(sorted(_REGISTRY.keys()))
raise KeyError(f"Unknown doc_type {doc_type!r}. Available: {available}")
return _REGISTRY[key]
def list_doc_types() -> list[str]:
"""Return all registered document types (used by GET /schemas)."""
return sorted(_REGISTRY.keys())
def get_json_schema(doc_type: str) -> dict:
"""Return the JSON Schema for a doc type — used by the API and by OpenAI structured outputs."""
schema_cls = get_schema(doc_type)
return schema_cls.model_json_schema()
|