Backport PR jupyterlab#628: Add OpenAI text-embedding-3-small, -large…

… models (jupyterlab#629) Co-authored-by: Jason Weill <[email protected]>
JasonWeill · Feb 7, 2024 · bd52241 · bd52241
1 parent 04ef454
commit bd52241
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/embedding_providers.py
@@ -71,7 +71,11 @@ def __init__(self, *args, **kwargs):
 class OpenAIEmbeddingsProvider(BaseEmbeddingsProvider, OpenAIEmbeddings):
     id = "openai"
     name = "OpenAI"
-    models = ["text-embedding-ada-002"]
+    models = [
+        "text-embedding-ada-002",
+        "text-embedding-3-small",
+        "text-embedding-3-large",
+    ]
     model_id_key = "model"
     pypi_package_deps = ["openai"]
     auth_strategy = EnvAuthStrategy(name="OPENAI_API_KEY")