From 7287d2081f68c81cc28a8126ae42aa75bb0aef4e Mon Sep 17 00:00:00 2001
From: Facundo Goiriz <81929318+Fakamoto@users.noreply.github.com>
Date: Thu, 21 Mar 2024 09:13:33 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=85=B0=EF=B8=8F=20feat:=20Add=20support?=
 =?UTF-8?q?=20for=20Azure=20OpenAI=20Embeddings=20(#4)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: fgoiriz <fgoiriz@contractor.networkninja.com>
---
 README.md | Bin 7164 -> 7768 bytes
 config.py |  12 ++++++++++--
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index ae7d0a4e0c1496b98a978c6b30be47886347424f..708efc0a8a87f1d1c15a4bdcc17d0f0681c5db71 100644
GIT binary patch
delta 440
zcma)&y-EX75QUG0h01DSvkJ-*5exHE3c+T<h(bt$Xrsk$NY)>S*|1qfZ1x%CmZtFy
z1Yg5mAH^bt-(6)<u#j8K+_^K~J#(fz@4L?vceyN2C3SVL8zqX!%jVJDo{j9yJh^M;
z)t$~2?<QM$Tkv(N6CJ6cz-&4j=3{Z)bR750FDRbc8bKS%*C9WQ$rIh`N<Ffw9L=jj
zJ><P$G1GzB%4dQ++oRHk&@vB8Tbbqc;l^-=N}0DRn#6Nz3Gu%Q?ADgi(}2DEd454H
z)g}78+=($v>&cg5PmN4>7<G++0Znz<LfTW#hz0AaaiU*6l|Mj*mDmt}o76=>m5*ZR
UHK;!gI?%oX9b>~wxH_5n1oXpJsQ>@~

delta 30
mcmca%^T&L{9Er&rB(zvF8Il<?CjXZBJK0OBZ}TRZWsCsJy9~Mj

diff --git a/config.py b/config.py
index 974965e..104e961 100644
--- a/config.py
+++ b/config.py
@@ -4,7 +4,7 @@
 import logging
 from dotenv import find_dotenv, load_dotenv
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_openai import OpenAIEmbeddings
+from langchain_openai import AzureOpenAIEmbeddings, OpenAIEmbeddings
 from store_factory import get_vector_store
 
 load_dotenv(find_dotenv())
@@ -55,13 +55,17 @@ def get_env_variable(var_name: str, default_value: str = None) -> str:
 ## Credentials 
 
 OPENAI_API_KEY = get_env_variable("OPENAI_API_KEY", "")
+AZURE_OPENAI_API_KEY = get_env_variable("AZURE_OPENAI_API_KEY", "")
+AZURE_OPENAI_ENDPOINT = get_env_variable("AZURE_OPENAI_ENDPOINT", "")
 HF_TOKEN = get_env_variable("HF_TOKEN", "")
 
 ## Embeddings
 
 def init_embeddings(provider, model):
     if provider == "openai":
-        return OpenAIEmbeddings(model=model)
+        return OpenAIEmbeddings(model=model, api_key=OPENAI_API_KEY)
+    elif provider == "azure":
+        return AzureOpenAIEmbeddings(model=model, api_key=AZURE_OPENAI_API_KEY) # AZURE_OPENAI_ENDPOINT is being grabbed from the environment
     elif provider == "huggingface":
         return HuggingFaceEmbeddings(model_name=model,  encode_kwargs={'normalize_embeddings': True})
     else:
@@ -71,6 +75,10 @@ def init_embeddings(provider, model):
 
 if EMBEDDINGS_PROVIDER == "openai":
     EMBEDDINGS_MODEL = get_env_variable("EMBEDDINGS_MODEL", "text-embedding-3-small")
+
+elif EMBEDDINGS_PROVIDER == "azure":
+    EMBEDDINGS_MODEL = get_env_variable("EMBEDDINGS_MODEL", "text-embedding-3-small")
+
 elif EMBEDDINGS_PROVIDER == "huggingface":
     EMBEDDINGS_MODEL = get_env_variable("EMBEDDINGS_MODEL", "sangmini/msmarco-cotmae-MiniLM-L12_en-ko-ja")
 else: