Merge branch 'main' into kg-builder-strict-mode

Author: NathalieCharbel

CHANGELOG.md

@@ -5,6 +5,7 @@
 ### Added
 
 - Added optional schema enforcement as a validation layer after entity and relation extraction. 
+- Introduced SearchQueryParseError for handling invalid Lucene query strings in HybridRetriever and HybridCypherRetriever.
 
 ## 1.5.0
 
@@ -310,3 +311,4 @@
 
 -   Updated documentation to include new custom exceptions.
 -   Improved the use of Pydantic for input data validation for retriever objects.
+-   Fixed config loading after module reload (usage in jupyter notebooks)

README.md

@@ -371,6 +371,13 @@ When you're finished with your changes, create a pull request (PR) using the fol
 
 ## 🧪 Tests
 
+To be able to run all tests, all extra packages needs to be installed.  
+This is achieved by:
+
+```bash
+poetry install --all-extras
+```
+
 ### Unit Tests
 
 Install the project dependencies then run the following command to run the unit tests locally:

docs/source/api.rst

@@ -445,6 +445,16 @@ Errors
 
   * :class:`neo4j_graphrag.exceptions.LLMGenerationError`
 
+  * :class:`neo4j_graphrag.exceptions.SchemaValidationError`
+
+  * :class:`neo4j_graphrag.exceptions.PdfLoaderError`
+
+  * :class:`neo4j_graphrag.exceptions.PromptMissingPlaceholderError`
+
+  * :class:`neo4j_graphrag.exceptions.InvalidHybridSearchRankerError`
+
+  * :class:`neo4j_graphrag.exceptions.SearchQueryParseError`
+
   * :class:`neo4j_graphrag.experimental.pipeline.exceptions.PipelineDefinitionError`
 
   * :class:`neo4j_graphrag.experimental.pipeline.exceptions.PipelineMissingDependencyError`
@@ -559,6 +569,41 @@ LLMGenerationError
    :show-inheritance:
 
 
+SchemaValidationError
+=====================
+
+.. autoclass:: neo4j_graphrag.exceptions.SchemaValidationError
+   :show-inheritance:
+
+
+PdfLoaderError
+==============
+
+.. autoclass:: neo4j_graphrag.exceptions.PdfLoaderError
+   :show-inheritance:
+
+
+PromptMissingPlaceholderError
+=============================
+
+.. autoclass:: neo4j_graphrag.exceptions.PromptMissingPlaceholderError
+   :show-inheritance:
+
+
+InvalidHybridSearchRankerError
+==============================
+
+.. autoclass:: neo4j_graphrag.exceptions.InvalidHybridSearchRankerError
+   :show-inheritance:
+
+
+SearchQueryParseError
+=====================
+
+.. autoclass:: neo4j_graphrag.exceptions.SearchQueryParseError
+   :show-inheritance:
+
+
 PipelineDefinitionError
 =======================
 

docs/source/user_guide_rag.rst

@@ -223,7 +223,10 @@ it can be queried using the following:
 .. code:: python
 
     from neo4j_graphrag.llm import OllamaLLM
-    llm = OllamaLLM(model_name="orca-mini")
+    llm = OllamaLLM(
+        model_name="orca-mini",
+        # host="...",  # when using a remote server
+    )
     llm.invoke("say something")
 
 

examples/customize/embeddings/ollama_embeddings.py

@@ -6,6 +6,7 @@
 
 embeder = OllamaEmbeddings(
     model="<model_name>",
+    # host="...",  # if using a remote server
 )
 res = embeder.embed_query("my question")
 print(res[:10])

examples/customize/llms/ollama_llm.py

@@ -6,6 +6,7 @@
 
 llm = OllamaLLM(
     model_name="<model_name>",
+    # host="...",  # if using a remote server
 )
 res: LLMResponse = llm.invoke("What is the additive color model?")
 print(res.content)

poetry.lock

@@ -32,8 +32,8 @@ python = "^3.9.0"
 neo4j = "^5.17.0"
 pydantic = "^2.6.3"
 fsspec = "^2024.9.0"
-pypdf = "^4.3.1"
-json-repair = "^0.30.2"
+pypdf = "^5.1.0"
+json-repair = "^0.39.1"
 pyyaml = "^6.0.2"
 types-pyyaml = "^6.0.12.20240917"
 # optional deps
@@ -48,7 +48,7 @@ google-cloud-aiplatform = {version = "^1.66.0", optional = true }
 cohere = {version = "^5.9.0", optional = true}
 mistralai = {version = "^1.0.3", optional = true}
 qdrant-client = {version = "^1.11.3", optional = true}
