Pydantic search model refactoring (#69)

* Refactored Pydantic search models

* Added VectorCypherSearchModel and HybridCypherSearchModel back in

* Minor Pinecone and Weaviate Pydantic model refactoring

* Removed unneeded imports

* Removed filters from parameters in vector search

* Changed List, Dict type hints to list, dict

Author: alexthomas93

src/neo4j_genai/retrievers/external/pinecone/pinecone.py

@@ -179,15 +179,14 @@ def get_search_results(
             RawSearchResult: The results of the search query as a list of neo4j.Record and an optional metadata dict
         """
 
-        pinecone_filters = kwargs.get("pinecone_filters")
+        pinecone_filter = kwargs.get("pinecone_filter")
 
         try:
             validated_data = PineconeSearchModel(
-                vector_index_name=self.index_name,
                 query_vector=query_vector,
                 query_text=query_text,
                 top_k=top_k,
-                pinecone_filter=pinecone_filters,
+                pinecone_filter=pinecone_filter,
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors()) from e

src/neo4j_genai/retrievers/external/pinecone/types.py

@@ -13,7 +13,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from __future__ import annotations
-from typing import Any, Callable, Optional
+
+from typing import Any, Callable, Optional, Union
 
 import neo4j
 from pinecone import Pinecone
@@ -27,7 +28,9 @@
 
 
 class PineconeSearchModel(VectorSearchModel):
-    pinecone_filter: Optional[dict[str, Any]] = None
+    pinecone_filter: Optional[
+        dict[str, Union[str, float, int, bool, list[Any], dict[Any, Any]]]
+    ] = None
 
 
 class PineconeClientModel(BaseModel):

src/neo4j_genai/retrievers/external/weaviate/types.py

@@ -13,21 +13,19 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from __future__ import annotations
-from typing import Optional, Any, Callable
+
+from typing import Callable, Optional
 
 import neo4j
 from pydantic import (
-    field_validator,
     BaseModel,
-    PositiveInt,
-    model_validator,
     ConfigDict,
+    field_validator,
 )
 from weaviate.client import WeaviateClient
 from weaviate.collections.classes.filters import _Filters
 
-from neo4j_genai.utils import validate_search_query_input
-from neo4j_genai.types import Neo4jDriverModel, EmbedderModel
+from neo4j_genai.types import EmbedderModel, Neo4jDriverModel, VectorSearchModel
 
 
 class WeaviateModel(BaseModel):
@@ -55,10 +53,7 @@ class WeaviateNeo4jRetrieverModel(BaseModel):
     result_formatter: Optional[Callable[[neo4j.Record], str]] = None
 
 
-class WeaviateNeo4jSearchModel(BaseModel):
-    top_k: PositiveInt = 5
-    query_vector: Optional[list[float]] = None
-    query_text: Optional[str] = None
+class WeaviateNeo4jSearchModel(VectorSearchModel):
     weaviate_filters: Optional[_Filters] = None
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -69,12 +64,3 @@ def check_weaviate_filters(cls, value: _Filters) -> _Filters:
                 "Provided filters need to be of type weaviate.collections.classes.filters._Filters"
             )
         return value
-
-    @model_validator(mode="before")
-    def check_query(cls, values: dict[str, Any]) -> dict[str, Any]:
-        """
-        Validates that one of either query_vector or query_text is provided exclusively.
-        """
-        query_vector, query_text = values.get("query_vector"), values.get("query_text")
-        validate_search_query_input(query_text, query_vector)
-        return values

src/neo4j_genai/retrievers/hybrid.py

@@ -13,31 +13,32 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from __future__ import annotations
-from typing import Optional, Any, Callable
+
+import logging
+from typing import Any, Callable, Optional
 
 import neo4j
 from pydantic import ValidationError
 
 from neo4j_genai.embedder import Embedder
 from neo4j_genai.exceptions import (
+    EmbeddingRequiredError,
     RetrieverInitializationError,
     SearchValidationError,
-    EmbeddingRequiredError,
 )
+from neo4j_genai.neo4j_queries import get_search_query
 from neo4j_genai.retrievers.base import Retriever
 from neo4j_genai.types import (
-    HybridSearchModel,
-    SearchType,
-    HybridCypherSearchModel,
-    Neo4jDriverModel,
     EmbedderModel,
-    HybridRetrieverModel,
     HybridCypherRetrieverModel,
+    HybridCypherSearchModel,
+    HybridRetrieverModel,
+    HybridSearchModel,
+    Neo4jDriverModel,
     RawSearchResult,
     RetrieverResultItem,
+    SearchType,
 )
-from neo4j_genai.neo4j_queries import get_search_query
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -146,16 +147,16 @@ def get_search_results(
         """
         try:
             validated_data = HybridSearchModel(
-                vector_index_name=self.vector_index_name,
-                fulltext_index_name=self.fulltext_index_name,
-                top_k=top_k,
                 query_vector=query_vector,
                 query_text=query_text,
+                top_k=top_k,
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors()) from e
 
         parameters = validated_data.model_dump(exclude_none=True)
+        parameters["vector_index_name"] = self.vector_index_name
+        parameters["fulltext_index_name"] = self.fulltext_index_name
 
         if query_text and not query_vector:
             if not self.embedder:
@@ -276,17 +277,17 @@ def get_search_results(
         """
         try:
             validated_data = HybridCypherSearchModel(
-                vector_index_name=self.vector_index_name,
-                fulltext_index_name=self.fulltext_index_name,
-                top_k=top_k,
                 query_vector=query_vector,
                 query_text=query_text,
+                top_k=top_k,
                 query_params=query_params,
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors()) from e
 
         parameters = validated_data.model_dump(exclude_none=True)
+        parameters["vector_index_name"] = self.vector_index_name
+        parameters["fulltext_index_name"] = self.fulltext_index_name
 
         if query_text and not query_vector:
             if not self.embedder:

src/neo4j_genai/retrievers/vector.py

@@ -13,32 +13,32 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from __future__ import annotations
-from typing import Optional, Any, Callable
+
+import logging
+from typing import Any, Callable, Optional
 
 import neo4j
+from pydantic import ValidationError
 
+from neo4j_genai.embedder import Embedder
 from neo4j_genai.exceptions import (
+    EmbeddingRequiredError,
     RetrieverInitializationError,
     SearchValidationError,
-    EmbeddingRequiredError,
 )
+from neo4j_genai.neo4j_queries import get_search_query
 from neo4j_genai.retrievers.base import Retriever
-from pydantic import ValidationError
-
-from neo4j_genai.embedder import Embedder
 from neo4j_genai.types import (
-    VectorSearchModel,
-    VectorCypherSearchModel,
-    SearchType,
-    Neo4jDriverModel,
     EmbedderModel,
-    VectorRetrieverModel,
-    VectorCypherRetrieverModel,
+    Neo4jDriverModel,
     RawSearchResult,
     RetrieverResultItem,
+    SearchType,
+    VectorCypherRetrieverModel,
+    VectorCypherSearchModel,
+    VectorRetrieverModel,
+    VectorSearchModel,
 )
-from neo4j_genai.neo4j_queries import get_search_query
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -146,15 +146,18 @@ def get_search_results(
         """
         try:
             validated_data = VectorSearchModel(
-                vector_index_name=self.index_name,
-                top_k=top_k,
                 query_vector=query_vector,
                 query_text=query_text,
+                top_k=top_k,
+                filters=filters,
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors()) from e
 
         parameters = validated_data.model_dump(exclude_none=True)
+        parameters["vector_index_name"] = self.index_name
+        if filters:
+            del parameters["filters"]
 
         if query_text:
             if not self.embedder:
@@ -275,16 +278,19 @@ def get_search_results(
         """
         try:
             validated_data = VectorCypherSearchModel(
-                vector_index_name=self.index_name,
-                top_k=top_k,
                 query_vector=query_vector,
                 query_text=query_text,
+                top_k=top_k,
                 query_params=query_params,
+                filters=filters,
             )
         except ValidationError as e:
             raise SearchValidationError(e.errors()) from e
 
         parameters = validated_data.model_dump(exclude_none=True)
+        parameters["vector_index_name"] = self.index_name
+        if filters:
+            del parameters["filters"]
 
         if query_text:
             if not self.embedder:

src/neo4j_genai/types.py

@@ -116,10 +116,10 @@ def check_node_properties_not_empty(cls, v: list[Any]) -> list[Any]:
 
 
 class VectorSearchModel(BaseModel):
-    vector_index_name: str
-    top_k: PositiveInt = 5
     query_vector: Optional[list[float]] = None
     query_text: Optional[str] = None
+    top_k: PositiveInt = 5
+    filters: Optional[dict[str, Any]] = None
 
     @model_validator(mode="before")
     def check_query(cls, values: dict[str, Any]) -> dict[str, Any]:
@@ -136,11 +136,9 @@ class VectorCypherSearchModel(VectorSearchModel):
 
 
 class HybridSearchModel(BaseModel):
-    vector_index_name: str
-    fulltext_index_name: str
     query_text: str
-    top_k: PositiveInt = 5
     query_vector: Optional[list[float]] = None
+    top_k: PositiveInt = 5
 
 
 class HybridCypherSearchModel(HybridSearchModel):