from typing import List

import requests
from exceptions import EmbeddingError
from resolver import pattern

from patterns.embedding import EmbeddingModel


@pattern(name="Local_Embedding_Private")
class LocalEmbeddingPrivate(EmbeddingModel):
    """
    Embedding model using local embedding model.
    """

    def __init__(
        self,
        embedding_endpoint: str,
        embedding_model: str,
        prefix: str = "[QRY]",
        suffix: str = "",
    ):
        self.embedding_api = f"{embedding_endpoint}/embedding-models/{embedding_model}"
        self.prefix = prefix
        self.suffix = suffix

    def embedding(self, text: str) -> List[float]:
        """
        Get the embedding of the text.

        Args:
            text (str): The text to be embedded.

        Returns:
            List[float]: The embedding of the text.
        """
        response = requests.post(
            self.embedding_api, json={"text": self.prefix + text + self.suffix}
        )
        if response.status_code != 200:
            raise EmbeddingError(f"Failed to embed text: {text}")
        return response.json()