_utils/models/gerar_documento.py

from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass
import numpy as np


@dataclass
class DocumentChunk:
    content: str
    page_number: int
    chunk_id: str
    start_char: int
    end_char: int
    id_do_processo: int = 0
    contextual_summary: str = ""


@dataclass
class RetrievalConfig:
    num_chunks: int = 5
    embedding_weight: float = 0.5
    bm25_weight: float = 0.5
    context_window: int = 3
    chunk_overlap: int = 200
    chunk_size: int = 1000


@dataclass
class ContextualizedChunk(DocumentChunk):
    context: str = ""
    embedding: Optional[np.ndarray] = None
    bm25_score: Optional[float] = None