from dataclasses import dataclass, field
from typing import List, Optional, Union
from rest_framework import serializers


class GerarDocumentoParametros:
    num_chunks_retrieval = serializers.IntegerField(default=20)
    embedding_weight = serializers.FloatField(default=0.5)
    bm25_weight = serializers.FloatField(default=0.5)
    context_window = serializers.IntegerField(default=3)
    num_k_rerank = serializers.IntegerField(default=20)
    more_initial_chunks_for_reranking = serializers.IntegerField(default=100)
    gpt_temperature = serializers.FloatField(default=0)


@dataclass
class GerarDocumentoParametrosData:
    num_chunks_retrieval: int = 20
    embedding_weight: float = 0.5
    bm25_weight: float = 0.5
    context_window: int = 3
    num_k_rerank: int = 20
    more_initial_chunks_for_reranking: int = 100
    gpt_temperature: float = 0.0