from statistics import mean

from llama_index.core import VectorStoreIndex
from llama_index.core.embeddings import BaseEmbedding
from llama_index.core.evaluation import RetrieverEvaluator
from llama_index.core.schema import TextNode


class KeywordEmbedding(BaseEmbedding):
    def _vector_for_text(self, text: str) -> list[float]:
        text = text.lower()
        if "billing" in text and "release" in text:
            return [0.8, 0.6, 0.0]
        if "billing" in text or "7421" in text:
            return [1.0, 0.0, 0.0]
        if "release" in text:
            return [0.0, 1.0, 0.0]
        return [0.0, 0.0, 1.0]

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

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

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