afulara/PythonicRAG-FastAPI-React

1from dotenv import load_dotenv2from openai import AsyncOpenAI, OpenAI3import openai4from typing import List5import os6import asyncio7 8 9class EmbeddingModel:10    def __init__(self, embeddings_model_name: str = "text-embedding-3-small"):11        load_dotenv()12        self.openai_api_key = os.getenv("OPENAI_API_KEY")13        self.async_client = AsyncOpenAI()14        self.client = OpenAI()15 16        if self.openai_api_key is None:17            raise ValueError(18                "OPENAI_API_KEY environment variable is not set. Please set it to your OpenAI API key."19            )20        openai.api_key = self.openai_api_key21        self.embeddings_model_name = embeddings_model_name22 23    async def async_get_embeddings(self, list_of_text: List[str]) -> List[List[float]]:24        embedding_response = await self.async_client.embeddings.create(25            input=list_of_text, model=self.embeddings_model_name26        )27 28        return [embeddings.embedding for embeddings in embedding_response.data]29 30    async def async_get_embedding(self, text: str) -> List[float]:31        embedding = await self.async_client.embeddings.create(32            input=text, model=self.embeddings_model_name33        )34 35        return embedding.data[0].embedding36 37    def get_embeddings(self, list_of_text: List[str]) -> List[List[float]]:38        embedding_response = self.client.embeddings.create(39            input=list_of_text, model=self.embeddings_model_name40        )41 42        return [embeddings.embedding for embeddings in embedding_response.data]43 44    def get_embedding(self, text: str) -> List[float]:45        embedding = self.client.embeddings.create(46            input=text, model=self.embeddings_model_name47        )48 49        return embedding.data[0].embedding50 51 52if __name__ == "__main__":53    embedding_model = EmbeddingModel()54    print(asyncio.run(embedding_model.async_get_embedding("Hello, world!")))55    print(56        asyncio.run(57            embedding_model.async_get_embeddings(["Hello, world!", "Goodbye, world!"])58        )59    )60