-llama-index = {version = "^0.10.55", optional = true }
+llama-index = {version = "^0.12.0", optional = true }
 openai = {version = "^1.51.1", optional = true }
 anthropic = { version = "^0.36.0", optional = true}
 sentence-transformers = {version = "^3.0.0", optional = true }

pyproject.toml

@@ -128,3 +128,7 @@ class PromptMissingPlaceholderError(Neo4jGraphRagError):
 
 class InvalidHybridSearchRankerError(Neo4jGraphRagError):
     """Exception raised when an invalid ranker type for Hybrid Search is provided."""
+
+
+class SearchQueryParseError(Neo4jGraphRagError):
+    """Exception raised when there is a query parse error in the text search string."""

src/neo4j_graphrag/exceptions.py

@@ -21,6 +21,7 @@
 from pydantic import BaseModel
 
 from neo4j_graphrag.experimental.pipeline.exceptions import PipelineDefinitionError
+from neo4j_graphrag.utils.validation import issubclass_safe
 
 
 class DataModel(BaseModel):
@@ -52,7 +53,7 @@ def __new__(
                     f"The run method return type must be annotated in {name}"
                 )
             # the type hint must be a subclass of DataModel
-            if not issubclass(return_model, DataModel):
+            if not issubclass_safe(return_model, DataModel):
                 raise PipelineDefinitionError(
                     f"The run method must return a subclass of DataModel in {name}"
                 )

src/neo4j_graphrag/experimental/pipeline/component.py

@@ -56,6 +56,8 @@
     ParamConfig,
 )
 from neo4j_graphrag.llm import LLMInterface
+from neo4j_graphrag.utils.validation import issubclass_safe
+
 
 logger = logging.getLogger(__name__)
 
@@ -131,9 +133,9 @@ def parse(self, resolved_data: dict[str, Any] | None = None) -> T:
         self._global_data = resolved_data or {}
         logger.debug(f"OBJECT_CONFIG: parsing {self} using {resolved_data}")
         if self.class_ is None:
-            raise ValueError(f"`class_` is not required to parse object {self}")
+            raise ValueError(f"`class_` is required to parse object {self}")
         klass = self._get_class(self.class_, self.get_module())
-        if not issubclass(klass, self.get_interface()):
+        if not issubclass_safe(klass, self.get_interface()):
             raise ValueError(
                 f"Invalid class '{klass}'. Expected a subclass of '{self.get_interface()}'"
             )

src/neo4j_graphrag/experimental/pipeline/config/object_config.py

@@ -95,9 +95,11 @@ async def set_task_status(self, task_name: str, status: RunStatus) -> None:
         async with asyncio.Lock():
             current_status = await self.get_status_for_component(task_name)
             if status == current_status:
-                raise PipelineStatusUpdateError(f"Status is already '{status}'")
-            if status == RunStatus.RUNNING and current_status == RunStatus.DONE:
-                raise PipelineStatusUpdateError("Can't go from DONE to RUNNING")
+                raise PipelineStatusUpdateError(f"Status is already {status}")
+            if status not in current_status.possible_next_status():
+                raise PipelineStatusUpdateError(
+                    f"Can't go from {current_status} to {status}"
+                )
             return await self.pipeline.store.add_status_for_component(
                 self.run_id, task_name, status.value
             )

src/neo4j_graphrag/experimental/pipeline/orchestrator.py

@@ -54,6 +54,15 @@ class RunStatus(enum.Enum):
     RUNNING = "RUNNING"
     DONE = "DONE"
 
+    def possible_next_status(self) -> list[RunStatus]:
+        if self == RunStatus.UNKNOWN:
+            return [RunStatus.RUNNING]
+        if self == RunStatus.RUNNING:
+            return [RunStatus.DONE]
+        if self == RunStatus.DONE:
+            return []
+        return []
+
 
 class RunResult(BaseModel):
     status: RunStatus = RunStatus.DONE

src/neo4j_graphrag/experimental/pipeline/types.py

