feat: Initial version if Graph RAG in KAGGLE scenario (#301)

* Initial version if Graph RAG in KAGGLE scenario

* fix CI

* fix a small bug

* fix CI

* fix CI

* fix CI

rdagent/components/coder/factor_coder/CoSTEER/knowledge_management.py

@@ -22,9 +22,9 @@ from rdagent.components.knowledge_management.graph import (
 )
 from rdagent.core.evolving_framework import (
     EvolvableSubjects,
+    EvolvingKnowledgeBase,
     EvoStep,
     Knowledge,
-    KnowledgeBase,
     QueriedKnowledge,
     RAGStrategy,
 )
@@ -71,12 +71,13 @@ class FactorQueriedKnowledge(QueriedKnowledge):
         self.failed_task_info_set = failed_task_info_set
 
 
-class FactorKnowledgeBaseV1(KnowledgeBase):
-    def __init__(self) -> None:
+class FactorKnowledgeBaseV1(EvolvingKnowledgeBase):
+    def __init__(self, path: str | Path = None) -> None:
         self.implementation_trace: dict[str, FactorKnowledge] = dict()
         self.success_task_info_set: set[str] = set()
 
         self.task_to_embedding = dict()
+        super().__init__(path)
 
     def query(self) -> QueriedKnowledge | None:
         """
@@ -746,12 +747,12 @@ class FactorGraphRAGStrategy(RAGStrategy):
         return factor_implementation_queried_graph_knowledge
 
 
-class FactorGraphKnowledgeBase(KnowledgeBase):
-    def __init__(self, init_component_list=None, data_set_knowledge_path=None) -> None:
+class FactorGraphKnowledgeBase(EvolvingKnowledgeBase):
+    def __init__(self, init_component_list=None, path: str | Path = None, data_set_knowledge_path=None) -> None:
         """
         Load knowledge, offer brief information of knowledge and common handle interfaces
         """
-        self.graph: UndirectedGraph = UndirectedGraph.load(Path.cwd() / "graph.pkl")
+        self.graph: UndirectedGraph = UndirectedGraph(Path.cwd() / "graph.pkl")
         logger.info(f"Knowledge Graph loaded, size={self.graph.size()}")
 
         if init_component_list:
@@ -780,6 +781,7 @@ class FactorGraphKnowledgeBase(KnowledgeBase):
         if data_set_knowledge_path:
             with open(data_set_knowledge_path, "r") as f:
                 self.data_set_knowledge_dict = json.load(f)
+        super().__init__(path)
 
     def get_all_nodes_by_label(self, label: str) -> list[UndirectedNode]:
         return self.graph.get_all_nodes_by_label(label)

rdagent/components/coder/model_coder/CoSTEER/knowledge_management.py

@@ -1,11 +1,13 @@
+from pathlib import Path
+
 from rdagent.components.coder.model_coder.conf import MODEL_IMPL_SETTINGS
 from rdagent.components.coder.model_coder.CoSTEER.evaluators import ModelCoderFeedback
 from rdagent.components.coder.model_coder.model import ModelTask
 from rdagent.core.evolving_framework import (
     EvolvableSubjects,
+    EvolvingKnowledgeBase,
     EvoStep,
     Knowledge,
-    KnowledgeBase,
     QueriedKnowledge,
     RAGStrategy,
 )
@@ -49,13 +51,15 @@ class ModelQueriedKnowledge(QueriedKnowledge):
         self.working_task_to_similar_successful_knowledge_dict = dict()
 
 
-class ModelKnowledgeBase(KnowledgeBase):
-    def __init__(self) -> None:
+class ModelKnowledgeBase(EvolvingKnowledgeBase):
+    def __init__(self, path: str | Path = None) -> None:
         self.implementation_trace: dict[str, ModelKnowledge] = dict()
         self.success_task_info_set: set[str] = set()
 
         self.task_to_embedding = dict()
 
+        super().__init__(path)
+
     def query(self) -> QueriedKnowledge | None:
         """
         Query the knowledge base to get the queried knowledge. So far is handled in RAG strategy.

rdagent/components/knowledge_management/graph.py

@@ -12,6 +12,7 @@ from rdagent.components.knowledge_management.vector_base import (
     VectorBase,
     cosine,
 )
+from rdagent.core.knowledge_base import KnowledgeBase
 from rdagent.oai.llm_utils import APIBackend
 
 Node = KnowledgeMetaData
@@ -47,14 +48,14 @@ class UndirectedNode(Node):
         )
 
 
-class Graph:
+class Graph(KnowledgeBase):
     """
     base Graph class for Knowledge Graph Search
     """
 
     def __init__(self, path: str | Path | None = None) -> None:
