Fix mypy issues

Author: NathalieCharbel

src/neo4j_graphrag/experimental/components/schema.py

@@ -20,7 +20,6 @@
 from pathlib import Path
 
 from pydantic import BaseModel, ValidationError, model_validator, validate_call
-from requests.exceptions import InvalidJSONError
 from typing_extensions import Self
 
 from neo4j_graphrag.exceptions import SchemaValidationError
@@ -336,10 +335,10 @@ async def run(
         return self.create_schema_model(entities, relations, potential_schema)
 
 
-class SchemaFromText(SchemaBuilder):
+class SchemaFromText(Component):
     """
-    A builder class for constructing SchemaConfig objects from the output of an LLM after
-     automatic schema extraction from text.
+    A component for constructing SchemaConfig objects from the output of an LLM after
+    automatic schema extraction from text.
     """
 
     def __init__(
@@ -348,7 +347,6 @@ def __init__(
         prompt_template: Optional[PromptTemplate] = None,
         llm_params: Optional[Dict[str, Any]] = None,
     ) -> None:
-        super().__init__()
         self._llm: LLMInterface = llm
         self._prompt_template: PromptTemplate = (
             prompt_template or SchemaExtractionTemplate()
@@ -369,7 +367,7 @@ async def run(self, text: str, examples: str = "", **kwargs: Any) -> SchemaConfi
         """
         prompt: str = self._prompt_template.format(text=text, examples=examples)
 
-        response = await self._llm.invoke(prompt, **self._llm_params)
+        response = await self._llm.ainvoke(prompt, **self._llm_params)
         content: str = (
             response
             if isinstance(response, str)
@@ -381,8 +379,12 @@ async def run(self, text: str, examples: str = "", **kwargs: Any) -> SchemaConfi
         except json.JSONDecodeError as exc:
             raise ValueError("LLM response is not valid JSON.") from exc
 
-        extracted_entities: List[dict] = extracted_schema.get("entities", [])
-        extracted_relations: Optional[List[dict]] = extracted_schema.get("relations")
+        extracted_entities: List[Dict[str, Any]] = (
+            extracted_schema.get("entities") or []
+        )
+        extracted_relations: Optional[List[Dict[str, Any]]] = extracted_schema.get(
+            "relations"
+        )
         potential_schema: Optional[List[Tuple[str, str, str]]] = extracted_schema.get(
             "potential_schema"
         )
@@ -394,7 +396,7 @@ async def run(self, text: str, examples: str = "", **kwargs: Any) -> SchemaConfi
             else None
         )
 
-        return await super().run(
+        return SchemaBuilder.create_schema_model(
             entities=entities,
             relations=relations,
             potential_schema=potential_schema,

src/neo4j_graphrag/experimental/pipeline/config/template_pipeline/simple_kg_builder.py

@@ -81,7 +81,7 @@ class SimpleKGPipelineConfig(TemplatePipelineConfig):
     entities: Sequence[EntityInputType] = []
     relations: Sequence[RelationInputType] = []
     potential_schema: Optional[list[tuple[str, str, str]]] = None
-    schema: Optional[Union[SchemaConfig, dict[str, list]]] = None
+    schema: Optional[Union[SchemaConfig, dict[str, list[Any]]]] = None
     enforce_schema: SchemaEnforcementMode = SchemaEnforcementMode.NONE
     on_error: OnError = OnError.IGNORE
     prompt_template: Union[ERExtractionTemplate, str] = ERExtractionTemplate()

src/neo4j_graphrag/experimental/pipeline/kg_builder.py

@@ -15,7 +15,7 @@
 
 from __future__ import annotations
 
-from typing import List, Optional, Sequence, Union
+from typing import List, Optional, Sequence, Union, Any
 import logging
 
 import neo4j
@@ -92,7 +92,7 @@ def __init__(
         entities: Optional[Sequence[EntityInputType]] = None,
         relations: Optional[Sequence[RelationInputType]] = None,
         potential_schema: Optional[List[tuple[str, str, str]]] = None,
-        schema: Optional[Union[SchemaConfig, dict[str, list]]] = None,
+        schema: Optional[Union[SchemaConfig, dict[str, list[Any]]]] = None,
         enforce_schema: str = "NONE",
         from_pdf: bool = True,
         text_splitter: Optional[TextSplitter] = None,

tests/unit/experimental/components/test_schema.py

@@ -110,7 +110,7 @@ def schema_config(
     valid_entities: list[SchemaEntity],
     valid_relations: list[SchemaRelation],
     potential_schema: list[tuple[str, str, str]],
-):
+) -> SchemaConfig:
     return schema_builder.create_schema_model(
         valid_entities, valid_relations, potential_schema
     )
@@ -445,14 +445,14 @@ def test_create_schema_model_missing_relations(
 
 
 @pytest.fixture
-def mock_llm():
+def mock_llm() -> AsyncMock:
     mock = AsyncMock()
     mock.invoke = AsyncMock()
     return mock
 
 
 @pytest.fixture
-def valid_schema_json():
+def valid_schema_json() -> str:
     return """
     {
         "entities": [
@@ -485,7 +485,7 @@ def valid_schema_json():
 
 
 @pytest.fixture
-def invalid_schema_json():
+def invalid_schema_json() -> str:
     return """
     {
         "entities": [
@@ -499,14 +499,14 @@ def invalid_schema_json():
 
 
 @pytest.fixture
-def schema_from_text(mock_llm):
+def schema_from_text(mock_llm: AsyncMock) -> SchemaFromText:
     return SchemaFromText(llm=mock_llm)
 
 
 @pytest.mark.asyncio
 async def test_schema_from_text_run_valid_response(
-    schema_from_text, mock_llm, valid_schema_json
-):
+    schema_from_text: SchemaFromText, mock_llm: AsyncMock, valid_schema_json: str
+) -> None:
     # configure the mock LLM to return a valid schema JSON
     mock_llm.invoke.return_value = valid_schema_json
 
@@ -534,8 +534,8 @@ async def test_schema_from_text_run_valid_response(
 
 @pytest.mark.asyncio
 async def test_schema_from_text_run_invalid_json(
-    schema_from_text, mock_llm, invalid_schema_json
-):
+    schema_from_text: SchemaFromText, mock_llm: AsyncMock, invalid_schema_json: str
+) -> None:
     # configure the mock LLM to return invalid JSON
     mock_llm.invoke.return_value = invalid_schema_json
 
@@ -547,7 +547,9 @@ async def test_schema_from_text_run_invalid_json(
 
 
 @pytest.mark.asyncio
-async def test_schema_from_text_custom_template(mock_llm, valid_schema_json):
+async def test_schema_from_text_custom_template(
+    mock_llm: AsyncMock, valid_schema_json: str
+) -> None:
     # create a  custom template
     custom_prompt = "This is a custom prompt with text: {text}"
     custom_template = PromptTemplate(template=custom_prompt, expected_inputs=["text"])
@@ -567,7 +569,9 @@ async def test_schema_from_text_custom_template(mock_llm, valid_schema_json):
 
 
 @pytest.mark.asyncio
-async def test_schema_from_text_llm_params(mock_llm, valid_schema_json):
+async def test_schema_from_text_llm_params(
+    mock_llm: AsyncMock, valid_schema_json: str
+) -> None:
     # configure custom LLM parameters
     llm_params = {"temperature": 0.1, "max_tokens": 500}
 
@@ -588,7 +592,7 @@ async def test_schema_from_text_llm_params(mock_llm, valid_schema_json):
 
 
 @pytest.mark.asyncio
-async def test_schema_config_store_as_json(schema_config):
+async def test_schema_config_store_as_json(schema_config: SchemaConfig) -> None:
     with tempfile.TemporaryDirectory() as temp_dir:
         # create file path
         json_path = os.path.join(temp_dir, "schema.json")
@@ -614,7 +618,7 @@ async def test_schema_config_store_as_json(schema_config):
 
 
 @pytest.mark.asyncio
-async def test_schema_config_store_as_yaml(schema_config):
+async def test_schema_config_store_as_yaml(schema_config: SchemaConfig) -> None:
     with tempfile.TemporaryDirectory() as temp_dir:
         # Create file path
         yaml_path = os.path.join(temp_dir, "schema.yaml")
@@ -640,7 +644,7 @@ async def test_schema_config_store_as_yaml(schema_config):
 
 
 @pytest.mark.asyncio
-async def test_schema_config_from_file(schema_config):
+async def test_schema_config_from_file(schema_config: SchemaConfig) -> None:
     with tempfile.TemporaryDirectory() as temp_dir:
         # create file paths with different extensions
         json_path = os.path.join(temp_dir, "schema.json")