pinecone_text.dense.openai_encoder

  1import os
  2from typing import Union, List, Any, Optional, Dict
  3from pinecone_text.dense.base_dense_ecoder import BaseDenseEncoder
  4
  5try:
  6    import openai
  7    from openai import OpenAIError
  8except (OSError, ImportError, ModuleNotFoundError) as e:
  9    _openai_installed = False
 10else:
 11    _openai_installed = True
 12
 13
 14class OpenAIEncoder(BaseDenseEncoder):
 15    """
 16    OpenAI's text embedding wrapper. See https://platform.openai.com/docs/guides/embeddings
 17
 18    Note: this method reflects the OpenAI client initialization behaviour (See https://github.com/openai/openai-python/blob/main/src/openai/_client.py)
 19      On initialization, You may explicitly pass any argument that the OpenAI client accepts, or use the following environment variables:
 20
 21    - `OPENAI_API_KEY` as `api_key`
 22    - `OPENAI_ORG_ID` as `organization`
 23    - `OPENAI_BASE_URL` as `base_url`
 24
 25    Example:
 26        Using environment variables:
 27        >>> import os
 28        >>> from pinecone_text.dense import OpenAIEncoder
 29        >>> os.environ['OPENAI_API_KEY'] = "sk-..."
 30        >>> encoder = OpenAIEncoder()
 31        >>> encoder.encode_documents(["some text", "some other text"])
 32
 33        Passing arguments explicitly:
 34        >>> from pinecone_text.dense import OpenAIEncoder
 35        >>> encoder = OpenAIEncoder(api_key="sk-...")
 36    """  # noqa: E501
 37
 38    def __init__(
 39        self,
 40        model_name: str = "text-embedding-3-small",
 41        *,
 42        dimension: Optional[int] = None,
 43        **kwargs: Any,
 44    ):
 45        if not _openai_installed:
 46            raise ImportError(
 47                "Failed to import openai. Make sure you install openai extra "
 48                "dependencies by running: "
 49                "`pip install pinecone-text[openai]"
 50            )
 51        super().__init__()
 52
 53        if dimension is not None:
 54            assert dimension > 0, "dimension must be a positive integer"
 55
 56        self._model_name = model_name
 57        self._dimension = dimension
 58        self._client = self._create_client(**kwargs)
 59
 60    @staticmethod
 61    def _create_client(**kwargs: Any) -> Union[openai.OpenAI, openai.AzureOpenAI]:
 62        return openai.OpenAI(**kwargs)
 63
 64    def encode_documents(
 65        self, texts: Union[str, List[str]]
 66    ) -> Union[List[float], List[List[float]]]:
 67        return self._encode(texts)
 68
 69    def encode_queries(
 70        self, texts: Union[str, List[str]]
 71    ) -> Union[List[float], List[List[float]]]:
 72        return self._encode(texts)
 73
 74    def _encode(
 75        self, texts: Union[str, List[str]]
 76    ) -> Union[List[float], List[List[float]]]:
 77        if isinstance(texts, str):
 78            texts_input = [texts]
 79        elif isinstance(texts, list):
 80            texts_input = texts
 81        else:
 82            raise ValueError(
 83                f"texts must be a string or list of strings, got: {type(texts)}"
 84            )
 85
 86        try:
 87            params: Dict[str, Any] = dict(
 88                input=texts_input,
 89                model=self._model_name,
 90            )
 91            if self._dimension is not None:
 92                params["dimensions"] = self._dimension
 93            response = self._client.embeddings.create(**params)
 94        except OpenAIError as e:
 95            # TODO: consider wrapping external provider errors
 96            raise e
 97
 98        if isinstance(texts, str):
 99            return response.data[0].embedding
100        return [result.embedding for result in response.data]
101
102
103class AzureOpenAIEncoder(OpenAIEncoder):
104    """
105    Initialize the Azure OpenAI encoder.
106
107    Note: this method reflects the AzureOpenAI client initialization behaviour (See https://github.com/openai/openai-python/blob/main/src/openai/lib/azure.py).
108           You may explicitly pass any argument that the AzureOpenAI client accepts, or use the following environment variables:
109
110        - `AZURE_OPENAI_API_KEY` as `api_key`
111        - `AZURE_OPENAI_ENDPOINT` as `azure_endpoint`
112        - `OPENAI_API_VERSION` as `api_version`
113        - `OPENAI_ORG_ID` as `organization`
114        - `AZURE_OPENAI_AD_TOKEN` as `azure_ad_token`
115
116    In addition, you must pass the `model_name` argument with the name of the deployment you wish to use in your own Azure account.
117
118    Example:
119        Using environment variables:
120        >>> import os
121        >>> from pinecone_text.dense import AzureOpenAIEncoder
122        >>> os.environ['AZURE_OPENAI_API_KEY'] = "sk-..."
123        >>> os.environ['AZURE_OPENAI_ENDPOINT'] = "https://.....openai.azure.com/"
124        >>> os.environ['OPENAI_API_VERSION'] = "2023-12-01-preview"
125        >>> encoder = AzureOpenAIEncoder(model_name="my-ada-002-deployment")
126        >>> encoder.encode_documents(["some text", "some other text"])
127
128        Passing arguments explicitly:
129        >>> from pinecone_text.dense import AzureOpenAIEncoder
130        >>> encoder = AzureOpenAIEncoder(api_key="sk-...", azure_endpoint="https://.....openai.azure.com/", api_version="2023-12-01-preview")
131    """  # noqa: E501
132
133    def __init__(self, model_name: str, **kwargs: Any):
134        super().__init__(model_name=model_name, **kwargs)
135
136    @staticmethod
137    def _create_client(**kwargs: Any) -> openai.AzureOpenAI:
138        return openai.AzureOpenAI(**kwargs)