AzureOpenai access from behind company proxies. (#1459)

embedchain/docs/api-reference/advanced/configuration.mdx

@@ -55,6 +55,7 @@ embedder:
   config:
     model: 'text-embedding-ada-002'
     api_key: sk-xxx
+    http_client_proxies: http://testproxy.mem0.net:8000
 
 chunker:
   chunk_size: 2000
@@ -106,7 +107,8 @@ cache:
     "provider": "openai",
     "config": {
       "model": "text-embedding-ada-002",
-      "api_key": "sk-xxx"
+      "api_key": "sk-xxx",
+      "http_client_proxies": "http://testproxy.mem0.net:8000",
     }
   },
   "chunker": {
@@ -168,7 +170,8 @@ config = {
         'provider': 'openai',
         'config': {
             'model': 'text-embedding-ada-002',
-            'api_key': 'sk-xxx'
+            'api_key': 'sk-xxx',
+            "http_client_proxies": "http://testproxy.mem0.net:8000",
         }
     },
     'chunker': {
@@ -236,6 +239,8 @@ Alright, let's dive into what each key means in the yaml config above:
         - `title` (String): The title for the embedding model for Google Embedder.
         - `task_type` (String): The task type for the embedding model for Google Embedder.
         - `model_kwargs` (Dict): Used to pass extra arguments to embedders.
+        - `http_client_proxies` (Dict | String): The proxy server settings used to create `self.http_client` using `httpx.Client(proxies=http_client_proxies)`
+        - `http_async_client_proxies` (Dict | String): The proxy server settings for async calls used to create `self.http_async_client` using `httpx.AsyncClient(proxies=http_async_client_proxies)`
 5. `chunker` Section:
     - `chunk_size` (Integer): The size of each chunk of text that is sent to the language model.
     - `chunk_overlap` (Integer): The amount of overlap between each chunk of text.

embedchain/embedchain/config/embedder/base.py

@@ -1,4 +1,6 @@
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
+
+import httpx
 
 from embedchain.helpers.json_serializable import register_deserializable
 
@@ -14,6 +16,8 @@ def __init__(
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
         model_kwargs: Optional[Dict[str, Any]] = None,
+        http_client_proxies: Optional[Union[Dict, str]] = None,
+        http_async_client_proxies: Optional[Union[Dict, str]] = None,
     ):
         """
         Initialize a new instance of an embedder config class.
@@ -32,6 +36,11 @@ def __init__(
         :type api_base: Optional[str], optional
         :param model_kwargs: key-value arguments for the embedding model, defaults a dict inside init.
         :type model_kwargs: Optional[Dict[str, Any]], defaults a dict inside init.
+        :param http_client_proxies: The proxy server settings used to create self.http_client, defaults to None
+        :type http_client_proxies: Optional[Dict | str], optional
+        :param http_async_client_proxies: The proxy server settings for async calls used to create
+        self.http_async_client, defaults to None
+        :type http_async_client_proxies: Optional[Dict | str], optional
         """
         self.model = model
         self.deployment_name = deployment_name
@@ -40,3 +49,7 @@ def __init__(
         self.api_key = api_key
         self.api_base = api_base
         self.model_kwargs = model_kwargs or {}
+        self.http_client = httpx.Client(proxies=http_client_proxies) if http_client_proxies else None
+        self.http_async_client = (
+            httpx.AsyncClient(proxies=http_async_client_proxies) if http_async_client_proxies else None
+        )

embedchain/embedchain/embedder/azure_openai.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from langchain_community.embeddings import AzureOpenAIEmbeddings
+from langchain_openai import AzureOpenAIEmbeddings
 
 from embedchain.config import BaseEmbedderConfig
 from embedchain.embedder.base import BaseEmbedder
@@ -14,7 +14,11 @@ def __init__(self, config: Optional[BaseEmbedderConfig] = None):
         if self.config.model is None:
             self.config.model = "text-embedding-ada-002"
 
-        embeddings = AzureOpenAIEmbeddings(deployment=self.config.deployment_name)
+        embeddings = AzureOpenAIEmbeddings(
+            deployment=self.config.deployment_name,
+            http_client=self.config.http_client,
+            http_async_client=self.config.http_async_client,
+        )
         embedding_fn = BaseEmbedder._langchain_default_concept(embeddings)
 
         self.set_embedding_fn(embedding_fn=embedding_fn)

embedchain/embedchain/llm/azure_openai.py

@@ -30,6 +30,8 @@ def _get_answer(prompt: str, config: BaseLlmConfig) -> str:
             temperature=config.temperature,
             max_tokens=config.max_tokens,
             streaming=config.stream,
+            http_client=config.http_client,
+            http_async_client=config.http_async_client,
         )
 
         if config.top_p and config.top_p != 1:

embedchain/embedchain/utils/misc.py

@@ -479,6 +479,8 @@ def validate_config(config_data):
                     Optional("base_url"): str,
                     Optional("endpoint"): str,
                     Optional("model_kwargs"): dict,
+                    Optional("http_client_proxies"): Or(str, dict),
+                    Optional("http_async_client_proxies"): Or(str, dict),
                 },
             },
             Optional("embedding_model"): {

embedchain/tests/embedder/test_azure_openai_embedder.py

@@ -0,0 +1,52 @@
+from unittest.mock import patch, Mock
+
+import httpx
+
+from embedchain.config import BaseEmbedderConfig
+from embedchain.embedder.azure_openai import AzureOpenAIEmbedder
+
+
+def test_azure_openai_embedder_with_http_client(monkeypatch):
+    mock_http_client = Mock(spec=httpx.Client)
+    mock_http_client_instance = Mock(spec=httpx.Client)
+    mock_http_client.return_value = mock_http_client_instance
+
+    with patch("embedchain.embedder.azure_openai.AzureOpenAIEmbeddings") as mock_embeddings, patch(
+        "httpx.Client", new=mock_http_client
+    ) as mock_http_client:
+        config = BaseEmbedderConfig(
+            deployment_name="text-embedding-ada-002",
+            http_client_proxies="http://testproxy.mem0.net:8000",
+        )
+
+        _ = AzureOpenAIEmbedder(config=config)
+
+        mock_embeddings.assert_called_once_with(
+            deployment="text-embedding-ada-002",
+            http_client=mock_http_client_instance,
+            http_async_client=None,
+        )
+        mock_http_client.assert_called_once_with(proxies="http://testproxy.mem0.net:8000")
+
+
+def test_azure_openai_embedder_with_http_async_client(monkeypatch):
+    mock_http_async_client = Mock(spec=httpx.AsyncClient)
+    mock_http_async_client_instance = Mock(spec=httpx.AsyncClient)
+    mock_http_async_client.return_value = mock_http_async_client_instance
+
+    with patch("embedchain.embedder.azure_openai.AzureOpenAIEmbeddings") as mock_embeddings, patch(
+        "httpx.AsyncClient", new=mock_http_async_client
+    ) as mock_http_async_client:
+        config = BaseEmbedderConfig(
+            deployment_name="text-embedding-ada-002",
+            http_async_client_proxies={"http://": "http://testproxy.mem0.net:8000"},
+        )
+
+        _ = AzureOpenAIEmbedder(config=config)
+
+        mock_embeddings.assert_called_once_with(
+            deployment="text-embedding-ada-002",
+            http_client=None,
+            http_async_client=mock_http_async_client_instance,
+        )
+        mock_http_async_client.assert_called_once_with(proxies={"http://": "http://testproxy.mem0.net:8000"})

embedchain/tests/llm/test_azure_openai.py

@@ -1,5 +1,6 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import Mock, MagicMock, patch
 
+import httpx
 import pytest
 from langchain.schema import HumanMessage, SystemMessage
 
@@ -43,6 +44,8 @@ def test_get_answer(azure_openai_llm):
             temperature=azure_openai_llm.config.temperature,
             max_tokens=azure_openai_llm.config.max_tokens,
             streaming=azure_openai_llm.config.stream,
+            http_client=None,
+            http_async_client=None,
         )
 
 
