Mypy

Author: stellasia

src/neo4j_graphrag/embeddings/base.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from __future__ import annotations
+from typing import Any
 
 from abc import ABC, abstractmethod
 
@@ -25,7 +26,7 @@ class Embedder(ABC):
 
     @abstractmethod
     def embed_query(
-        self, text: str, dimensions: int | None = None, **kwargs
+        self, text: str, dimensions: int | None = None, **kwargs: Any
     ) -> list[float]:
         """Embed query text.
 

src/neo4j_graphrag/embeddings/ollama.py

@@ -41,7 +41,7 @@ def __init__(self, model: str, **kwargs: Any) -> None:
         self.model = model
         self.client = ollama.Client(**kwargs)
 
-    def embed_query(self, text: str, **kwargs: Any) -> list[float]:
+    def embed_query(self, text: str, **kwargs: Any) -> list[float]:  # type: ignore[override]
         """
         Generate embeddings for a given query using an Ollama text embedding model.
 

src/neo4j_graphrag/embeddings/openai.py

@@ -18,6 +18,8 @@
 import abc
 from typing import TYPE_CHECKING, Any
 
+from openai import NotGiven
+
 from neo4j_graphrag.embeddings.base import Embedder
 
 if TYPE_CHECKING:
@@ -63,8 +65,9 @@ def embed_query(
 
             **kwargs (Any): Additional arguments to pass to the OpenAI embedding generation function.
         """
+        d = dimensions or NotGiven()
         response = self.client.embeddings.create(
-            input=text, model=self.model, dimensions=dimensions, **kwargs
+            input=text, model=self.model, dimensions=d, **kwargs
         )
         embedding: list[float] = response.data[0].embedding
         return embedding

src/neo4j_graphrag/embeddings/sentence_transformers.py

@@ -35,7 +35,7 @@ def __init__(
         self.np = np
         self.model = sentence_transformers.SentenceTransformer(model, *args, **kwargs)
 
-    def embed_query(self, text: str, **kwargs) -> Any:
+    def embed_query(self, text: str, **kwargs: Any) -> Any:  # type: ignore[override]
         result = self.model.encode([text])
         if isinstance(result, self.torch.Tensor) or isinstance(result, self.np.ndarray):
             return result.flatten().tolist()

src/neo4j_graphrag/embeddings/vertexai.py

@@ -42,7 +42,7 @@ def __init__(self, model: str = "text-embedding-004") -> None:
             )
         self.model = TextEmbeddingModel.from_pretrained(model)
 
-    def embed_query(
+    def embed_query(  # type: ignore[override]
         self,
         text: str,
         task_type: str = "RETRIEVAL_QUERY",
@@ -62,4 +62,4 @@ def embed_query(
         embeddings = self.model.get_embeddings(
             inputs, output_dimensionality=dimensions, **kwargs
         )
-        return embeddings[0].values  # type: ignore
+        return embeddings[0].values