Merge pull request #1635 from Imccccc/feature/embedding-with-retry

Embedding Improvement

autogpt/llm_utils.py

@@ -113,3 +113,35 @@ def create_chat_completion(
         raise RuntimeError(f"Failed to get response after {num_retries} retries")
 
     return response.choices[0].message["content"]
+
+
+def create_embedding_with_ada(text) -> list:
+    """Create a embedding with text-ada-002 using the OpenAI SDK"""
+    num_retries = 10
+    for attempt in range(num_retries):
+        backoff = 2 ** (attempt + 2)
+        try:
+            if CFG.use_azure:
+                return openai.Embedding.create(input=[text],
+                    engine=CFG.get_azure_deployment_id_for_model("text-embedding-ada-002"),
+                )["data"][0]["embedding"]
+            else:
+                return openai.Embedding.create(input=[text], model="text-embedding-ada-002")[
+                    "data"
+                ][0]["embedding"]
+        except RateLimitError:
+            pass
+        except APIError as e:
+            if e.http_status == 502:
+                pass
+            else:
+                raise
+            if attempt == num_retries - 1:
+                raise
+        if CFG.debug_mode:
+            print(
+                Fore.RED + "Error: ",
+                f"API Bad gateway. Waiting {backoff} seconds..." + Fore.RESET,
+            )
+        time.sleep(backoff)
+