from typing import List
from openai import AsyncOpenAI


class SimpleOpenAIEmbedding:
    def __init__(
        self,
        model,
        api_key,
        api_base=None,
    ) -> None:
        self.client = AsyncOpenAI(api_key=api_key, base_url=api_base)
        self.model = model

    async def get_embedding(self, text) -> List[float]:
        """
        获取文本的嵌入
        """
        embedding = await self.client.embeddings.create(input=text, model=self.model)
        return embedding.data[0].embedding