Fix tests

neo4j · Dec 4, 2024 · d192667 · d192667
1 parent b070f6b
commit d192667
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 47 deletions.
diff --git a/src/neo4j_graphrag/experimental/components/schema.py b/src/neo4j_graphrag/experimental/components/schema.py
@@ -14,9 +14,10 @@
 #  limitations under the License.
 from __future__ import annotations
 
-from typing import Any, Dict, List, Literal, Optional, Self, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from pydantic import BaseModel, ValidationError, model_validator, validate_call
+from typing_extensions import Self
 
 from neo4j_graphrag.exceptions import SchemaValidationError
 from neo4j_graphrag.experimental.pipeline.component import Component, DataModel

diff --git a/src/neo4j_graphrag/experimental/pipeline/config/config_poc.py b/src/neo4j_graphrag/experimental/pipeline/config/config_poc.py
@@ -14,7 +14,6 @@
     Generic,
     Literal,
     Optional,
-    Self,
     TypeVar,
     Union,
 )
@@ -31,6 +30,7 @@
     field_validator,
 )
 from pydantic.v1.utils import deep_update
+from typing_extensions import Self
 
 from neo4j_graphrag.embeddings import Embedder
 from neo4j_graphrag.experimental.components.embedder import TextChunkEmbedder
@@ -742,43 +742,3 @@ async def run(self, data: dict[str, Any]) -> PipelineResult:
         else:
             run_param = deep_update(self.run_params, data)
         return await self.pipeline.run(data=run_param)
-
-
-if __name__ == "__main__":
-    from dotenv import load_dotenv
-
-    load_dotenv()
-
-    import asyncio
-
-    file_path = "examples/customize/build_graph/pipeline/pipeline_config.json"
-    runner = PipelineRunner.from_config_file(file_path)
-    print(runner)
-    # print(asyncio.run(runner.run({"splitter": {"text": "blabla"}})))
-
-    config = SimpleKGPipelineConfig.model_validate(
-        {
-            "template_": PipelineType.SIMPLE_KG_PIPELINE.value,
-            "neo4j_config": neo4j.GraphDatabase.driver("bolt://", auth=("", "")),
-            "from_pdf": True,
-        }
-    )
-    print(config)
-    runner = PipelineRunner.from_config(config)
-    print(runner.pipeline._nodes)
-
-
-"""
-        {
-            "name_": "embedder",
-            "class_": "embedder.TextChunkEmbedder",
-            "params_": {
-                "embedder": {
-                    "resolver_": "CONFIG_ARRAY",
-                    "array_": "embedder_config",
-                    "name_": "openai"
-                }
-            }
-        },
-
-"""
diff --git a/src/neo4j_graphrag/experimental/pipeline/kg_builder.py b/src/neo4j_graphrag/experimental/pipeline/kg_builder.py
@@ -21,10 +21,6 @@
 from pydantic import ValidationError
 
 from neo4j_graphrag.embeddings import Embedder
-from neo4j_graphrag.experimental.components.schema import (
-    SchemaEntity,
-    SchemaRelation,
-)
 from neo4j_graphrag.experimental.components.types import LexicalGraphConfig
 from neo4j_graphrag.experimental.pipeline.config.config_poc import (
     PipelineRunner,

diff --git a/tests/unit/experimental/pipeline/test_kg_builder.py b/tests/unit/experimental/pipeline/test_kg_builder.py
@@ -197,7 +197,7 @@ def test_simple_kg_pipeline_on_error_invalid_value() -> None:
     driver = MagicMock(spec=neo4j.Driver)
     embedder = MagicMock(spec=Embedder)
 
-    with pytest.raises(PipelineDefinitionError) as exc_info:
+    with pytest.raises(PipelineDefinitionError):
         SimpleKGPipeline(
             llm=llm,
             driver=driver,