feat: enable MACER multi-step subgraph refinement and make num_steps …

README.md

@@ -47,12 +47,12 @@ pip install -e .
 
 ## Usage
 ```bash
-bash run.sh naive_rag/graph_rag/mm_rag
+bash run.sh ToG3
 ```
 or
 
 ```bash
-python main.py --config examples/graphrag/config.yaml
+python main.py --config examples/ToG3/config.yaml
 ```
 
 

examples/ToG3/config.yaml

@@ -39,7 +39,8 @@ rag:
   insert_community_nodes: True
   num_workers: 4 # 并行处理chunk的worker数
   similarity_top_k: 10 # 检索到的top_k个节点
-  stages: ["inference","evaluation"]
+  stages: ["create","inference","evaluation"]
   # graph_rag参数
   max_paths_per_chunk: 2 # 每个chunk的最大path数, 也就是每个chunk抽取的max_knowledge_triplets
   max_cluster_size: 5 # 对graph进行聚类以获得commuities
+  num_steps: 3 # MACER 多步迭代的最大步数

rag_factory/args.py

@@ -52,6 +52,7 @@ class RAGConfig:
     stages: List[str] = field(default_factory=lambda: ["create", "inference", "evaluation"])
     max_paths_per_chunk: int = 10
     max_cluster_size: int = 50
+    num_steps: int = 3
 
 @dataclass
 class Query:

rag_factory/llms/openai_compatible.py

@@ -104,7 +104,7 @@ class OpenAICompatible(OpenAI):
         description=LLMMetadata.model_fields["context_window"].description,
     )
     is_chat_model: bool = Field(
-        default=False,
+        default=True,
         description=LLMMetadata.model_fields["is_chat_model"].description,
     )
     is_function_calling_model: bool = Field(

run.sh

@@ -18,6 +18,9 @@ elif [ "$solution" == "graph_rag" ]; then
 elif [ "$solution" == "mm_rag" ]; then
     echo "Starting MultiModalRAG example..."
     python main.py --config examples/multimodal_rag/config.yaml
+elif [ "$solution" == "ToG3" ]; then
+    echo "Starting ToG3 example..."
+    python main.py --config examples/ToG3/config.yaml
 else
     echo "Unknown solution: $solution"
     exit 1

tog3_main.py

@@ -34,6 +34,8 @@
 
 from rag_factory.graph_constructor import GraphRAGConstructor
 from rag_factory.retrivers.graphrag_query_engine import GraphRAGQueryEngine
+from rag_factory.retrivers.refinable_multistep_query_engine import RefinableMultiStepQueryEngine
+from llama_index.core.indices.query.query_transform.base import StepDecomposeQueryTransform
 
 
 def read_args(config_path: Union[str, Path]) -> Tuple[DatasetConfig, LLMConfig, EmbeddingConfig, StorageConfig, RAGConfig]:
@@ -332,8 +334,15 @@ def _query_task(retriever, query_engine, query: Query, solution="naive_rag") ->
         elif rag_config.solution == "multi_modal_rag":
             # TODO: Implement Multi-modal RAG solution
             raise NotImplementedError("Multi-modal RAG solution is not implemented yet.")
-        elif rag_config.solution == "tog3":
-            query_engine = index.as_query_engine()
+        elif rag_config.solution == "ToG3":
+            # 使用 RefinableMultiStepQueryEngine 实现 MACER 子图精化
+            query_transform = StepDecomposeQueryTransform(llm=llm, verbose=True)
+            query_engine = RefinableMultiStepQueryEngine(
+                query_engine=index.as_query_engine(),
+                query_transform=query_transform,
+                index=index,
+                num_steps=rag_config.num_steps,
+            )
         else:
             raise ValueError(f"Unsupported RAG solution: {rag_config.solution}")