doc/search_methods (neo4j#68)

* Improve search method documentation

* Merge with upstream/main

* Fix import for older versions of Python

* Ruff

* Rename method according to discussion

* Update docstring in the right method

* Update CHANGELOG

Author: stellasia

CHANGELOG.md

@@ -11,11 +11,13 @@
 -   Added PromptTemplate and RagTemplate for customizable prompt generation.
 -   Added LLMInterface with implementation for OpenAI LLM.
 -   Updated project configuration to support multiple Python versions (3.8 to 3.12) in CI workflows.
+-   Improved developer experience by copying the docstring from the `Retriever.get_search_results` method to the `Retriever.search` method
 
 ### Changed
 -   Refactored import paths for retrievers to neo4j_genai.retrievers.
 -   Implemented exception chaining for all re-raised exceptions to improve stack trace readability.
 -   Made error messages in `index.py` more consistent.
+-   Renamed `Retriever._get_search_results` to `Retriever.get_search_results`
 
 ## 0.2.0
 

docs/source/api.rst

@@ -18,27 +18,27 @@ VectorRetriever
 ===============
 
 .. autoclass:: neo4j_genai.retrievers.vector.VectorRetriever
-    :members:
+    :members: search
 
 VectorCypherRetriever
 =====================
 
 .. autoclass:: neo4j_genai.retrievers.vector.VectorCypherRetriever
-   :members:
+    :members: search
 
 
 HybridRetriever
 ===============
 
 .. autoclass:: neo4j_genai.retrievers.hybrid.HybridRetriever
-   :members:
+    :members: search
 
 
 HybridCypherRetriever
 =====================
 
 .. autoclass:: neo4j_genai.retrievers.hybrid.HybridCypherRetriever
-   :members:
+    :members: search
 
 
 
@@ -53,14 +53,14 @@ WeaviateNeo4jRetriever
 ======================
 
 .. autoclass:: neo4j_genai.retrievers.external.weaviate.weaviate.WeaviateNeo4jRetriever
-   :members:
+    :members: search
 
 
 PineconeNeo4jRetriever
 ======================
 
 .. autoclass:: neo4j_genai.retrievers.external.pinecone.pinecone.PineconeNeo4jRetriever
-   :members:
+    :members: search
 
 
 ******

src/neo4j_genai/retrievers/base.py

@@ -13,15 +13,66 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 from __future__ import annotations
-from abc import ABC, abstractmethod
-from typing import Optional, Callable, Any
+import types
+import inspect
+from abc import ABC, abstractmethod, ABCMeta
+from typing import Optional, Callable, Any, TypeVar
+from typing_extensions import ParamSpec
+
 import neo4j
 
 from neo4j_genai.types import RawSearchResult, RetrieverResult, RetrieverResultItem
 from neo4j_genai.exceptions import Neo4jVersionError
 
+T = ParamSpec("T")
+P = TypeVar("P")
+
+
+def copy_function(f: Callable[T, P]) -> Callable[T, P]:
+    """Based on https://stackoverflow.com/a/30714299"""
+    g = types.FunctionType(
+        f.__code__,
+        f.__globals__,
+        name=f.__name__,
+        argdefs=f.__defaults__,
+        closure=f.__closure__,
+    )
+    # in case f was given attrs (note this dict is a shallow copy):
+    g.__dict__.update(f.__dict__)
+    return g
+
+
+class RetrieverMetaclass(ABCMeta):
+    """This metaclass is used to copy the docstring from the
+    `get_search_results` method, instantiated in all subclasses,
+    to the `search` method in the base class.
+    """
+
+    def __new__(
+        meta, name: str, bases: tuple[type, ...], attrs: dict[str, Any]
+    ) -> type:
+        if "search" in attrs:
+            # search method was explicitly overridden, do nothing
+            return type.__new__(meta, name, bases, attrs)
+        # otherwise, we copy the signature and doc of the get_search_results
+        # method to a copy of the search method
+        get_search_results_method = attrs.get("get_search_results")
+        search_method = None
+        for b in bases:
+            search_method = getattr(b, "search", None)
+            if search_method is not None:
+                break
+        if search_method and get_search_results_method:
+            new_search_method = copy_function(search_method)
+            new_search_method.__doc__ = get_search_results_method.__doc__
+            new_search_method.__signature__ = inspect.signature(  # type: ignore
+                get_search_results_method
+            )
+            attrs["search"] = new_search_method
+        return type.__new__(meta, name, bases, attrs)
+
 
-class Retriever(ABC):
+class Retriever(ABC, metaclass=RetrieverMetaclass):
     """
     Abstract class for Neo4j retrievers
     """
@@ -78,11 +129,11 @@ def _fetch_index_infos(self) -> None:
             raise Exception(f"No index with name {self.index_name} found") from e
 
     def search(self, *args: Any, **kwargs: Any) -> RetrieverResult:
+        """Search method. Call the `get_search_results` method that returns
+        a list of `neo4j.Record`, and format them using the function returned by
+        `get_result_formatter` to return `RetrieverResult`.
         """
-        Search method. Call the get_search_result method that returns
-        a list of neo4j.Record, and format them to return RetrieverResult.
-        """
-        raw_result = self._get_search_results(*args, **kwargs)
+        raw_result = self.get_search_results(*args, **kwargs)
         formatter = self.get_result_formatter()
         search_items = [formatter(record) for record in raw_result.records]
         metadata = raw_result.metadata or {}
@@ -93,7 +144,20 @@ def search(self, *args: Any, **kwargs: Any) -> RetrieverResult:
         )
 
     @abstractmethod
-    def _get_search_results(self, *args: Any, **kwargs: Any) -> RawSearchResult:
+    def get_search_results(self, *args: Any, **kwargs: Any) -> RawSearchResult:
+        """This method must be implemented in each child class. It will
+        receive the same parameters provided to the public interface via
+        the `search` method, after validation. It returns a `RawSearchResult`
+        object which comprises a list of `neo4j.Record` objects and an optional
+        `metadata` dictionary that can contain retriever-level information.
+
+        Note that, even though this method is not intended to be called from
+        outside the class, we make it public to make it clearer for the developers
+        that it should be implemented in child classes.
+
+        Returns:
+            RawSearchResult: List of Neo4j Records and optional metadata dict
+        """
         pass
 
     def get_result_formatter(self) -> Callable[[neo4j.Record], RetrieverResultItem]:
@@ -127,7 +191,7 @@ def __init__(
         self.id_property_neo4j = id_property_neo4j
 
     @abstractmethod
-    def _get_search_results(
+    def get_search_results(
         self,
         query_vector: Optional[list[float]] = None,
         query_text: Optional[str] = None,
@@ -137,7 +201,7 @@ def _get_search_results(
         """
 
         Returns:
-                list[neo4j.Record]: List of Neo4j Records
+                RawSearchResult: List of Neo4j Records and optional metadata dict
 
         """
         pass

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

@@ -129,7 +129,7 @@ def __init__(
         self.retrieval_query = validated_data.retrieval_query
         self.result_formatter = validated_data.result_formatter
 
-    def _get_search_results(
+    def get_search_results(
         self,
         query_vector: Optional[list[float]] = None,
         query_text: Optional[str] = None,

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

@@ -117,7 +117,7 @@ def __init__(
         self.retrieval_query = validated_data.retrieval_query
         self.result_formatter = validated_data.result_formatter
 
-    def _get_search_results(
+    def get_search_results(
         self,
         query_vector: Optional[list[float]] = None,
         query_text: Optional[str] = None,

src/neo4j_genai/retrievers/hybrid.py

@@ -115,7 +115,7 @@ def default_record_formatter(self, record: neo4j.Record) -> RetrieverResultItem:
             metadata=metadata,
         )
 
-    def _get_search_results(
+    def get_search_results(
         self,
         query_text: str,
         query_vector: Optional[list[float]] = None,
@@ -243,7 +243,7 @@ def __init__(
         )
         self.result_formatter = result_formatter
 
-    def _get_search_results(
+    def get_search_results(
         self,
         query_text: str,
         query_vector: Optional[list[float]] = None,

src/neo4j_genai/retrievers/text2cypher.py

@@ -95,7 +95,7 @@ def __init__(
                 f"Failed to fetch schema for Text2CypherRetriever: {error_message}"
             ) from e
 
-    def _get_search_results(
+    def get_search_results(
         self,
         query_text: str,
     ) -> RawSearchResult:

src/neo4j_genai/retrievers/vector.py

@@ -118,7 +118,7 @@ def default_record_formatter(self, record: neo4j.Record) -> RetrieverResultItem:
             metadata=metadata,
         )
 
-    def _get_search_results(
+    def get_search_results(
         self,
         query_vector: Optional[list[float]] = None,
         query_text: Optional[str] = None,
@@ -245,7 +245,7 @@ def __init__(
         self._embedding_dimension = None
         self._fetch_index_infos()
 
-    def _get_search_results(
+    def get_search_results(
         self,
         query_vector: Optional[list[float]] = None,
         query_text: Optional[str] = None,

tests/unit/retrievers/test_base.py

@@ -14,12 +14,16 @@
 #  limitations under the License.
 from __future__ import annotations  # Reminder: May be removed after Python 3.9 is EOL.
 
+import inspect
+
 import pytest
 
+from typing import Union, Any
+from unittest.mock import MagicMock, patch
+
 from neo4j_genai.exceptions import Neo4jVersionError
 from neo4j_genai.retrievers.base import Retriever
-from unittest.mock import MagicMock
-from typing import Union, Any
+from neo4j_genai.types import RawSearchResult, RetrieverResult
 
 
 @pytest.mark.parametrize(
@@ -37,12 +41,62 @@ def test_retriever_version_support(
     expected_exception: Union[type[ValueError], None],
 ) -> None:
     class MockRetriever(Retriever):
-        def _get_search_results(self, *args: Any, **kwargs: Any) -> None:  # type: ignore
-            pass
+        def get_search_results(self, *args: Any, **kwargs: Any) -> RawSearchResult:
+            return RawSearchResult(records=[])
 
     driver.execute_query.return_value = [[{"versions": db_version}], None, None]
     if expected_exception:
         with pytest.raises(expected_exception):
             MockRetriever(driver=driver)
     else:
         MockRetriever(driver=driver)
+
+
+@patch("neo4j_genai.retrievers.base.Retriever._verify_version")
+def test_retriever_search_docstring_copied(
+    _verify_version_mock: MagicMock,
+    driver: MagicMock,
+) -> None:
+    class MockRetriever(Retriever):
+        def get_search_results(self, query: str, top_k: int = 10) -> RawSearchResult:
+            """My fabulous docstring"""
+            return RawSearchResult(records=[])
+
+    retriever = MockRetriever(driver=driver)
+    assert retriever.search.__doc__ == "My fabulous docstring"
+    signature = inspect.signature(retriever.search)
+    assert "query" in signature.parameters
+    query_param = signature.parameters["query"]
+    assert query_param.default == query_param.empty
+    assert query_param.annotation == "str"
+    assert "top_k" in signature.parameters
+    top_k_param = signature.parameters["top_k"]
+    assert top_k_param.default == 10
+    assert top_k_param.annotation == "int"
+
+
+@patch("neo4j_genai.retrievers.base.Retriever._verify_version")
+def test_retriever_search_docstring_unchanged(
+    _verify_version_mock: MagicMock,
+    driver: MagicMock,
+) -> None:
+    class MockRetrieverForNoise(Retriever):
+        def get_search_results(self, query: str, top_k: int = 10) -> RawSearchResult:
+            """My fabulous docstring"""
+            return RawSearchResult(records=[])
+
+    class MockRetriever(Retriever):
+        def get_search_results(self, *args: Any, **kwargs: Any) -> RawSearchResult:
+            return RawSearchResult(records=[])
+
+        def search(self, query: str, top_k: int = 10) -> RetrieverResult:
+            """My fabulous docstring that I do not want to be updated"""
+            return RetrieverResult(items=[])
+
+    assert MockRetrieverForNoise.search is not MockRetriever.search
+
+    retriever = MockRetriever(driver=driver)
+    assert (
+        retriever.search.__doc__
+        == "My fabulous docstring that I do not want to be updated"
+    )