Harrison/azure embeddings (#1787)

Co-authored-by: Hemant <4627288+ghaccount@users.noreply.github.com>

docs/modules/indexes/examples/embeddings.ipynb

@@ -76,6 +76,131 @@
     "doc_result = embeddings.embed_documents([text])"
    ]
   },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "bb61bbeb",
+   "metadata": {},
+   "source": [
+    "Let's load the OpenAI Embedding class with first generation models (e.g. text-search-ada-doc-001/text-search-ada-query-001). Note: These are not recommended models - see [here](https://platform.openai.com/docs/guides/embeddings/what-are-embeddings)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0b072cc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.embeddings.openai import OpenAIEmbeddings"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a56b70f5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "embeddings = OpenAIEmbeddings(model_name=\"ada\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14aefb64",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text = \"This is a test document.\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3c39ed33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "query_result = embeddings.embed_query(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e3221db6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "doc_result = embeddings.embed_documents([text])"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "c3852491",
+   "metadata": {},
+   "source": [
+    "## AzureOpenAI\n",
+    "\n",
+    "Let's load the OpenAI Embedding class with environment variables set to indicate to use Azure endpoints."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1b40f827",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# set the environment variables needed for openai package to know to reach out to azure\n",
+    "import os\n",
+    "\n",
+    "os.environ[\"OPENAI_API_TYPE\"] = \"azure\"\n",
+    "os.environ[\"OPENAI_API_BASE\"] = \"https://<your-endpoint.openai.azure.com/\"\n",
+    "os.environ[\"OPENAI_API_KEY\"] = \"your AzureOpenAI key\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bb36d16c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "embeddings = OpenAIEmbeddings(model=\"your-embeddings-deployment-name\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "228abcbb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text = \"This is a test document.\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "60dd7fad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "query_result = embeddings.embed_query(text)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "83bc1a72",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "doc_result = embeddings.embed_documents([text])"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "42f76e43",
@@ -86,6 +211,13 @@
     "Let's load the Cohere Embedding class."
    ]
   },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "id": "ca9e2b3a",
+   "metadata": {},
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": 1,
@@ -103,7 +235,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "embeddings = CohereEmbeddings(cohere_api_key= cohere_api_key)"
+    "embeddings = CohereEmbeddings(cohere_api_key=cohere_api_key)"
    ]
   },
   {
@@ -290,7 +422,9 @@
     }
    ],
    "source": [
-    "embeddings = HuggingFaceInstructEmbeddings(query_instruction=\"Represent the query for retrieval: \")"
+    "embeddings = HuggingFaceInstructEmbeddings(\n",
+    "    query_instruction=\"Represent the query for retrieval: \"\n",
+    ")"
    ]
   },
   {
@@ -332,9 +466,9 @@
    "outputs": [],
    "source": [
     "from langchain.embeddings import (\n",
-    "    SelfHostedEmbeddings, \n",
-    "    SelfHostedHuggingFaceEmbeddings, \n",
-    "    SelfHostedHuggingFaceInstructEmbeddings\n",
+    "    SelfHostedEmbeddings,\n",
+    "    SelfHostedHuggingFaceEmbeddings,\n",
+    "    SelfHostedHuggingFaceInstructEmbeddings,\n",
     ")\n",
     "import runhouse as rh"
    ]
@@ -353,7 +487,7 @@
     "# gpu = rh.cluster(name='rh-a10x', instance_type='g5.2xlarge', provider='aws')\n",
     "\n",
     "# For an existing cluster\n",
-    "# gpu = rh.cluster(ips=['<ip of the cluster>'], \n",
+    "# gpu = rh.cluster(ips=['<ip of the cluster>'],\n",
     "#                  ssh_creds={'ssh_user': '...', 'ssh_private_key':'<path_to_key>'},\n",
     "#                  name='my-cluster')"
    ]
