from rest_framework import serializers
from setup.environment import default_model
# from _utils.utils import DEFAULT_SYSTEM_PROMPT

prompt_template = """
    Based on the following context, provide multiple key points from the document.
    For each point, create a new paragraph.
    Each paragraph should be a complete, self-contained insight.
    
    Context: {context}
    
    Key points:
    """

class ResumoPDFSerializer(serializers.Serializer):
    files = serializers.ListField(child=serializers.FileField(), required=True)
    system_prompt = serializers.CharField(required=False)
    user_message = serializers.CharField(required=False, default="")
    model = serializers.CharField(required=False)
    iterative_refinement = serializers.BooleanField(required=False, default=False)

class ResumoCursorSerializer(serializers.Serializer):
    files = serializers.ListField(child=serializers.FileField(), required=True)
    system_prompt = serializers.CharField(required=False, default=prompt_template)
    user_message = serializers.CharField(required=False, default="")
    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=1000)
    chunk_overlap = serializers.IntegerField(required=False, default=200)


system_prompt = """
            Based on the following context, provide multiple key points from the document.
            For each point, create a new paragraph.
            Each paragraph should be a complete, self-contained insight.
            Include any relevant context provided.
            
            Context: {context}
            
            Modelo do usuário: {modelo_usuario}
            
            Key points:
            """
user_message = "What are the main points of this document?"
class ResumoCursorCompeltoSerializer(ResumoCursorSerializer):
    system_prompt = serializers.CharField(required=False, default=system_prompt)
    user_message = serializers.CharField(required=False, default=user_message)
    num_chunks_retrieval = serializers.IntegerField(default=5)
    embedding_weight = serializers.FloatField(default=0.5)
    bm25_weight = serializers.FloatField(default=0.5)
    context_window = serializers.IntegerField(default=3)
    chunk_overlap = serializers.IntegerField(default=200)
    num_k_rerank = serializers.IntegerField(default=5)
    model_cohere_rerank = serializers.CharField(required=False, default="rerank-english-v2.0")
    more_initial_chunks_for_reranking = serializers.IntegerField(default=20)
    claude_context_model = serializers.CharField(required=False, default="claude-3-haiku-20240307")
    gpt_temperature = serializers.FloatField(default=0)
    id_modelo_do_usuario = serializers.IntegerField(required=True)