schemas.py 1.11 KB
from __future__ import annotations

from dataclasses import asdict, dataclass, field
from typing import Any


@dataclass
class ParsedDocument:
    doc_id: str
    source_file: str
    file_type: str
    content: str
    page: int | None = None
    sheet: str | None = None
    row_index: int | None = None
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class TestsetRecord:
    sample_id: str
    user_input: str
    reference: str
    reference_contexts: list[str]
    source_file: str | None = None
    gold_chunk_ids: list[str] = field(default_factory=list)
    question_type: str = "single_hop"
    review_status: str = "pending"

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass
class WeKnoraAnswer:
    sample_id: str
    user_input: str
    response: str
    retrieved_contexts: list[str]
    weknora_references: list[dict[str, Any]]
    session_id: str | None = None
    request_id: str | None = None
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)