From 7c046eafab20e8714ee985bb1cf9873c4e9ae3bf Mon Sep 17 00:00:00 2001
From: hajdul88 <52442977+hajdul88@users.noreply.github.com>
Date: Wed, 17 Sep 2025 13:28:11 +0200
Subject: [PATCH] feat: adds get_ontology_resolver + typed dict to hold params

---
 .../utils/expand_with_nodes_and_edges.py      |  5 +--
 .../ontology/get_default_ontology_resolver.py | 18 ----------
 .../modules/ontology/get_ontology_resolver.py | 36 +++++++++++++++++++
 cognee/modules/ontology/ontology_config.py    | 16 +++++++++
 4 files changed, 55 insertions(+), 20 deletions(-)
 delete mode 100644 cognee/modules/ontology/get_default_ontology_resolver.py
 create mode 100644 cognee/modules/ontology/get_ontology_resolver.py
 create mode 100644 cognee/modules/ontology/ontology_config.py

diff --git a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py b/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
index b3e8e8029..e18860744 100644
--- a/cognee/modules/graph/utils/expand_with_nodes_and_edges.py
+++ b/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 = {}
diff --git a/cognee/modules/ontology/get_default_ontology_resolver.py b/cognee/modules/ontology/get_default_ontology_resolver.py
deleted file mode 100644
index 9dc5c59ba..000000000
--- a/cognee/modules/ontology/get_default_ontology_resolver.py
+++ /dev/null
@@ -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)
diff --git a/cognee/modules/ontology/get_ontology_resolver.py b/cognee/modules/ontology/get_ontology_resolver.py
new file mode 100644
index 000000000..01377c162
--- /dev/null
+++ b/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
diff --git a/cognee/modules/ontology/ontology_config.py b/cognee/modules/ontology/ontology_config.py
new file mode 100644
index 000000000..e28da9f92
--- /dev/null
+++ b/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]