@@ -84,4 +87,78 @@ def test_with_api_version():
             temperature=0.7,
             max_tokens=50,
             streaming=False,
+            http_client=None,
+            http_async_client=None,
         )
+
+
+def test_get_llm_model_answer_with_http_client_proxies():
+    mock_http_client = Mock(spec=httpx.Client)
+    mock_http_client_instance = Mock(spec=httpx.Client)
+    mock_http_client.return_value = mock_http_client_instance
+
+    with patch("langchain_openai.AzureChatOpenAI") as mock_chat, patch(
+        "httpx.Client", new=mock_http_client
+    ) as mock_http_client:
+        mock_chat.return_value.invoke.return_value.content = "Mocked response"
+
+        config = BaseLlmConfig(
+            deployment_name="azure_deployment",
+            temperature=0.7,
+            max_tokens=50,
+            stream=False,
+            system_prompt="System prompt",
+            model="gpt-3.5-turbo",
+            http_client_proxies="http://testproxy.mem0.net:8000",
+        )
+
+        llm = AzureOpenAILlm(config)
+        llm.get_llm_model_answer("Test query")
+
+        mock_chat.assert_called_once_with(
+            deployment_name="azure_deployment",
+            openai_api_version="2024-02-01",
+            model_name="gpt-3.5-turbo",
+            temperature=0.7,
+            max_tokens=50,
+            streaming=False,
+            http_client=mock_http_client_instance,
+            http_async_client=None,
+        )
+        mock_http_client.assert_called_once_with(proxies="http://testproxy.mem0.net:8000")
+
+
+def test_get_llm_model_answer_with_http_async_client_proxies():
+    mock_http_async_client = Mock(spec=httpx.AsyncClient)
+    mock_http_async_client_instance = Mock(spec=httpx.AsyncClient)
+    mock_http_async_client.return_value = mock_http_async_client_instance
+
+    with patch("langchain_openai.AzureChatOpenAI") as mock_chat, patch(
+        "httpx.AsyncClient", new=mock_http_async_client
+    ) as mock_http_async_client:
+        mock_chat.return_value.invoke.return_value.content = "Mocked response"
+
+        config = BaseLlmConfig(
+            deployment_name="azure_deployment",
+            temperature=0.7,
+            max_tokens=50,
+            stream=False,
+            system_prompt="System prompt",
+            model="gpt-3.5-turbo",
+            http_async_client_proxies={"http://": "http://testproxy.mem0.net:8000"},
+        )
+
+        llm = AzureOpenAILlm(config)
+        llm.get_llm_model_answer("Test query")
+
+        mock_chat.assert_called_once_with(
+            deployment_name="azure_deployment",
+            openai_api_version="2024-02-01",
+            model_name="gpt-3.5-turbo",
+            temperature=0.7,
+            max_tokens=50,
+            streaming=False,
+            http_client=None,
+            http_async_client=mock_http_async_client_instance,
+        )
+        mock_http_async_client.assert_called_once_with(proxies={"http://": "http://testproxy.mem0.net:8000"})