Tweaked order in init

Signed-off-by: Mark Sze <[email protected]>
ag2ai · Dec 3, 2024 · 5ae2fcc · 5ae2fcc
1 parent 2537fbf
commit 5ae2fcc
Showing 1 changed file with 12 additions and 13 deletions.
diff --git a/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py b/autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py
@@ -81,16 +81,6 @@ def init_db(self, input_doc: List[Document]):
         """
         Build the knowledge graph with input documents.
         """
-        self.knowledge_graph = KnowledgeGraph(
-            name=self.name,
-            host=self.host,
-            port=self.port,
-            username=self.username,
-            password=self.password,
-            model_config=KnowledgeGraphModelConfig.with_model(self.model),
-            ontology=self.ontology,
-        )
-
         sources = []
         for doc in input_doc:
             if os.path.exists(doc.path_or_url):
@@ -104,14 +94,23 @@ def init_db(self, input_doc: List[Document]):
                     model=self.model,
                 )
 
+            self.knowledge_graph = KnowledgeGraph(
+                name=self.name,
+                host=self.host,
+                port=self.port,
+                username=self.username,
+                password=self.password,
+                model_config=KnowledgeGraphModelConfig.with_model(self.model),
+                ontology=self.ontology,
+            )
+
             self.knowledge_graph.process_sources(sources)
 
+            # Establishing a chat session will maintain the history
+            self._chat_session = self.knowledge_graph.chat_session()
         else:
             raise ValueError("No input documents could be loaded.")
 
-        # Establishing a chat session will maintain the history
-        self._chat_session = self.knowledge_graph.chat_session()
-
     def add_records(self, new_records: List) -> bool:
         raise NotImplementedError("This method is not supported by FalkorDB SDK yet.")