from llama_index.core import Document, VectorStoreIndex
from llama_index.core.embeddings import BaseEmbedding


class SupportKeywordEmbedding(BaseEmbedding):
    @classmethod
    def class_name(cls) -> str:
        return "SupportKeywordEmbedding"

    def _vector(self, text: str) -> list[float]:
        terms = [
            "billing",
            "refund",
            "escalation",
            "support",
            "inventory",
            "warehouse",
            "restart",
            "sync",
        ]
        lowered = text.lower()
        return [1.0 if term in lowered else 0.0 for term in terms]

    def _get_text_embedding(self, text: str) -> list[float]:
        return self._vector(text)

    def _get_query_embedding(self, query: str) -> list[float]:
        return self._vector(query)

    async def _aget_query_embedding(self, query: str) -> list[float]:
        return self._vector(query)