Fix init class of OpenAIEmbeddings

Author: willtai

examples/graphrag.py

@@ -19,7 +19,7 @@
 URI = "neo4j://localhost:7687"
 AUTH = ("neo4j", "password")
 DATABASE = "neo4j"
-INDEX = "moviePlotsEmbedding"
+INDEX = "vector-index-name"
 
 
 # setup logger config

src/neo4j_genai/embeddings/openai.py

@@ -6,7 +6,7 @@
 
 
 class OpenAIEmbeddings(Embedder):
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, model: str = "text-embedding-ada-002") -> None:
         try:
             import openai
         except ImportError:
@@ -15,10 +15,11 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
                 "Please install it with `pip install openai`."
             )
 
-        self.model = openai.OpenAI(*args, **kwargs)
+        self.openai_model = openai.OpenAI()
+        self.model = model
 
-    def embed_query(
-        self, text: str, model: str = "text-embedding-ada-002", **kwargs: Any
-    ) -> list[float]:
-        response = self.model.embeddings.create(input=text, model=model, **kwargs)
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+        response = self.openai_model.embeddings.create(
+            input=text, model=self.model, **kwargs
+        )
         return response.data[0].embedding