gerar_documento/serializer.py

from dataclasses import dataclass, field
from typing import List, Optional, Union
from rest_framework import serializers
from _utils.gerar_documento_utils.prompts import (
    prompt_gerar_documento,
)
from gerar_documento.serializer_base import (
    GerarDocumentoParametros,
    GerarDocumentoParametrosData,
)
from setup.environment import default_model
from django.core.files.uploadedfile import UploadedFile

user_message = "What are the main points of this document?"


class GerarDocumentoInitialSerializer(serializers.Serializer):
    files = serializers.ListField(child=serializers.FileField(), required=True)
    user_message = serializers.CharField(required=False, default=user_message)
    model = serializers.CharField(required=False, default=default_model)
    hf_embedding = serializers.CharField(required=False, default="all-MiniLM-L6-v2")
    chunk_size = serializers.IntegerField(required=False, default=3500)
    chunk_overlap = serializers.IntegerField(required=False, default=800)
    prompt_gerar_documento = serializers.CharField(
        required=False, default=prompt_gerar_documento
    )
    prompt_gerar_documento_etapa_2 = serializers.CharField(required=False)
    prompt_gerar_documento_etapa_3 = serializers.CharField(required=False)
    model_cohere_rerank = serializers.CharField(
        required=False, default="rerank-english-v2.0"
    )
    claude_context_model = serializers.CharField(
        required=False, default="claude-3-haiku-20240307"
    )
    should_have_contextual_chunks = serializers.BooleanField(default=False)  # type: ignore
    should_use_llama_parse = serializers.BooleanField(required=False, default=False)  # type: ignore
    llm_ultimas_requests = serializers.CharField(
        required=False, default="gemini-2.0-flash"
    )


@dataclass
class GerarDocumentoInitialSerializerData:
    files: List[dict]
    user_message: str = ""
    model: str = default_model
    hf_embedding: str = "all-MiniLM-L6-v2"
    chunk_size: int = 3500
    chunk_overlap: int = 800
    prompt_gerar_documento: str = ""
    prompt_gerar_documento_etapa_2: Union[str, None] = None
    prompt_gerar_documento_etapa_3: Union[str, None] = None
    model_cohere_rerank: str = "rerank-english-v2.0"
    claude_context_model: str = "claude-3-haiku-20240307"
    should_have_contextual_chunks: bool = False
    should_use_llama_parse: bool = False
    llm_ultimas_requests: str = "gemini-2.0-flash"


class FileInfoSerializer(serializers.Serializer):
    unique_id = serializers.CharField(max_length=255)
    tipo_arquivo = serializers.CharField(max_length=255)
    link_arquivo = serializers.URLField()


@dataclass
class FileInfoSerializerData:
    unique_id: str
    tipo_arquivo: str
    link_arquivo: str


class GerarDocumentoSerializer(
    GerarDocumentoInitialSerializer, GerarDocumentoParametros
):
    files = serializers.CharField(required=True)
    bubble_editor_version = serializers.CharField(
        required=False, default="version-test"
    )  # Será o valor utilizado dentro da URL da requisição pro Bubble
    doc_id = serializers.CharField(required=True)
    form_response_id = serializers.CharField(required=True)
    version = serializers.CharField(required=True)

    def get_obj(self):
        return GerarDocumentoSerializerData(**self.validated_data)  # type: ignore


@dataclass
class GerarDocumentoSerializerData(
    GerarDocumentoParametrosData, GerarDocumentoInitialSerializerData
):
    files: str
    bubble_editor_version: str = "version-test"

    doc_id: str = ""
    form_response_id: str = ""
    version: str = ""


class GerarDocumentoComPDFProprioSerializer(
    GerarDocumentoInitialSerializer, GerarDocumentoParametros
):
    files = serializers.ListField(required=True)

    def get_obj(self):
        return GerarDocumentoSerializerData(**self.validated_data)  # type: ignore


@dataclass
class GerarDocumentoComPDFProprioSerializerData(
    GerarDocumentoParametrosData, GerarDocumentoInitialSerializerData
):
    files: List[FileInfoSerializerData]