@@ -424,16 +558,22 @@
    "outputs": [],
    "source": [
     "def get_pipeline():\n",
-    "    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline  # Must be inside the function in notebooks\n",
+    "    from transformers import (\n",
+    "        AutoModelForCausalLM,\n",
+    "        AutoTokenizer,\n",
+    "        pipeline,\n",
+    "    )  # Must be inside the function in notebooks\n",
+    "\n",
     "    model_id = \"facebook/bart-base\"\n",
     "    tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
     "    model = AutoModelForCausalLM.from_pretrained(model_id)\n",
     "    return pipeline(\"feature-extraction\", model=model, tokenizer=tokenizer)\n",
     "\n",
+    "\n",
     "def inference_fn(pipeline, prompt):\n",
     "    # Return last hidden state of the model\n",
     "    if isinstance(prompt, list):\n",
-    "        return [emb[0][-1] for emb in pipeline(prompt)] \n",
+    "        return [emb[0][-1] for emb in pipeline(prompt)]\n",
     "    return pipeline(prompt)[0][-1]"
    ]
   },
@@ -445,10 +585,10 @@
    "outputs": [],
    "source": [
     "embeddings = SelfHostedEmbeddings(\n",
-    "    model_load_fn=get_pipeline, \n",
+    "    model_load_fn=get_pipeline,\n",
     "    hardware=gpu,\n",
     "    model_reqs=[\"./\", \"torch\", \"transformers\"],\n",
-    "    inference_fn=inference_fn\n",
+    "    inference_fn=inference_fn,\n",
     ")"
    ]
   },

langchain/embeddings/openai.py

@@ -65,9 +65,35 @@ class OpenAIEmbeddings(BaseModel, Embeddings):
 
             from langchain.embeddings import OpenAIEmbeddings
             openai = OpenAIEmbeddings(openai_api_key="my-api-key")
+
+    In order to use the library with Microsoft Azure endpoints, you need to set
+    the OPENAI_API_TYPE, OPENAI_API_BASE, OPENAI_API_KEY and optionally and
+    API_VERSION.
+    The OPENAI_API_TYPE must be set to 'azure' and the others correspond to
+    the properties of your endpoint.
+    In addition, the deployment name must be passed as the model parameter.
+
+    Example:
+        .. code-block:: python
+
+            import os
+            os.environ["OPENAI_API_TYPE"] = "azure"
+            os.environ["OPENAI_API_BASE"] = "https://<your-endpoint.openai.azure.com/"
+            os.environ["OPENAI_API_KEY"] = "your AzureOpenAI key"
+
+            from langchain.embeddings.openai import OpenAIEmbeddings
+            embeddings = OpenAIEmbeddings(model="your-embeddings-deployment-name")
+            text = "This is a test query."
+            query_result = embeddings.embed_query(text)
+
     """
 
     client: Any  #: :meta private:
+    model: str = "text-embedding-ada-002"
+
+    # TODO: deprecate these two in favor of model
+    #  https://community.openai.com/t/api-update-engines-models/18597
+    #  https://github.com/openai/openai-python/issues/132
     document_model_name: str = "text-embedding-ada-002"
     query_model_name: str = "text-embedding-ada-002"
     embedding_ctx_length: int = -1
@@ -85,6 +111,14 @@ class OpenAIEmbeddings(BaseModel, Embeddings):
     # TODO: deprecate this
     @root_validator(pre=True)
     def get_model_names(cls, values: Dict) -> Dict:
+        # model_name is for first generation, and model is for second generation.
+        # Both are not allowed together.
+        if "model_name" in values and "model" in values:
+            raise ValueError(
+                "Both `model_name` and `model` were provided, "
+                "but only one should be."
+            )
+
         """Get model names from just old model name."""
         if "model_name" in values:
             if "document_model_name" in values:
@@ -100,6 +134,23 @@ class OpenAIEmbeddings(BaseModel, Embeddings):
             model_name = values.pop("model_name")
             values["document_model_name"] = f"text-search-{model_name}-doc-001"
             values["query_model_name"] = f"text-search-{model_name}-query-001"
+
+        # Set document/query model names from model parameter.
+        if "model" in values:
+            if "document_model_name" in values:
+                raise ValueError(
+                    "Both `model` and `document_model_name` were provided, "
+                    "but only one should be."
+                )
+            if "query_model_name" in values:
+                raise ValueError(
+                    "Both `model` and `query_model_name` were provided, "
+                    "but only one should be."
+                )
+            model = values.get("model")
+            values["document_model_name"] = model
+            values["query_model_name"] = model
+
         return values
 
     @root_validator()