feat: adds get_ontology_resolver + typed dict to hold params

cognee/modules/graph/utils/expand_with_nodes_and_edges.py

@@ -9,7 +9,7 @@ from cognee.modules.engine.utils import (
 )
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
-from cognee.modules.ontology.get_default_ontology_resolver import get_default_ontology_resolver
+from cognee.modules.ontology.get_ontology_resolver import get_ontology_resolver
 
 
 def _create_node_key(node_id: str, category: str) -> str:
@@ -321,7 +321,8 @@ def expand_with_nodes_and_edges(
         existing_edges_map = {}
 
     if ontology_resolver is None:
-        ontology_resolver = get_default_ontology_resolver()
+        config = get_ontology_resolver()
+        ontology_resolver = config["resolver"]
 
     added_nodes_map = {}
     added_ontology_nodes_map = {}

cognee/modules/ontology/get_default_ontology_resolver.py

@@ -1,18 +0,0 @@
-from typing import Optional
-
-from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
-from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
-from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
-
-
-def get_default_ontology_resolver(ontology_file: Optional[str] = None) -> BaseOntologyResolver:
-    """Get the default ontology resolver (RDFLib with fuzzy matching).
-
-    Args:
-        ontology_file: Optional path to ontology file
-
-    Returns:
-        Default RDFLib ontology resolver with fuzzy matching strategy
-    """
-    fuzzy_strategy = FuzzyMatchingStrategy()
-    return RDFLibOntologyResolver(ontology_file=ontology_file, matching_strategy=fuzzy_strategy)

cognee/modules/ontology/get_ontology_resolver.py

@@ -0,0 +1,36 @@
+from typing import Optional
+
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
+from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+from cognee.modules.ontology.matching_strategies import MatchingStrategy, FuzzyMatchingStrategy
+from cognee.modules.ontology.ontology_config import OntologyConfig
+
+
+def get_ontology_resolver(
+    resolver: Optional[BaseOntologyResolver] = None,
+    matching_strategy: Optional[MatchingStrategy] = None,
+) -> OntologyConfig:
+    """Get ontology resolver configuration with default or custom objects.
+
+    Args:
+        resolver: Optional pre-configured ontology resolver instance
+        matching_strategy: Optional matching strategy instance
+
+    Returns:
+        Ontology configuration with default RDFLib resolver and fuzzy matching strategy,
+        or custom objects if provided
+    """
+    config: OntologyConfig = {}
+
+    if resolver is not None:
+        config["resolver"] = resolver
+    else:
+        default_strategy = matching_strategy or FuzzyMatchingStrategy()
+        config["resolver"] = RDFLibOntologyResolver(
+            ontology_file=None, matching_strategy=default_strategy
+        )
+
+    if matching_strategy is not None and resolver is None:
+        config["matching_strategy"] = matching_strategy
+
+    return config

cognee/modules/ontology/ontology_config.py

@@ -0,0 +1,16 @@
+from typing import TypedDict, Optional
+
+from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
+from cognee.modules.ontology.matching_strategies import MatchingStrategy
+
+
+class OntologyConfig(TypedDict, total=False):
+    """Configuration for ontology resolver.
+
+    Attributes:
+        resolver: The ontology resolver instance to use
+        matching_strategy: The matching strategy to use
+    """
+
+    resolver: Optional[BaseOntologyResolver]
+    matching_strategy: Optional[MatchingStrategy]