-        self.path = path
         self.nodes = {}
+        super().__init__(path=path)
 
     def size(self) -> int:
         return len(self.nodes)
@@ -77,22 +78,6 @@ class Graph:
                 return node
         return None
 
-    @classmethod
-    def load(cls: type[Graph], path: str | Path) -> Graph:
-        """use pickle as the default load method"""
-        path = path if isinstance(path, Path) else Path(path)
-        if not path.exists():
-            return cls(path=path)
-
-        with path.open("rb") as f:
-            return pickle.load(f)
-
-    def save(self, path: str | Path) -> None:
-        """use pickle as the default save method"""
-        Path.mkdir(path.parent, exist_ok=True)
-        with path.open("wb") as f:
-            pickle.dump(self, f)
-
     @staticmethod
     def batch_embedding(nodes: list[Node]) -> list[Node]:
         contents = [node.content for node in nodes]
@@ -119,8 +104,8 @@ class UndirectedGraph(Graph):
     """
 
     def __init__(self, path: str | Path | None = None) -> None:
-        super().__init__(path=path)
         self.vector_base: VectorBase = PDVectorBase()
+        super().__init__(path=path)
 
     def __str__(self) -> str:
         return f"UndirectedGraph(nodes={self.nodes})"
@@ -174,16 +159,6 @@ class UndirectedGraph(Graph):
 
             node.add_neighbor(neighbor)
 
-    @classmethod
-    def load(cls: type[UndirectedGraph], path: str | Path) -> UndirectedGraph:
-        """use pickle as the default load method"""
-        path = path if isinstance(path, Path) else Path(path)
-        if not path.exists():
-            return cls(path=path)
-
-        with path.open("rb") as f:
-            return pickle.load(f)
-
     def add_nodes(self, node: UndirectedNode, neighbors: list[UndirectedNode]) -> None:
         if not neighbors:
             self.add_node(node)

rdagent/components/knowledge_management/vector_base.py

@@ -5,6 +5,7 @@ from typing import List, Tuple, Union
 import pandas as pd
 from scipy.spatial.distance import cosine
 
+from rdagent.core.knowledge_base import KnowledgeBase
 from rdagent.log import rdagent_logger as logger
 from rdagent.oai.llm_utils import APIBackend
 
@@ -68,14 +69,11 @@ def contents_to_documents(contents: List[str], label: str = None) -> List[Docume
     return docs
 
 
-class VectorBase:
+class VectorBase(KnowledgeBase):
     """
     This class is used for handling vector storage and query
     """
 
-    def __init__(self, vector_df_path: Union[str, Path] = None, **kwargs):
-        pass
-
     def add(self, document: Union[Document, List[Document]]):
         """
         add new node to vector_df
@@ -104,28 +102,15 @@ class VectorBase:
         """
         pass
 
-    def load(self, **kwargs):
-        """load vector_df"""
-
-    def save(self, **kwargs):
-        """save vector_df"""
-
 
 class PDVectorBase(VectorBase):
     """
     Implement of VectorBase using Pandas
     """
 
-    def __init__(self, vector_df_path: Union[str, Path] = None):
-        super().__init__(vector_df_path)
-
-        if vector_df_path:
-            try:
-                self.vector_df = self.load(vector_df_path)
-            except FileNotFoundError:
-                self.vector_df = pd.DataFrame(columns=["id", "label", "content", "embedding"])
-        else:
-            self.vector_df = pd.DataFrame(columns=["id", "label", "content", "embedding"])
+    def __init__(self, path: Union[str, Path] = None):
+        self.vector_df = pd.DataFrame(columns=["id", "label", "content", "embedding"])
+        super().__init__(path)
 
     def shape(self):
         return self.vector_df.shape
@@ -196,10 +181,3 @@ class PDVectorBase(VectorBase):
         for _, similar_docs in most_similar_docs.iterrows():
             docs.append(Document().from_dict(similar_docs.to_dict()))
         return docs, searched_similarities.to_list()
-
-    def load(self, vector_df_path, **kwargs):
-        vector_df = pd.read_pickle(vector_df_path)
-        return vector_df
-
-    def save(self, vector_df_path, **kwargs):
-        self.vector_df.to_pickle(vector_df_path)

rdagent/components/workflow/conf.py

@@ -14,6 +14,8 @@ class BasePropSetting(BaseSettings):
     """
 
     scen: str = ""
+    knowledge_base: str = ""
+    knowledge_base_path: str = ""
     hypothesis_gen: str = ""
     hypothesis2experiment: str = ""
     coder: str = ""