From ec557f22bf7505a9cfa1952e9957ec3561f492bb Mon Sep 17 00:00:00 2001 From: Estelle Scifo Date: Fri, 26 Apr 2024 09:51:14 +0200 Subject: [PATCH] Simplified the Embedder return type - also fixes the retriever that was not reading the vector property (#21) --- src/neo4j_genai/embedder.py | 5 ++--- src/neo4j_genai/types.py | 4 ---- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/src/neo4j_genai/embedder.py b/src/neo4j_genai/embedder.py index be3e4c8fa..4cd695295 100644 --- a/src/neo4j_genai/embedder.py +++ b/src/neo4j_genai/embedder.py @@ -14,19 +14,18 @@ # limitations under the License. from abc import ABC, abstractmethod -from .types import EmbeddingVector class Embedder(ABC): """Interface for embedding models.""" @abstractmethod - def embed_query(self, text: str) -> EmbeddingVector: + def embed_query(self, text: str) -> list[float]: """Embed query text. Args: text (str): Text to convert to vector embedding Returns: - EmbeddingVector: A vector embedding. + list[float]: A vector embedding. """ diff --git a/src/neo4j_genai/types.py b/src/neo4j_genai/types.py index ebe8cd166..6db895fd1 100644 --- a/src/neo4j_genai/types.py +++ b/src/neo4j_genai/types.py @@ -23,10 +23,6 @@ class VectorSearchRecord(BaseModel): score: float -class EmbeddingVector(BaseModel): - vector: list[float] - - class IndexModel(BaseModel): driver: Any