langchain_wrappers/langchain_embedding_models.py

import os
from dotenv import load_dotenv
from typing import Optional
from langchain_openai.embeddings import OpenAIEmbeddings

# import inspect
# load_dotenv(os.path.join(os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) , '.env'))

load_dotenv()
class MyOpenAIEmbeddings:
    @classmethod
    def from_model(
        cls, 
        model: str = 'small',
        *,
        dimensions: Optional[int] = None,
        max_retries: int = 1,
        **kwargs
        )-> OpenAIEmbeddings:
        
        if model in ['text-embedding-3-small', 'TEXT-EMBEDDING-3-SMALL', 'small', 'SMALL']:
            model = 'text-embedding-3-small'
            dimensions = 1536 if dimensions is None else dimensions
        elif model in ['text-embedding-3-large', 'TEXT-EMBEDDING-3-LARGE', 'large', 'LARGE']:
            model = 'text-embedding-3-large'
            dimensions = 3072 if dimensions is None else dimensions
        else:
            raise ValueError(f"Model {model} is currently not supported. Supported models are: ['text-embedding-3-small', 'text-embedding-3-large']")

        
        return OpenAIEmbeddings(
            openai_api_key=os.getenv("OPENAI_API_KEY"), 
            model=model,
            dimensions=dimensions,
            max_retries=max_retries,
            **kwargs
            )