Pass in httpx client in Azure OpenAI embedding initialisation & increase default timeouts (#101)

* COnfigure timeputs for httpx clients and increase embedding timeout

* Update src/unstract/sdk/adapters/embedding/azure_open_ai/src/azure_open_ai.py

Co-authored-by: Ritwik G <[email protected]>
Signed-off-by: Gayathri <[email protected]>

* Add timeout changes to OpenAI

* Add model field for OpenAI

---------

Signed-off-by: Gayathri <[email protected]>
Co-authored-by: Ritwik G <[email protected]>

Author: gaya3-zipstack

src/unstract/sdk/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.48.2"
+__version__ = "0.49.0"
 
 
 def get_sdk_version():

src/unstract/sdk/adapters/embedding/azure_open_ai/src/azure_open_ai.py

@@ -1,6 +1,7 @@
 import os
 from typing import Any
 
+import httpx
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
 
@@ -18,7 +19,7 @@ class Constants:
     DEPLOYMENT_NAME = "deployment_name"
     API_TYPE = "azure"
     TIMEOUT = "timeout"
-    DEFAULT_TIMEOUT = 60
+    DEFAULT_TIMEOUT = 240
 
 
 class AzureOpenAI(EmbeddingAdapter):
@@ -59,6 +60,8 @@ def get_embedding_instance(self) -> BaseEmbedding:
                 config=self.config
             )
             timeout = int(self.config.get(Constants.TIMEOUT, Constants.DEFAULT_TIMEOUT))
+            httpx_timeout = httpx.Timeout(timeout, connect=60.0)
+            httpx_client = httpx.Client(timeout=httpx_timeout)
             embedding: BaseEmbedding = AzureOpenAIEmbedding(
                 model=str(self.config.get(Constants.MODEL)),
                 deployment_name=str(self.config.get(Constants.DEPLOYMENT_NAME)),
@@ -68,6 +71,7 @@ def get_embedding_instance(self) -> BaseEmbedding:
                 embed_batch_size=embedding_batch_size,
                 api_type=Constants.API_TYPE,
                 timeout=timeout,
+                http_client=httpx_client,
             )
             return embedding
         except Exception as e:

src/unstract/sdk/adapters/embedding/azure_open_ai/src/static/json_schema.json

@@ -59,7 +59,7 @@
       "minimum": 0,
       "multipleOf": 1,
       "title": "Timeout",
-      "default": 60,
+      "default": 240,
       "description": "Timeout for each request in seconds"
     }
   }

src/unstract/sdk/adapters/embedding/open_ai/src/open_ai.py

@@ -1,6 +1,7 @@
 import os
 from typing import Any
 
+import httpx
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.embeddings.openai import OpenAIEmbedding
 
@@ -11,12 +12,13 @@
 
 class Constants:
     API_KEY = "api_key"
+    MODEL = "model"
     API_BASE_VALUE = "https://api.openai.com/v1/"
     API_BASE_KEY = "api_base"
     ADAPTER_NAME = "adapter_name"
     API_TYPE = "openai"
     TIMEOUT = "timeout"
-    DEFAULT_TIMEOUT = 60
+    DEFAULT_TIMEOUT = 240
 
 
 class OpenAI(EmbeddingAdapter):
@@ -54,13 +56,17 @@ def get_json_schema() -> str:
     def get_embedding_instance(self) -> BaseEmbedding:
         try:
             timeout = int(self.config.get(Constants.TIMEOUT, Constants.DEFAULT_TIMEOUT))
+            httpx_timeout = httpx.Timeout(10.0, connect=60.0)
+            httpx_client = httpx.Client(timeout=httpx_timeout)
             embedding: BaseEmbedding = OpenAIEmbedding(
                 api_key=str(self.config.get(Constants.API_KEY)),
                 api_base=str(
                     self.config.get(Constants.API_BASE_KEY, Constants.API_BASE_VALUE)
                 ),
+                model=str(self.config.get(Constants.MODEL)),
                 api_type=Constants.API_TYPE,
                 timeout=timeout,
+                http_client=httpx_client,
             )
             return embedding
         except Exception as e:

src/unstract/sdk/adapters/embedding/open_ai/src/static/json_schema.json

@@ -12,6 +12,12 @@
       "default": "",
       "description": "Provide a unique name for this adapter instance. Example: openai-emb-1"
     },
+    "model": {
+      "type": "string",
+      "title": "Model",
+      "default": "text-embedding-ada-002",
+      "description": "Provide the name of the model."
+    },
     "api_key": {
       "type": "string",
       "title": "API Key",
@@ -36,7 +42,7 @@
       "minimum": 0,
       "multipleOf": 1,
       "title": "Timeout",
-      "default": 60,
+      "default": 240,
       "description": "Timeout in seconds"
     }
   }