from llama_index.core import Document, Settings, VectorStoreIndex
from llama_index.core.embeddings import MockEmbedding
from llama_index.core.llms import CompletionResponse, CustomLLM, LLMMetadata
from llama_index.core.llms.callbacks import llm_completion_callback
from llama_index.core.postprocessor.types import BaseNodePostprocessor


class TicketReranker(BaseNodePostprocessor):
    top_n: int = 2

    def _postprocess_nodes(self, nodes, query_bundle=None):
        query = query_bundle.query_str.lower() if query_bundle else ""

        for item in nodes:
            text = item.node.get_content().lower()
            score = 0.0
            score += 2.0 if "ticket 7421" in text else 0.0
            score += 2.0 if "who owns" in query and "belongs to" in text else 0.0
            score += 1.0 if "escalated" in query and "escalated" in text else 0.0
            item.score = score

        return sorted(nodes, key=lambda item: item.score or 0.0, reverse=True)[
            : self.top_n
        ]