Fix import and init issue with using OpenAIEmbeddings

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Next
 
+### Fixed
+-   Corrected initialization to allow specifying the embedding model name.
+-   Removed sentence_transformers from embeddings/__init__.py to avoid ImportError when the package is not installed.
+
 ## 0.3.0
 
 ### Added

docs/source/user_guide.rst

@@ -266,7 +266,7 @@ The `OpenAIEmbedder` was illustrated previously. Here is how to use the `Sentenc
 
 .. code:: python
 
-    from neo4j_genai.embeddings import SentenceTransformerEmbeddings
+    from neo4j_genai.embeddings.sentence_transformers import SentenceTransformerEmbeddings
 
     embedder = SentenceTransformerEmbeddings(model="all-MiniLM-L6-v2")  # Note: this is the default model
 

src/neo4j_genai/embeddings/__init__.py

@@ -1,5 +1,14 @@
-from .sentence_transformers import SentenceTransformerEmbeddings
-
-__all__ = [
-    "SentenceTransformerEmbeddings",
-]
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.

src/neo4j_genai/embeddings/openai.py

@@ -1,3 +1,18 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 from __future__ import annotations
 
 from typing import Any
@@ -6,7 +21,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 +30,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

src/neo4j_genai/embeddings/sentence_transformers.py

@@ -1,3 +1,18 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
 from typing import Any
 
 import numpy as np

tests/unit/embeddings/test_sentence_transformers.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 from neo4j_genai.embedder import Embedder
-from neo4j_genai.embeddings import SentenceTransformerEmbeddings
+from neo4j_genai.embeddings.sentence_transformers import SentenceTransformerEmbeddings
 
 
 @patch("sentence_transformers.SentenceTransformer")