from abc import ABC, abstractmethod
from typing import Optional, Any

from pydantic import BaseModel
from sqlmodel import Field
from pytidb.schema import VectorField


class BaseEmbeddingFunction(BaseModel, ABC):
    model_name: str = Field(
        None, description="The name of embedding model used for embedding"
    )
    dimensions: Optional[int] = Field(
        None, description="Dimensions of the vector generated by the embedding model"
    )

    def __init__(self, /, **data: Any):
        super().__init__(**data)

    def VectorField(self, source_field: Optional[str] = None, **kwargs):
        return VectorField(
            embed_fn=self,
            dimensions=self.dimensions,
            source_field=source_field,
            **kwargs,
        )

    @abstractmethod
    def get_query_embedding(self, query: str) -> list[float]:
        raise NotImplementedError()

    @abstractmethod
    def get_source_embedding(self, source: str) -> list[float]:
        raise NotImplementedError()

    @abstractmethod
    def get_source_embeddings(self, source: list[str]) -> list[list[float]]:
        raise NotImplementedError()