@@ -211,7 +211,7 @@ def get_search_results(
         for point in points:
             assert point.payload is not None
             result_tuples.append(
-                [f"{point.payload[self.id_property_external]}", point.score]
+                [point.payload.get(self.id_property_external, point.id), point.score]
             )
 
         search_query = get_match_query(

src/neo4j_graphrag/retrievers/external/qdrant/qdrant.py

@@ -26,6 +26,7 @@
     EmbeddingRequiredError,
     RetrieverInitializationError,
     SearchValidationError,
+    SearchQueryParseError,
 )
 from neo4j_graphrag.neo4j_queries import get_search_query
 from neo4j_graphrag.retrievers.base import Retriever
@@ -218,12 +219,19 @@ def get_search_results(
         logger.debug("HybridRetriever Cypher parameters: %s", sanitized_parameters)
         logger.debug("HybridRetriever Cypher query: %s", search_query)
 
-        records, _, _ = self.driver.execute_query(
-            search_query,
-            parameters,
-            database_=self.neo4j_database,
-            routing_=neo4j.RoutingControl.READ,
-        )
+        try:
+            records, _, _ = self.driver.execute_query(
+                search_query,
+                parameters,
+                database_=self.neo4j_database,
+                routing_=neo4j.RoutingControl.READ,
+            )
+        except neo4j.exceptions.ClientError as e:
+            if "org.apache.lucene.queryparser.classic.ParseException" in str(e):
+                raise SearchQueryParseError(
+                    f"Invalid Lucene query generated from query_text: {query_text}"
+                ) from e
+            raise
         return RawSearchResult(
             records=records,
         )
@@ -395,12 +403,19 @@ def get_search_results(
         logger.debug("HybridRetriever Cypher parameters: %s", sanitized_parameters)
         logger.debug("HybridRetriever Cypher query: %s", search_query)
 
-        records, _, _ = self.driver.execute_query(
-            search_query,
-            parameters,
-            database_=self.neo4j_database,
-            routing_=neo4j.RoutingControl.READ,
-        )
+        try:
+            records, _, _ = self.driver.execute_query(
+                search_query,
+                parameters,
+                database_=self.neo4j_database,
+                routing_=neo4j.RoutingControl.READ,
+            )
+        except neo4j.exceptions.ClientError as e:
+            if "org.apache.lucene.queryparser.classic.ParseException" in str(e):
+                raise SearchQueryParseError(
+                    f"Invalid Lucene query generated from query_text: {query_text}"
+                ) from e
+            raise
         return RawSearchResult(
             records=records,
         )

src/neo4j_graphrag/retrievers/hybrid.py

@@ -14,11 +14,32 @@
 #  limitations under the License.
 from __future__ import annotations
 
-from typing import Optional
+import importlib
+from typing import Optional, Tuple, Union, cast, Type
 
 
 def validate_search_query_input(
     query_text: Optional[str] = None, query_vector: Optional[list[float]] = None
 ) -> None:
     if not (bool(query_vector) ^ bool(query_text)):
         raise ValueError("You must provide exactly one of query_vector or query_text.")
+
+
+def issubclass_safe(
+    cls: Type[object], class_or_tuple: Union[Type[object], Tuple[Type[object]]]
+) -> bool:
+    if isinstance(class_or_tuple, tuple):
+        return any(issubclass_safe(cls, base) for base in class_or_tuple)
+
+    if issubclass(cls, class_or_tuple):
+        return True
+
+    # Handle case where module was reloaded
+    cls_module = importlib.import_module(cls.__module__)
+    # Get the latest version of the base class from the module
+    latest_base = getattr(cls_module, class_or_tuple.__name__, None)
+    latest_base = cast(Union[tuple[Type[object], ...], Type[object]], latest_base)
+    if issubclass(cls, latest_base):
+        return True
+
+    return False

src/neo4j_graphrag/utils/validation.py

@@ -12,6 +12,10 @@
 #  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.
+import importlib
+import sys
+from abc import ABC
+from typing import ClassVar
 from unittest.mock import patch
 
 import neo4j
@@ -58,6 +62,30 @@ def test_get_class_wrong_path() -> None:
         c._get_class("MyClass")
 
 
+class _MyClass:
+    def __init__(self, param: str) -> None:
+        self.param = param
+
+
+class _MyInterface(ABC): ...
+
+
+def test_parse_after_module_reload() -> None:
+    class MyClassConfig(ObjectConfig[_MyClass]):
+        DEFAULT_MODULE: ClassVar[str] = __name__
+        INTERFACE: ClassVar[type] = _MyClass
+
+    param_value = "value"
+    config = MyClassConfig.model_validate(
+        {"class_": f"{__name__}.{_MyClass.__name__}", "params_": {"param": param_value}}
+    )
+    importlib.reload(sys.modules[__name__])
+
+    my_obj = config.parse()
+    assert isinstance(my_obj, _MyClass)
+    assert my_obj.param == param_value
+
+
 def test_neo4j_driver_config() -> None:
     config = Neo4jDriverConfig.model_validate(
         {